From 86eb73614937735ef38534baeb8b00b2341e587d Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 17:52:45 -0700
Subject: [PATCH 1/6] feat: track `simp` cache in `grind`

---
 src/Lean/Meta/Tactic/Grind/Intro.lean   |  1 +
 src/Lean/Meta/Tactic/Grind/Main.lean    |  2 +-
 src/Lean/Meta/Tactic/Grind/SearchM.lean |  1 +
 src/Lean/Meta/Tactic/Grind/Simp.lean    |  7 ++++---
 src/Lean/Meta/Tactic/Grind/Types.lean   | 15 ++++++++++++++-
 src/Lean/Meta/Tactic/Simp/Main.lean     | 22 ++++++++++++++++------
 6 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Intro.lean b/src/Lean/Meta/Tactic/Grind/Intro.lean
index e68140189dbd..a65a7f321193 100644
--- a/src/Lean/Meta/Tactic/Grind/Intro.lean
+++ b/src/Lean/Meta/Tactic/Grind/Intro.lean
@@ -224,6 +224,7 @@ def intros (generation : Nat) : SearchM Unit := withCurrGoalContext do
   repeat
     if (← isInconsistent) then
       return ()
+    resetSimpCache
     match (← introNext (← getGoal) generation) with
     | .done goal =>
       let goal ← exfalsoIfNotProp goal
diff --git a/src/Lean/Meta/Tactic/Grind/Main.lean b/src/Lean/Meta/Tactic/Grind/Main.lean
index 4a2a8ee5d21f..5762b657a25b 100644
--- a/src/Lean/Meta/Tactic/Grind/Main.lean
+++ b/src/Lean/Meta/Tactic/Grind/Main.lean
@@ -167,7 +167,7 @@ def main (mvarId : MVarId) (params : Params) (fallback : Fallback) : MetaM Resul
     let issues   := (← get).issues
     let trace    := (← get).trace
     let counters := (← get).counters
-    let simp     := (← get).simpStats
+    let simp     := { (← get).simp with }
     if failure?.isNone then
       -- If there are no failures and diagnostics are enabled, we still report the performance counters.
       if (← isDiagnosticsEnabled) then
diff --git a/src/Lean/Meta/Tactic/Grind/SearchM.lean b/src/Lean/Meta/Tactic/Grind/SearchM.lean
index d62279a0d809..bfc9b89ccbe7 100644
--- a/src/Lean/Meta/Tactic/Grind/SearchM.lean
+++ b/src/Lean/Meta/Tactic/Grind/SearchM.lean
@@ -172,6 +172,7 @@ def nextGoal? : SearchM (Option Nat) := do
   let mut choices := (← get).choiceStack
   if choices.isEmpty then
     return none -- done
+  resetSimpCache
   let goal := (← get).goal
   assert! goal.inconsistent
   let some falseProof ← getFalseProof? goal.mvarId
diff --git a/src/Lean/Meta/Tactic/Grind/Simp.lean b/src/Lean/Meta/Tactic/Grind/Simp.lean
index 18b3c387ec4a..7388c2b3ab84 100644
--- a/src/Lean/Meta/Tactic/Grind/Simp.lean
+++ b/src/Lean/Meta/Tactic/Grind/Simp.lean
@@ -14,11 +14,12 @@ import Lean.Meta.Tactic.Grind.MarkNestedProofs
 import Lean.Meta.Tactic.Grind.Canon
 
 namespace Lean.Meta.Grind
+
 /-- Simplifies the given expression using the `grind` simprocs and normalization theorems. -/
 private def simpCore (e : Expr) : GrindM Simp.Result := do
-  let simpStats := (← get).simpStats
-  let (r, simpStats) ← Meta.simp e (← readThe Context).simp (← readThe Context).simprocs (stats := simpStats)
-  modify fun s => { s with simpStats }
+  let simp ← modifyGet fun s => (s.simp, { s with simp := {} })
+  let (r, simp) ← Meta.simpCore e (← readThe Context).simp (← readThe Context).simprocs (s := { simp with numSteps := 0 })
+  modify fun s => { s with simp }
   return r
 
 /--
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 34389a477358..0c44dbbeb436 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -124,7 +124,7 @@ structure State where
   Remark: we currently do not reuse congruence theorems
   -/
   congrThms  : PHashMap CongrTheoremCacheKey CongrTheorem := {}
