feat: reuse simp cache in grind (#8483)

This PR ensures `grind` reuses the `simp` cache between different calls.
Recall that `grind` uses `simp` to normalize terms during
internalization.

Author: leodemoura

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

@@ -54,16 +54,29 @@ 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
   let (bfalseExpr, scState) := shareCommonAlpha (mkConst ``Bool.false) scState
   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
@@ -167,7 +180,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

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
@@ -286,6 +288,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

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

@@ -14,11 +14,13 @@ 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 }
+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)
+  modify fun s => { s with simp }
   return r
 
 /--

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
@@ -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

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

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.