perf: reuse simplifer cache

src/Lean/Meta/Sym/Simp/Goal.lean

@@ -35,6 +35,28 @@ public def SimpGoalResult.toOption : SimpGoalResult → CoreM (Option MVarId)
   | .closed => return none
   | .goal mvarId => return some mvarId
 
+public def SimpGoalResult.ignoreNoProgress : SimpGoalResult → MVarId → SimpGoalResult
+  | .noProgress, mvarId => .goal mvarId
+  | r, _ => r
+
+/--
+Converts a `Simp.Result` value into `SimpGoalResult`.
+-/
+public def Simp.Result.toSimpGoalResult (result : Simp.Result) (mvarId : MVarId) : SymM SimpGoalResult := do
+  let decl ← mvarId.getDecl
+  match result with
+  | .rfl _ => return .noProgress
+  | .step target' h _ =>
+    let mvarNew ← mkFreshExprSyntheticOpaqueMVar target' decl.userName
+    let u ← getLevel decl.type
+    let h := mkApp4 (mkConst ``Eq.mpr [u]) decl.type target' h mvarNew
+    mvarId.assign h
+    if target'.isTrue then
+      mvarNew.mvarId!.assign (mkConst ``True.intro)
+      return .closed
+    else
+      return .goal mvarNew.mvarId!
+
 /--
 Simplifies the target of `mvarId` using `Sym.simp`.
 Returns `.closed` if the target simplifies to `True`, `.simp mvarId'` if simplified