-  simpStats  : Simp.Stats := {}
+  simp       : Simp.State := {}
   trueExpr   : Expr
   falseExpr  : Expr
   natZExpr   : Expr
@@ -151,6 +151,19 @@ instance : Nonempty MethodsRef := MethodsRefPointed.property
 
 abbrev GrindM := ReaderT MethodsRef $ ReaderT Context $ StateRefT State MetaM
 
+/--
+Reset the `simp` cache maintained by `grind`.
+Recall that the default discharger used in `simp` may take assumptions into account.
+So, we use this function to reset the cache whenever we add new assumptions and
+switch to a different branch of the search tree.
+Note that we do not reset the `dsimp` cache nor the congruence cache since they
+are not affected.
+-/
+-- TODO: consider using a discharger that does not take assumptions into account,
+-- or track their use more carefully to minimize the number of resets.
+def resetSimpCache : GrindM Unit :=
+  modify fun s => { s with simp.cache := {} }
+
 @[inline] def mapGrindM [MonadControlT GrindM m] [Monad m] (f : {α : Type} → GrindM α → GrindM α) {α} (x : m α) : m α :=
   controlAt GrindM fun runInBase => f <| runInBase x
 
diff --git a/src/Lean/Meta/Tactic/Simp/Main.lean b/src/Lean/Meta/Tactic/Simp/Main.lean
index d3da7ca55189..9a3a7e87d794 100644
--- a/src/Lean/Meta/Tactic/Simp/Main.lean
+++ b/src/Lean/Meta/Tactic/Simp/Main.lean
@@ -822,17 +822,22 @@ private def updateUsedSimpsWithZetaDelta (ctx : Context) (stats : Stats) : MetaM
   else
     x
 
-def main (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods := {}) : MetaM (Result × Stats) := do
+def mainCore (e : Expr) (ctx : Context) (s : State := {}) (methods : Methods := {}) : MetaM (Result × State) := do
   let ctx ← ctx.setLctxInitIndices
   withSimpContext ctx do
-    let (r, s) ← go e methods.toMethodsRef ctx |>.run { stats with }
+    let (r, s) ← go e methods.toMethodsRef ctx |>.run s
     trace[Meta.Tactic.simp.numSteps] "{s.numSteps}"
-    let s ← updateUsedSimpsWithZetaDelta ctx { s with }
+    let stats ← updateUsedSimpsWithZetaDelta ctx { s with }
+    let s := { s with diag := stats.diag, usedTheorems := stats.usedTheorems }
     return (r, s)
 where
   go (e : Expr) : SimpM Result :=
     withCatchingRuntimeEx (simp e)
 
+def main (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods := {}) : MetaM (Result × Stats) := do
+  let (r, s) ← mainCore e ctx { stats with } methods
+  return (r, { s with })
+
 def dsimpMain (e : Expr) (ctx : Context) (stats : Stats := {}) (methods : Methods := {}) : MetaM (Expr × Stats) := do
   withSimpContext ctx do
     let (r, s) ← go e methods.toMethodsRef ctx |>.run { stats with }
@@ -845,11 +850,16 @@ where
 end Simp
 open Simp (SimprocsArray Stats)
 
