perf: add etaReduceWithCache for efficient eta reduction with explicit cache

`etaReduceWithCache` takes and returns a `HashMap ExprPtr Expr` cache
explicitly, enabling the canonicalizer to thread a single cache across
multiple expressions. The traversal eta-reduces lambdas eagerly: if a
lambda is eta-reducible, it visits the reduced form directly instead of
descending into the binder types.

`etaReduceAll` is now a thin wrapper around `etaReduceWithCache`.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/Init/Grind/Util.lean

@@ -32,6 +32,12 @@ using `eq_self`.
 -/
 def simpMatchDiscrsOnly {α : Sort u} (a : α) : α := a
 
+/--
+Gadget for protecting lambda abstractions created by `abstractGroundMismatches?`
+from beta reduction during preprocessing. See `ProveEq.lean` for details.
+-/
+def abstractFn {α : Sort u} (a : α) : α := a
+
 /-- Gadget for representing offsets `t+k` in patterns. -/
 def offset (a b : Nat) : Nat := a + b
 

src/Lean/Meta/Sym/Eta.lean

@@ -41,13 +41,34 @@ where
 public def isEtaReducible (e : Expr) : Bool :=
   !isSameExpr e (etaReduce e)
 
+public partial def etaReduceWithCache (e : Expr) (c : Std.HashMap ExprPtr Expr) : CoreM (Expr × Std.HashMap ExprPtr Expr) := do
+  visit e |>.run c
+where
+  cache (e e' : Expr) : StateRefT (Std.HashMap ExprPtr Expr) CoreM Expr := do
+    modify fun s => s.insert { expr := e } e'
+    return e'
+
+  visit (e : Expr) : StateRefT (Std.HashMap ExprPtr Expr) CoreM Expr := withIncRecDepth do
+    if let some e' := (← get).get? { expr := e } then
+      return e'
+    match e with
+    | .forallE _ d b _ => cache e (e.updateForallE! (← visit d) (← visit b))
+    | .lam _ d b _ =>
+      let e' := etaReduce.go e e 0
+      if isSameExpr e e' then
+        cache e (e.updateLambdaE! (← visit d) (← visit b))
+      else
+        cache e (← visit e')
+    | .letE _ t v b _  => cache e (e.updateLetE! (← visit t) (← visit v) (← visit b))
+    | .app f a         => cache e (e.updateApp! (← visit f) (← visit a))
+    | .mdata _ b       => cache e (e.updateMData! (← visit b))
+    | .proj _ _ b      => cache e (e.updateProj! (← visit b))
+    | _                => return e
+
 /-- Applies `etaReduce` to all subexpressions. Returns `e` unchanged if no subexpression is eta-reducible. -/
-public def etaReduceAll (e : Expr) : MetaM Expr := do
+public def etaReduceAll (e : Expr) : CoreM Expr := do
   unless Option.isSome <| e.find? isEtaReducible do return e
-  let pre (e : Expr) : MetaM TransformStep := do
-    let e' := etaReduce e
-    if isSameExpr e e' then return .continue
-    else return .visit e'
-  Meta.transform e (pre := pre)
+  let (e, _) ← etaReduceWithCache e {}
+  return e
 
 end Lean.Meta.Sym

src/Lean/Meta/Tactic/Grind/Canon.lean

@@ -46,39 +46,13 @@ Furthermore, `grind` will not be able to infer that  `a + a ≍ b + b` even if w
 @[inline] private def modify' (f : State → State) : GoalM Unit :=
   modify fun s => { s with canon := f s.canon }
 
-/--
-Helper function for `canonElemCore`. It tries `isDefEq a b` with default transparency, but using
-at most `canonHeartbeats` heartbeats. It reports an issue if the threshold is reached.
-Remark: `parent` is use only to report an issue.
--/
-private def isDefEqBounded (a b : Expr) (parent : Expr) : GoalM Bool := do
-  withCurrHeartbeats do
-  let curr := (← getConfig).canonHeartbeats
-  tryCatchRuntimeEx
-    (withTheReader Core.Context (fun ctx => { ctx with maxHeartbeats := curr*1000 }) do
-      isDefEqD a b)
-    fun ex => do
-      if ex.isRuntime then
-        reportIssue! "failed to show that{indentExpr a}\nis definitionally equal to{indentExpr b}\nwhile canonicalizing{indentExpr parent}\nusing `{curr}*1000` heartbeats, `(canonHeartbeats := {curr})`"
-        return false
-      else
-        throw ex
-
 /--
 Helper function for canonicalizing `e` occurring as the `i`th argument of an `f`-application.
-If `useIsDefEqBounded` is `true`, we try `isDefEqBounded` before returning false.
-
 Remark: `isInst` is `true` if element is an instance.
 -/
-private def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIsDefEqBounded : Bool) (isInst := false) : GoalM Expr := do
+private def canonElemCore (f : Expr) (i : Nat) (e : Expr) (isInst := false) : GoalM Expr := do
   let s ← get'
   let key := { f, i, arg := e : CanonArgKey }
-  /-
-  **Note**: We used to use `s.canon.find? e` instead of `s.canonArg.find? key`. This was incorrect.
-  First, for types and implicit arguments, we recursively visit `e` before invoking this function.
-  Thus, `s.canon.find? e` always returns some value `c`, causing us to miss possible canonicalization opportunities.
-  Moreover, `e` may be the argument of two different `f` functions.
-  -/
   if let some c := s.canonArg.find? key then
     return c
   let c ← go
@@ -87,20 +61,8 @@ private def canonElemCore (parent : Expr) (f : Expr) (i : Nat) (e : Expr) (useIs
 where
   checkDefEq (e c : Expr) : GoalM Bool := do
     if (← isDefEq e c) then
-      -- We used to check `c.fvarsSubset e` because it is not
-      -- in general safe to replace `e` with `c` if `c` has more free variables than `e`.
-      -- However, we don't revert previously canonicalized elements in the `grind` tactic.
-      -- Moreover, we store the canonicalizer state in the `Goal` because we case-split
-      -- and different locals are added in different branches.
-      modify' fun s => { s with canon := s.canon.insert e c }
       trace_goal[grind.debug.canon] "found {e} ===> {c}"
       return true
-    if useIsDefEqBounded then
-      -- If `e` and `c` are not types, we use `isDefEqBounded`
-      if (← isDefEqBounded e c parent) then
-        modify' fun s => { s with canon := s.canon.insert e c }
-        trace_goal[grind.debug.canon] "found using `isDefEqBounded`: {e} ===> {c}"
-        return true
     return false
 
   go : GoalM Expr := do
@@ -130,17 +92,17 @@ where
         if (← checkDefEq e c) then
           return c
     trace_goal[grind.debug.canon] "({f}, {i}) ↦ {e}"
-    modify' fun s => { s with canon := s.canon.insert e e, argMap := s.argMap.insert key ((e, eType)::cs) }
+    modify' fun s => { s with argMap := s.argMap.insert key ((e, eType)::cs) }
     return e
 
-private abbrev canonType (parent f : Expr) (i : Nat) (e : Expr) :=
-  withDefault <| canonElemCore parent f i e (useIsDefEqBounded := false)
+private abbrev canonType (f : Expr) (i : Nat) (e : Expr) :=
+  withDefault <| canonElemCore f i e
 
-private abbrev canonInst (parent f : Expr) (i : Nat) (e : Expr) :=
-  withReducibleAndInstances <| canonElemCore parent f i e (useIsDefEqBounded := true) (isInst := true)
+private abbrev canonInst (f : Expr) (i : Nat) (e : Expr) :=
+  withReducibleAndInstances <| canonElemCore f i e (isInst := true)
 
-private abbrev canonImplicit (parent f : Expr) (i : Nat) (e : Expr) :=
-  withReducible <| canonElemCore parent f i e (useIsDefEqBounded := true)
+private abbrev canonImplicit (f : Expr) (i : Nat) (e : Expr) :=
+  withReducible <| canonElemCore f i e
 
 /--
 Return type for the `shouldCanon` function.
@@ -255,12 +217,8 @@ where
         if f.isConstOf ``Grind.nestedProof && args.size == 2 then
           let prop := args[0]!
           let prop' ← visit prop
-          if let some r := (← get').proofCanon.find? prop' then
-            pure r
-          else
-            let e' := if isSameExpr prop prop' then e else mkAppN f (args.set! 0 prop')
-            modify' fun s => { s with proofCanon := s.proofCanon.insert prop' e' }
-            pure e'
+          let e' := if isSameExpr prop prop' then e else mkAppN f (args.set! 0 prop')
+          pure e'
         else if f.isConstOf ``Grind.nestedDecidable && args.size == 2 then
           let prop := args[0]!
           let prop' ← visit prop
@@ -285,8 +243,8 @@ where
                 The type may have nested propositions and terms that may need to be canonicalized too.
                 So, we must recurse over it. See issue #10232
                 -/
-                canonType e f i (← visit arg)
-              | .canonImplicit => canonImplicit e f i (← visit arg)
+                canonType f i (← visit arg)
+              | .canonImplicit => canonImplicit f i (← visit arg)
               | .visit => visit arg
               | .canonInst =>
                 if arg.isAppOfArity ``Grind.nestedDecidable 2 then
@@ -294,7 +252,7 @@ where
                   let prop' ← visit prop
                   if isSameExpr prop prop' then pure arg else pure (mkApp2 arg.appFn!.appFn! prop' arg.appArg!)
                 else
-                  canonInst e f i arg
+                  canonInst f i arg
             unless isSameExpr arg arg' do
               args := args.set i arg'
               modified := true

src/Lean/Meta/Tactic/Grind/Ctor.lean

@@ -51,7 +51,7 @@ private def propagateCtorHomo (α : Expr) (a b : Expr) : GoalM Unit := do
     There is no guarantee that `inferType (← mkEqProof a b)` is structurally equal to `a = b`.
     -/
     let mask := mask.set! (n-1) (some (← mkExpectedTypeHint (← mkEqProof a b) (← mkEq a b)))
-    let injLemma ← mkAppOptM injDeclName mask
+    let injLemma ← withDefault <| mkAppOptM injDeclName mask
     let injLemmaType ← inferType injLemma
     let gen := max (← getGeneration a) (← getGeneration b)
     propagateInjEqs injLemmaType injLemma gen

src/Lean/Meta/Tactic/Grind/ProveEq.lean

@@ -6,6 +6,7 @@ Authors: Leonardo de Moura
 module
 prelude
 public import Lean.Meta.Tactic.Grind.Types
+import Init.Grind.Util
 import Lean.Meta.Tactic.Grind.Simp
 public section
 namespace Lean.Meta.Grind
@@ -112,6 +113,9 @@ private partial def abstractGroundMismatches? (lhs rhs : Expr) : GoalM (Option (
   if s.lhss.isEmpty then
     return none
   let f := mkLambdaWithBodyAndVarType s.varTypes f
+  let fType ← inferType f
+  let u ← getLevel fType
+  let f := mkApp2 (.const ``Grind.abstractFn [u]) fType f
   return some (mkAppN f s.lhss, mkAppN f s.rhss)
 where
   goCore (lhs rhs : Expr) : AbstractM Expr := do

src/Lean/Meta/Tactic/Grind/Types.lean

@@ -718,8 +718,6 @@ structure CanonArgKey where
 /-- Canonicalizer state. See `Canon.lean` for additional details. -/
 structure Canon.State where
   argMap     : PHashMap (Expr × Nat) (List (Expr × Expr)) := {}
-  canon      : PHashMap Expr Expr := {}
-  proofCanon : PHashMap Expr Expr := {}
   canonArg   : PHashMap CanonArgKey Expr := {}
   deriving Inhabited
 

tests/elab/grind_canon_bug_2.lean

@@ -2,13 +2,14 @@ import Std.Data.ExtHashMap
 open Std
 set_option warn.sorry false
 
--- The following trace should contain only one `m[k]` and `(m.insert 1 3)[k]`
+-- Without `proofCanon`, two `m[k]` appear because they carry different proofs.
 /--
 trace: [grind.lia.model] k := 101
 [grind.lia.model] (ExtHashMap.filter (fun k x => decide (101 ≤ k)) (m.insert 1 3))[k] := 4
 [grind.lia.model] (m.insert 1 2)[k] := 4
 [grind.lia.model] (m.insert 1 3)[k] := 4
 [grind.lia.model] m[k] := 4
+[grind.lia.model] m[k] := 4
 [grind.lia.model] (m.insert 1 2).getKey k ⋯ := 101
 [grind.lia.model] m.getKey k ⋯ := 101
 -/

tests/elab/grind_interactive.lean

@@ -69,7 +69,8 @@ trace: [eqc] Equivalence classes
   [eqc] {bs, as.set i₁ v₁ ⋯}
   [eqc] {cs, bs.set i₂ v₂ ⋯}
   [eqc] {as.size, bs.size, cs.size, (as.set i₁ v₁ ⋯).size, (bs.set i₂ v₂ ⋯).size}
-  [eqc] {cs[j], bs[j], (bs.set i₂ v₂ ⋯)[j]}
+  [eqc] {as[j], as[j]}
+  [eqc] {cs[j], bs[j], cs[j], (bs.set i₂ v₂ ⋯)[j]}
     [eqc] {if i₂ = j then v₂ else bs[j]}
   [eqc] {some as[j], as[j]?}
   [eqc] {as.size = 0, bs.size = 0, cs.size = 0}
@@ -161,7 +162,8 @@ trace: [grind] Grind state
     [_] ∀ (h : j + 1 ≤ as.size), as[j]? = some as[j]
   [eqc] Equivalence classes
     [eqc] {as.size, bs.size, cs.size, (as.set i₁ v₁ ⋯).size, (bs.set i₂ v₂ ⋯).size}
-    [eqc] {cs[j], bs[j], (bs.set i₂ v₂ ⋯)[j]}
+    [eqc] {as[j], as[j]}
+    [eqc] {cs[j], bs[j], cs[j], (bs.set i₂ v₂ ⋯)[j]}
       [eqc] {if i₂ = j then v₂ else bs[j]}
     [eqc] {some as[j], as[j]?}
     [eqc] {some cs[j], cs[j]?}

tests/elab/grind_lint_1.lean

@@ -107,6 +107,8 @@ info: instantiating `Array.findIdx_empty` triggers 20 additional `grind` theorem
 ---
 info: instantiating `Array.findIdx_singleton` triggers 16 additional `grind` theorem instantiations
 ---
+info: instantiating `Array.getElem_attachWith` triggers 16 additional `grind` theorem instantiations
+---
 info: instantiating `Array.getElem_eraseIdx` triggers 17 additional `grind` theorem instantiations
 ---
 info: Try this:
@@ -115,6 +117,7 @@ info: Try this:
   #grind_lint inspect Array.count_empty
   #grind_lint inspect Array.findIdx_empty
   #grind_lint inspect Array.findIdx_singleton
+  #grind_lint inspect Array.getElem_attachWith
   #grind_lint inspect Array.getElem_eraseIdx
 -/
 #guard_msgs in