@@ -45,19 +67,7 @@ This function assumed the input goal is a valid `Sym` goal (e.g., expressions ar
 public def simpGoal (mvarId : MVarId) (methods :  Simp.Methods := {}) (config : Simp.Config := {})
     : SymM SimpGoalResult := mvarId.withContext do
   let decl ← mvarId.getDecl
-  let target := decl.type
-  match (← simp target methods config) with
-  | .rfl _ => return .noProgress
-  | .step target' h _ =>
-    let mvarNew ← mkFreshExprSyntheticOpaqueMVar target' decl.userName
-    let u ← getLevel target
-    let h := mkApp4 (mkConst ``Eq.mpr [u]) target target' h mvarNew
-    mvarId.assign h
-    if target'.isTrue then
-      mvarNew.mvarId!.assign (mkConst ``True.intro)
-      return .closed
-    else
-      return .goal mvarNew.mvarId!
+  (← simp decl.type methods config).toSimpGoalResult mvarId
 
 /--
 Similar to `simpGoal`, but returns `.goal mvarId` if no progress was made.

src/Lean/Meta/Sym/Simp/SimpM.lean

@@ -101,7 +101,7 @@ invalidating the cache and causing O(2^n) behavior on conditional trees.
 /-- Configuration options for the structural simplifier. -/
 structure Config where
   /-- Maximum number of steps that can be performed by the simplifier. -/
-  maxSteps : Nat := 100000
+  maxSteps : Nat := 100_000
   /--
   Maximum depth of reentrant simplifier calls through dischargers.
   Prevents infinite loops when conditional rewrite rules trigger recursive discharge attempts.
@@ -173,16 +173,13 @@ abbrev Cache := PHashMap ExprPtr Result
 
 /-- Mutable state for the simplifier. -/
 structure State where
+  /-- Number of steps performed so far. -/
+  numSteps := 0
   /--
   Cache of previously simplified expressions to avoid redundant work.
   **Note**: Consider moving to `SymM.State`
   -/
   cache : Cache := {}
-  /-- Stack of free variables available for reuse when re-entering binders.
-  Each entry is (type pointer, fvarId). -/
-  binderStack : List (ExprPtr × FVarId) := []
-  /-- Number of steps performed so far. -/
-  numSteps := 0
   /-- Cache for generated funext theorems -/
   funext : PHashMap ExprPtr Expr := {}
 
@@ -221,8 +218,13 @@ opaque MethodsRef.toMethods (m : MethodsRef) : Methods
 def getMethods : SimpM Methods :=
   return MethodsRef.toMethods (← read)
 
+/-- Runs a `SimpM` computation with the given theorems, configuration, and initial state -/
+def SimpM.run (x : SimpM α) (methods : Methods := {}) (config : Config := {}) (s : State := {}) : SymM (α × State) := do
+  let initialLCtxSize := (← getLCtx).decls.size
+  x methods.toMethodsRef { initialLCtxSize, config } |>.run s
+
 /-- Runs a `SimpM` computation with the given theorems and configuration. -/
-def SimpM.run (x : SimpM α) (methods : Methods := {}) (config : Config := {}) : SymM α := do
+def SimpM.run' (x : SimpM α) (methods : Methods := {}) (config : Config := {}) : SymM α := do
   let initialLCtxSize := (← getLCtx).decls.size
   x methods.toMethodsRef { initialLCtxSize, config } |>.run' {}
 
@@ -243,7 +245,8 @@ abbrev post : Simproc := fun e => do
 
 abbrev withoutModifyingCache (k : SimpM α) : SimpM α := do
   let cache ← getCache
-  try k finally modify fun s => { s with cache }
+  let funext := (← get).funext
+  try k finally modify fun s => { s with cache, funext }
 
 abbrev withoutModifyingCacheIfNotWellBehaved (k : SimpM α) : SimpM α := do
   if (← getMethods).wellBehavedMethods then k else withoutModifyingCache k
@@ -251,6 +254,6 @@ abbrev withoutModifyingCacheIfNotWellBehaved (k : SimpM α) : SimpM α := do
 end Simp
 
 abbrev simp (e : Expr) (methods :  Simp.Methods := {}) (config : Simp.Config := {}) : SymM Simp.Result := do
-  Simp.SimpM.run (Simp.simp e) methods config
+  Simp.SimpM.run' (Simp.simp e) methods config
 
 end Lean.Meta.Sym

tests/bench/sym/add_sub_cancel.lean

@@ -355,3 +355,59 @@ def runBenchUsingSym (simpEagerly : Bool := false) : MetaM Unit := do
 #eval solveUsingSym 400 true true
 #eval solveUsingSym 500 true true
 #eval solveUsingSym 600 true true
+#eval solveUsingSym 700 true true
+
+/-!
+`SymM` Solution + reusing simplifier cache
+-/
+
+/-- Simplifies the goal reusing the simplifier state. -/
+def simpGoal (mvarId : MVarId) (m : Sym.Simp.Methods) (s : Sym.Simp.State)
+    : SymM (SimpGoalResult × Sym.Simp.State) := mvarId.withContext do
+  Simp.SimpM.run (methods := m) (config := {}) (s := s) do
+    let decl ← mvarId.getDecl
+    return (← (← Simp.simp decl.type).toSimpGoalResult mvarId).ignoreNoProgress mvarId
+
+partial def solveReusingCache (mvarId : MVarId) : SymM Unit := do
+  let exec_cpsRule ← mkBackwardRuleFromDecl ``Exec.seq_cps
+  let inputRule ← mkBackwardRuleFromDecl ``Exec.input
+  let skipRule ← mkBackwardRuleFromDecl ``Exec.skip
+  let setRule ← mkBackwardRuleFromDecl ``Exec.set
+  let rflRule ← mkBackwardRuleFromDecl ``Eq.refl
+  let unfoldMethods ← mkMethods #[``generated_cmd.eq_1, ``repeated_cmds.eq_1, ``repeated_cmds.eq_2]
+  let simpMethods ← mkMethods #[``Expr.eval.eq_1, ``Expr.eval.eq_2, ``Expr.eval.eq_3,
+          ``PartialMap.get_put_diff, ``PartialMap.get_put, ``PartialMap.put_put, ``Binop.interp_add,
+          ``Binop.interp_sub, ``Word.add_sub_cancel, ``Option.some.injEq, ``not_false_eq_true, ``ne_eq]
+  let finalSimpMethods ← mkMethods #[``List.cons.injEq, ``IOEvent.IN.injEq, ``and_true, ``true_and,
+    ``PartialMap.put_put, ``PartialMap.get_put, ``PartialMap.get_put_diff, ``Option.some.injEq,
+    ``and_self, ``exists_eq_True, ``Word.add_sub_cancel]
+  -- ## Initialize
+  let mvarId ← preprocessMVar mvarId
+  let (_, mvarId) ← Sym.introN mvarId 2
+  let .goal mvarId ← Sym.simpGoal mvarId unfoldMethods | failure
+  let .goals [mvarId] ← exec_cpsRule.apply mvarId | failure
+  let .goals [mvarId] ← inputRule.apply mvarId | failure
+  let (_, mvarId) ← Sym.introN mvarId 1
+  -- ## Loop
+  let rec loop (mvarId : MVarId) (simpState : Sym.Simp.State) : SymM MVarId := do
+    let .goals [mvarId] ← exec_cpsRule.apply mvarId | return mvarId
+    let .goals [mvarId', mvarId, _] ← setRule.apply mvarId | failure
+    let (.goal mvarId', simpState) ← simpGoal mvarId' simpMethods simpState | failure
+    let .goals [] ← rflRule.apply mvarId' | failure
+    loop mvarId simpState
+  let mvarId ← loop mvarId {}
+  -- `apply Exec.skip`
+  let .goals [mvarId] ← skipRule.apply mvarId | failure
+  let .closed ← Sym.simpGoal mvarId finalSimpMethods | failure
+  return
+
+def solveUsingCSym (n : Nat) (check := true) : MetaM Unit := do
+  driver n check fun mvarId => SymM.run do solveReusingCache mvarId |>.run' {}
+
+def runBenchUsingCSym : MetaM Unit := do
+  IO.println "=== Symbolic Simulation Tests ==="
+  IO.println ""
+  for n in [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 200, 300, 400, 500, 600, 700] do
+    solveUsingCSym n
+
+#eval runBenchUsingCSym