+def simpCore (e : Expr) (ctx : Simp.Context) (simprocs : SimprocsArray := #[]) (discharge? : Option Simp.Discharge := none)
+    (s : Simp.State := {}) : MetaM (Simp.Result × Simp.State) := do profileitM Exception "simp" (← getOptions) do
+  match discharge? with
+  | none   => Simp.mainCore e ctx s (methods := Simp.mkDefaultMethodsCore simprocs)
+  | some d => Simp.mainCore e ctx s (methods := Simp.mkMethods simprocs d (wellBehavedDischarge := false))
+
 def simp (e : Expr) (ctx : Simp.Context) (simprocs : SimprocsArray := #[]) (discharge? : Option Simp.Discharge := none)
     (stats : Stats := {}) : MetaM (Simp.Result × Stats) := do profileitM Exception "simp" (← getOptions) do
-  match discharge? with
-  | none   => Simp.main e ctx stats (methods := Simp.mkDefaultMethodsCore simprocs)
-  | some d => Simp.main e ctx stats (methods := Simp.mkMethods simprocs d (wellBehavedDischarge := false))
+  let (r, s) ← simpCore e ctx simprocs discharge? { stats with }
+  return (r, { s with })
 
 def dsimp (e : Expr) (ctx : Simp.Context) (simprocs : SimprocsArray := #[])
     (stats : Stats := {}) : MetaM (Expr × Stats) := do profileitM Exception "dsimp" (← getOptions) do

From a25fb671e080ab45051c4d8d1b2c21bfc81ea952 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 18:10:44 -0700
Subject: [PATCH 2/6] feat: use `discharger?` that does not access the context

---
 src/Lean/Meta/Tactic/Grind/Intro.lean   |  1 -
 src/Lean/Meta/Tactic/Grind/Main.lean    | 15 ++++++++++++++-
 src/Lean/Meta/Tactic/Grind/SearchM.lean |  1 -
 src/Lean/Meta/Tactic/Grind/Simp.lean    |  3 ++-
 src/Lean/Meta/Tactic/Grind/Types.lean   | 15 +--------------
 5 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Intro.lean b/src/Lean/Meta/Tactic/Grind/Intro.lean
index a65a7f321193..e68140189dbd 100644
--- a/src/Lean/Meta/Tactic/Grind/Intro.lean
+++ b/src/Lean/Meta/Tactic/Grind/Intro.lean
@@ -224,7 +224,6 @@ def intros (generation : Nat) : SearchM Unit := withCurrGoalContext do
   repeat
     if (← isInconsistent) then
       return ()
-    resetSimpCache
     match (← introNext (← getGoal) generation) with
     | .done goal =>
       let goal ← exfalsoIfNotProp goal
diff --git a/src/Lean/Meta/Tactic/Grind/Main.lean b/src/Lean/Meta/Tactic/Grind/Main.lean
index 5762b657a25b..d704e93e435b 100644
--- a/src/Lean/Meta/Tactic/Grind/Main.lean
+++ b/src/Lean/Meta/Tactic/Grind/Main.lean
@@ -54,6 +54,18 @@ def mkMethods (fallback : Fallback) : CoreM Methods := do
        prop e
   }
 
+-- A `simp` discharger that does not use assumptions.
+-- We use it to make sure we don't have to reset the `simp` cache used in `grind`.
+private def discharge? (e : Expr) : SimpM (Option Expr) := do
+  let e := e.cleanupAnnotations
+  let r ← Simp.simp e
+  if let some p ← Simp.dischargeRfl r.expr then
+    return some (mkApp4 (mkConst ``Eq.mpr [levelZero]) e r.expr (← r.getProof) p)
+  else if r.expr.isTrue then
+    return some (← mkOfEqTrue (← r.getProof))
+  else
+    return none
+
 def GrindM.run (x : GrindM α) (params : Params) (fallback : Fallback) : MetaM α := do
   let (falseExpr, scState)  := shareCommonAlpha (mkConst ``False) {}
   let (trueExpr, scState)   := shareCommonAlpha (mkConst ``True) scState
@@ -61,9 +73,10 @@ def GrindM.run (x : GrindM α) (params : Params) (fallback : Fallback) : MetaM 
   let (btrueExpr, scState)  := shareCommonAlpha (mkConst ``Bool.true) scState
   let (natZExpr, scState)   := shareCommonAlpha (mkNatLit 0) scState
   let simprocs := params.normProcs
+  let simpMethods := Simp.mkMethods simprocs discharge? (wellBehavedDischarge := true)
   let simp := params.norm
   let config := params.config
-  x (← mkMethods fallback).toMethodsRef { config, simprocs, simp }
+  x (← mkMethods fallback).toMethodsRef { config, simpMethods, simp }
     |>.run' { scState, trueExpr, falseExpr, natZExpr, btrueExpr, bfalseExpr }
 
 private def mkCleanState (mvarId : MVarId) (params : Params) : MetaM Clean.State := mvarId.withContext do
diff --git a/src/Lean/Meta/Tactic/Grind/SearchM.lean b/src/Lean/Meta/Tactic/Grind/SearchM.lean
index bfc9b89ccbe7..d62279a0d809 100644
--- a/src/Lean/Meta/Tactic/Grind/SearchM.lean
+++ b/src/Lean/Meta/Tactic/Grind/SearchM.lean
@@ -172,7 +172,6 @@ def nextGoal? : SearchM (Option Nat) := do
   let mut choices := (← get).choiceStack
   if choices.isEmpty then
     return none -- done
-  resetSimpCache
   let goal := (← get).goal
   assert! goal.inconsistent
   let some falseProof ← getFalseProof? goal.mvarId
diff --git a/src/Lean/Meta/Tactic/Grind/Simp.lean b/src/Lean/Meta/Tactic/Grind/Simp.lean
index 7388c2b3ab84..b9186e0d4d75 100644
--- a/src/Lean/Meta/Tactic/Grind/Simp.lean
+++ b/src/Lean/Meta/Tactic/Grind/Simp.lean
@@ -18,7 +18,8 @@ namespace Lean.Meta.Grind
 /-- Simplifies the given expression using the `grind` simprocs and normalization theorems. -/
 private def simpCore (e : Expr) : GrindM Simp.Result := do
   let simp ← modifyGet fun s => (s.simp, { s with simp := {} })
-  let (r, simp) ← Meta.simpCore e (← readThe Context).simp (← readThe Context).simprocs (s := { simp with numSteps := 0 })
+  let ctx := (← readThe Context).simp
+  let (r, simp) ← Simp.mainCore e ctx simp (methods := (← readThe Context).simpMethods)
   modify fun s => { s with simp }
   return r
 
diff --git a/src/Lean/Meta/Tactic/Grind/Types.lean b/src/Lean/Meta/Tactic/Grind/Types.lean
index 0c44dbbeb436..ee2881585796 100644
--- a/src/Lean/Meta/Tactic/Grind/Types.lean
+++ b/src/Lean/Meta/Tactic/Grind/Types.lean
@@ -56,7 +56,7 @@ register_builtin_option grind.warning : Bool := {
 /-- Context for `GrindM` monad. -/
 structure Context where
   simp         : Simp.Context
-  simprocs     : Array Simp.Simprocs
+  simpMethods  : Simp.Methods
   config       : Grind.Config
   /--
   If `cheapCases` is `true`, `grind` only applies `cases` to types that contain
@@ -151,19 +151,6 @@ instance : Nonempty MethodsRef := MethodsRefPointed.property
 
 abbrev GrindM := ReaderT MethodsRef $ ReaderT Context $ StateRefT State MetaM
 
-/--
-Reset the `simp` cache maintained by `grind`.
-Recall that the default discharger used in `simp` may take assumptions into account.
-So, we use this function to reset the cache whenever we add new assumptions and
-switch to a different branch of the search tree.
-Note that we do not reset the `dsimp` cache nor the congruence cache since they
-are not affected.
--/
--- TODO: consider using a discharger that does not take assumptions into account,
--- or track their use more carefully to minimize the number of resets.
-def resetSimpCache : GrindM Unit :=
-  modify fun s => { s with simp.cache := {} }
-
 @[inline] def mapGrindM [MonadControlT GrindM m] [Monad m] (f : {α : Type} → GrindM α → GrindM α) {α} (x : m α) : m α :=
   controlAt GrindM fun runInBase => f <| runInBase x
 

From 16fd4c3e8f6dca0e115ed5156d2c9423447c9830 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 18:12:48 -0700
Subject: [PATCH 3/6] feat: profile `simp` invocations from `grind`

---
 src/Lean/Meta/Tactic/Grind/Simp.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/Lean/Meta/Tactic/Grind/Simp.lean b/src/Lean/Meta/Tactic/Grind/Simp.lean
index b9186e0d4d75..caec12c041e8 100644
--- a/src/Lean/Meta/Tactic/Grind/Simp.lean
+++ b/src/Lean/Meta/Tactic/Grind/Simp.lean
@@ -16,7 +16,7 @@ import Lean.Meta.Tactic.Grind.Canon
 namespace Lean.Meta.Grind
 
 /-- Simplifies the given expression using the `grind` simprocs and normalization theorems. -/
-private def simpCore (e : Expr) : GrindM Simp.Result := do
+private def simpCore (e : Expr) : GrindM Simp.Result := do profileitM Exception "grind simp" (← getOptions) do
   let simp ← modifyGet fun s => (s.simp, { s with simp := {} })
   let ctx := (← readThe Context).simp
   let (r, simp) ← Simp.mainCore e ctx simp (methods := (← readThe Context).simpMethods)

From 6cc3b363b256d78cdf771ca1de1432f5ab065891 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 20:32:49 -0700
Subject: [PATCH 4/6] fix: check types

---
 src/Lean/Meta/Tactic/Grind/MatchCond.lean | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/Lean/Meta/Tactic/Grind/MatchCond.lean b/src/Lean/Meta/Tactic/Grind/MatchCond.lean
index e818d5d8d855..2ac96ff7ab1b 100644
--- a/src/Lean/Meta/Tactic/Grind/MatchCond.lean
+++ b/src/Lean/Meta/Tactic/Grind/MatchCond.lean
@@ -286,6 +286,8 @@ where
       | reportIssue! "found term that has not been internalized{indentExpr lhs}\nwhile trying to construct a proof for `MatchCond`{indentExpr e}"
         return none
     let isHEq := α?.isSome
+    unless (← hasSameType root.self rhs) do
+      return none
     let h ← if isHEq then
       mkEqOfHEq (← mkHEqTrans (← mkHEqProof root.self lhs) h) (check := false)
     else

From 92e6585a2798d698e50f7fbfbe47be85562572e0 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 20:32:58 -0700
Subject: [PATCH 5/6] chore: fix test

---
 tests/lean/run/grind_heartbeats.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/lean/run/grind_heartbeats.lean b/tests/lean/run/grind_heartbeats.lean
index 9e158415c11b..8f09b6287374 100644
--- a/tests/lean/run/grind_heartbeats.lean
+++ b/tests/lean/run/grind_heartbeats.lean
@@ -12,7 +12,7 @@ macro_rules
   | `(gen! $n:num) => `(op (f $n) (gen! $(Lean.quote (n.getNat - 1))))
 
 /--
-trace: [grind.issues] (deterministic) timeout at `isDefEq`, maximum number of heartbeats (5000) has been reached
+trace: [grind.issues] (deterministic) timeout at `grind`, maximum number of heartbeats (5000) has been reached
     Use `set_option maxHeartbeats <num>` to set the limit.
     ⏎
     Additional diagnostic information may be available using the `set_option diagnostics true` command.

From 0d8836a2e84df1e551c10e55024fe8a450df49c3 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sun, 25 May 2025 20:33:51 -0700
Subject: [PATCH 6/6] chore: add TODO for offset equalities

---
 src/Lean/Meta/Tactic/Grind/MatchCond.lean | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/Lean/Meta/Tactic/Grind/MatchCond.lean b/src/Lean/Meta/Tactic/Grind/MatchCond.lean
index 2ac96ff7ab1b..318ae450e311 100644
--- a/src/Lean/Meta/Tactic/Grind/MatchCond.lean
+++ b/src/Lean/Meta/Tactic/Grind/MatchCond.lean
@@ -262,6 +262,8 @@ private partial def isStatisfied (e : Expr) : GoalM Bool := do
     e := b
   return false
 
+-- TODO: we don't have support for offset equalities
+
 /-- Constructs a proof for a satisfied `match`-expression condition. -/
 private partial def mkMatchCondProof? (e : Expr) : GoalM (Option Expr) := do
   let_expr Grind.MatchCond f ← e | return none