From 3056848819da54c8b5f88329eca08040af82c5e2 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 2 Aug 2025 14:57:25 +0200
Subject: [PATCH] fix: `unfoldReducible'` optimization regression in `grind`
 (#9682)

This PR fixes a regression introduced by an optimization in the
`unfoldReducible` step used by the `grind` normalizer. It also ensures
that projection functions are not reduced, as they are folded in a later
step.
---
 src/Lean/Meta/Tactic/Grind/Simp.lean          | 12 ++--
 src/Lean/Meta/Tactic/Grind/Util.lean          |  6 +-
 .../grind_unfold_reducible_regression.lean    | 68 +++++++++++++++++++
 3 files changed, 81 insertions(+), 5 deletions(-)
 create mode 100644 tests/lean/run/grind_unfold_reducible_regression.lean

diff --git a/src/Lean/Meta/Tactic/Grind/Simp.lean b/src/Lean/Meta/Tactic/Grind/Simp.lean
index f02ca4119a..9f8b1eb9d6 100644
--- a/src/Lean/Meta/Tactic/Grind/Simp.lean
+++ b/src/Lean/Meta/Tactic/Grind/Simp.lean
@@ -41,15 +41,19 @@ def dsimpCore (e : Expr) : GrindM Expr := do profileitM Exception "grind dsimp"
 
 /--
 Unfolds all `reducible` declarations occurring in `e`.
-Similar to `unfoldReducible`, but uses `inShareCommon` as an extra filter
+Similar to `unfoldReducible`, but uses `alreadyInternalized` as an extra filter
 -/
-def unfoldReducible' (e : Expr) : GrindM Expr := do
+def unfoldReducible' (e : Expr) : GoalM Expr := do
   if !(← isUnfoldReducibleTarget e) then return e
-  let pre (e : Expr) : GrindM TransformStep := do
-    if (← inShareCommon e) then return .done e
+  let pre (e : Expr) : GoalM TransformStep := do
+    -- We used to use `inShareCommon` here, but it was to correct.
+    -- `inShareCommon` is used in terms that have not been preprocessed (e.g., proofs).
+    if (← alreadyInternalized e) then return .done e
     let .const declName _ := e.getAppFn | return .continue
     unless (← isReducible declName) do return .continue
     if isGrindGadget declName then return .continue
+    -- See comment at isUnfoldReducibleTarget.
+    if (← getEnv).isProjectionFn declName then return .continue
     let some v ← unfoldDefinition? e | return .continue
     return .visit v
   Core.transform e (pre := pre)
diff --git a/src/Lean/Meta/Tactic/Grind/Util.lean b/src/Lean/Meta/Tactic/Grind/Util.lean
index 2c6f22b425..9924a220fa 100644
--- a/src/Lean/Meta/Tactic/Grind/Util.lean
+++ b/src/Lean/Meta/Tactic/Grind/Util.lean
@@ -59,7 +59,9 @@ def isUnfoldReducibleTarget (e : Expr) : CoreM Bool := do
   return Option.isSome <| e.find? fun e => Id.run do
     let .const declName _ := e | return false
     if getReducibilityStatusCore env declName matches .reducible then
-      return !isGrindGadget declName
+      -- Remark: it is wasteful to unfold projection functions since
+      -- kernel projections are folded again in the `foldProjs` preprocessing step.
+      return !isGrindGadget declName && !env.isProjectionFn declName
     else
       return false
 
@@ -72,6 +74,8 @@ def unfoldReducible (e : Expr) : MetaM Expr := do
     let .const declName _ := e.getAppFn | return .continue
     unless (← isReducible declName) do return .continue
     if isGrindGadget declName then return .continue
+    -- See comment at isUnfoldReducibleTarget.
+    if (← getEnv).isProjectionFn declName then return .continue
     let some v ← unfoldDefinition? e | return .continue
     return .visit v
   Core.transform e (pre := pre)
diff --git a/tests/lean/run/grind_unfold_reducible_regression.lean b/tests/lean/run/grind_unfold_reducible_regression.lean
new file mode 100644
index 0000000000..6ce98fd281
--- /dev/null
+++ b/tests/lean/run/grind_unfold_reducible_regression.lean
@@ -0,0 +1,68 @@
+import Std
+
+open Std
+
+/--
+An extensional tree map with a default value.
+
+To preserve extensionality, we require that the default value is not present in the tree.
+
+**Implementation note**: we use `Ord α` rather than a `cmp : α → α → Ordering` argument,
+because `grind` can not instantiate `ReflCmp` and `TransCmp` theorems because there is no constant to key on.
+-/
+structure TreeMapD (α : Type u) [Ord α] [TransOrd α] (β : Type v) (d : β) where
+  tree : ExtTreeMap α β compare
+  no_default : ∀ a : α, tree[a]? ≠ some d := by grind
+
+namespace TreeMapD
+
+variable {α : Type u} [Ord α] [TransOrd α] {β : Type v} {d : β}
+
+instance : GetElem (TreeMapD α β d) α β (fun _ _ => True) where
+  getElem := fun m a _ => m.tree[a]?.getD d
+
+@[local grind] private theorem getElem_mk
+    (tree : ExtTreeMap α β compare) (no_default : ∀ a : α, tree[a]? ≠ some d) (a : α) :
+    (TreeMapD.mk tree no_default)[a] = tree[a]?.getD d := rfl
+
+@[local grind] private theorem getElem?_tree [DecidableEq β] (m : TreeMapD α β d) (a : α) :
+    m.tree[a]? = if m[a] = d then none else some m[a] := by
+  grind [cases TreeMapD]
+
+@[local grind] private theorem mem_tree (m : TreeMapD α β d) (a : α) :
+    a ∈ m.tree ↔ m[a] ≠ d := by
+  grind [cases TreeMapD]
+
+@[ext, grind ext]
+theorem ext [LawfulEqOrd α] {m₁ m₂ : TreeMapD α β d} (h : ∀ a : α, m₁[a] = m₂[a]) : m₁ = m₂ := by
+  rcases m₁ with ⟨tree₁, no_default₁⟩
+  rcases m₂ with ⟨tree₂, no_default₂⟩
+  congr
+  ext a b
+  specialize h a
+  grind
+
+def empty : TreeMapD α β d where
+  tree := ∅
+
+instance : EmptyCollection (TreeMapD α β d) :=
+  ⟨empty⟩
+
+@[grind =] theorem empty_eq_emptyc : (empty : TreeMapD α β d) = ∅ := rfl
+
+instance : Inhabited (TreeMapD α β d) :=
+  ⟨empty⟩
+
+@[grind =] theorem getElem_empty (a : α) : (∅ : TreeMapD α β d)[a] = d := rfl
+
+variable [DecidableEq β]
+
+def insert (m : TreeMapD α β d) (a : α) (b : β) : TreeMapD α β d where
+  tree := if b = d then m.tree.erase a else m.tree.insert a b
+  no_default := by
+    -- `grind` can't do this split because of the dependent typing in the `xs[i]?` notation.
+    split <;> grind
+
+@[grind =] theorem getElem_insert [DecidableEq α] [LawfulEqOrd α] (m : TreeMapD α β d) (a : α) (b : β) :
+    (m.insert a b)[k] = if k = a then b else m[k] := by
+  grind [insert]