From bc72487aed06786a62d41015ffc4aad3ccdb4784 Mon Sep 17 00:00:00 2001
From: Leonardo de Moura <leomoura@amazon.com>
Date: Sat, 3 Jan 2026 18:17:23 -0800
Subject: [PATCH] refactor: `Sym.simp` (#11888)

This PR refactors `Sym.simp` to make it more general and customizable.
It also moves the code
to its own subdirectory `Meta/Sym/Simp`.
---
 src/Lean/Meta/Sym.lean                      |  7 +-
 src/Lean/Meta/Sym/EqTrans.lean              | 22 -----
 src/Lean/Meta/Sym/Simp.lean                 | 85 ++-----------------
 src/Lean/Meta/Sym/{ => Simp}/Congr.lean     | 69 ++++++++-------
 src/Lean/Meta/Sym/{ => Simp}/CongrInfo.lean |  0
 src/Lean/Meta/Sym/{ => Simp}/DiscrTree.lean |  0
 src/Lean/Meta/Sym/Simp/Main.lean            | 75 ++++++++++++++++
 src/Lean/Meta/Sym/Simp/Result.lean          | 25 ++++++
 src/Lean/Meta/Sym/{ => Simp}/Rewrite.lean   | 35 ++++----
 src/Lean/Meta/Sym/{ => Simp}/SimpM.lean     | 94 +++++++++++----------
 src/Lean/Meta/Sym/Simp/Simproc.lean         | 22 +++++
 src/Lean/Meta/Sym/Simp/Theorems.lean        | 44 ++++++++++
 src/Lean/Meta/Sym/SimpFun.lean              | 29 -------
 src/Lean/Meta/Sym/SimpResult.lean           | 15 ----
 tests/bench/sym/simp_1.lean                 | 34 +++++---
 15 files changed, 300 insertions(+), 256 deletions(-)
 delete mode 100644 src/Lean/Meta/Sym/EqTrans.lean
 rename src/Lean/Meta/Sym/{ => Simp}/Congr.lean (77%)
 rename src/Lean/Meta/Sym/{ => Simp}/CongrInfo.lean (100%)
 rename src/Lean/Meta/Sym/{ => Simp}/DiscrTree.lean (100%)
 create mode 100644 src/Lean/Meta/Sym/Simp/Main.lean
 create mode 100644 src/Lean/Meta/Sym/Simp/Result.lean
 rename src/Lean/Meta/Sym/{ => Simp}/Rewrite.lean (51%)
 rename src/Lean/Meta/Sym/{ => Simp}/SimpM.lean (74%)
 create mode 100644 src/Lean/Meta/Sym/Simp/Simproc.lean
 create mode 100644 src/Lean/Meta/Sym/Simp/Theorems.lean
 delete mode 100644 src/Lean/Meta/Sym/SimpFun.lean
 delete mode 100644 src/Lean/Meta/Sym/SimpResult.lean

diff --git a/src/Lean/Meta/Sym.lean b/src/Lean/Meta/Sym.lean
index cab323c7e2..874cbeeb21 100644
--- a/src/Lean/Meta/Sym.lean
+++ b/src/Lean/Meta/Sym.lean
@@ -20,13 +20,8 @@ public import Lean.Meta.Sym.AbstractS
 public import Lean.Meta.Sym.Pattern
 public import Lean.Meta.Sym.Apply
 public import Lean.Meta.Sym.InferType
-public import Lean.Meta.Sym.SimpM
-public import Lean.Meta.Sym.CongrInfo
-public import Lean.Meta.Sym.EqTrans
-public import Lean.Meta.Sym.Congr
-public import Lean.Meta.Sym.SimpResult
 public import Lean.Meta.Sym.Simp
-public import Lean.Meta.Sym.DiscrTree
+
 
 /-!
 # Symbolic simulation support.
diff --git a/src/Lean/Meta/Sym/EqTrans.lean b/src/Lean/Meta/Sym/EqTrans.lean
deleted file mode 100644
index 6e0393e55a..0000000000
--- a/src/Lean/Meta/Sym/EqTrans.lean
+++ /dev/null
@@ -1,22 +0,0 @@
-/-
-Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Leonardo de Moura
--/
-module
-prelude
-public import Lean.Meta.Sym.SimpM
-import Lean.Meta.Sym.InferType
-namespace Lean.Meta.Sym
-
-public def mkEqTrans (e₁ : Expr) (e₂ : Expr) (h₁? : Option Expr) (e₃ : Expr) (h₂? : Option Expr) : SymM (Option Expr) := do
-  match h₁?, h₂? with
-  | none,    none    => return none
-  | some _,  none    => return h₁?
-  | none,    some _  => return h₂?
-  | some h₁, some h₂ =>
-    let α ← inferType e₁
-    let u ← getLevel α
-    return mkApp6 (mkConst ``Eq.trans [u]) α e₁ e₂ e₃ h₁ h₂
-
-end Lean.Meta.Sym
diff --git a/src/Lean/Meta/Sym/Simp.lean b/src/Lean/Meta/Sym/Simp.lean
index bed1f2abb7..964a3e0b6b 100644
--- a/src/Lean/Meta/Sym/Simp.lean
+++ b/src/Lean/Meta/Sym/Simp.lean
@@ -5,79 +5,12 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Sym.SimpM
-import Lean.Meta.Tactic.Grind.AlphaShareBuilder
-import Lean.Meta.Sym.EqTrans
-import Lean.Meta.Sym.Rewrite
-import Lean.Meta.Sym.SimpResult
-import Lean.Meta.Sym.SimpFun
-import Lean.Meta.Sym.Congr
-namespace Lean.Meta.Sym.Simp
-open Grind
-
-def simpConst (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpLambda (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpForall (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpLet (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpFVar (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpMVar (e : Expr) : SimpM Result := do
-  -- **TODO**
-  return { expr := e }
-
-def simpApp (e : Expr) : SimpM Result := do
-  congrArgs e
-
-def simpStep : SimpFun := fun e => do
-  match e with
-  | .lit _ | .sort _ | .bvar _ => return { expr := e }
-  | .proj .. =>
-    reportIssue! "unexpected kernel projection term during simplification{indentExpr e}\npre-process and fold them as projection applications"
-    return { expr := e }
-  | .mdata m b =>
-    let r ← simp b
-    if isSameExpr r.expr b then return { expr := e }
-    else return { r with expr := (← mkMDataS m r.expr) }
-  | .const .. => simpConst e
-  | .lam .. => simpLambda e
-  | .forallE .. => simpForall e
-  | .letE .. => simpLet e
-  | .fvar .. => simpFVar e
-  | .mvar .. => simpMVar e
-  | .app .. => simpApp e
-
-def cacheResult (e : Expr) (r : Result) : SimpM Result := do
-  modify fun s => { s with cache := s.cache.insert { expr := e } r }
-  return r
-
-@[export lean_sym_simp]
-def simpImpl (e : Expr) : SimpM Result := do
-  if (← get).numSteps >= (← getConfig).maxSteps then
-    throwError "`simp` failed: maximum number of steps exceeded"
-  if let some result := (← getCache).find? { expr := e } then
-    return result
-  let r ← (simpStep >> rewrite) e
-  if isSameExpr r.expr e then
-    cacheResult e r
-  else
-    let r' ← simp r.expr
-    if isSameExpr r'.expr r.expr then
-      cacheResult e r
-    else
-      cacheResult e (← mkEqTrans e r r')
-
-end Lean.Meta.Sym.Simp
+public import Lean.Meta.Sym.Simp.Congr
+public import Lean.Meta.Sym.Simp.CongrInfo
+public import Lean.Meta.Sym.Simp.DiscrTree
+public import Lean.Meta.Sym.Simp.Main
+public import Lean.Meta.Sym.Simp.Result
+public import Lean.Meta.Sym.Simp.Rewrite
+public import Lean.Meta.Sym.Simp.SimpM
+public import Lean.Meta.Sym.Simp.Simproc
+public import Lean.Meta.Sym.Simp.Theorems
diff --git a/src/Lean/Meta/Sym/Congr.lean b/src/Lean/Meta/Sym/Simp/Congr.lean
similarity index 77%
rename from src/Lean/Meta/Sym/Congr.lean
rename to src/Lean/Meta/Sym/Simp/Congr.lean
index 1d2e4471f0..f268e04069 100644
--- a/src/Lean/Meta/Sym/Congr.lean
+++ b/src/Lean/Meta/Sym/Simp/Congr.lean
@@ -5,11 +5,11 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Sym.SimpM
+public import Lean.Meta.Sym.Simp.SimpM
 import Lean.Meta.Tactic.Grind.AlphaShareBuilder
 import Lean.Meta.Sym.InferType
-import Lean.Meta.Sym.SimpResult
-import Lean.Meta.Sym.CongrInfo
+import Lean.Meta.Sym.Simp.Result
+import Lean.Meta.Sym.Simp.CongrInfo
 namespace Lean.Meta.Sym.Simp
 open Grind
 
@@ -47,21 +47,21 @@ def mkCongr (e : Expr) (f a : Expr) (fr : Result) (ar : Result) (_ : e = .app f
     let β ← inferType e
     let v ← getLevel β
     return mkApp2 (mkConst declName [u, v]) α β
-  match isSameExpr fr.expr f, isSameExpr ar.expr a with
-  | true,  true =>
-    return { expr := e }
-  | false, true =>
-    let expr ← mkAppS fr.expr a
-    let proof? := mkApp4 (← mkCongrPrefix ``congrFun') f fr.expr (← fr.getProof) a
-    return { expr, proof? }
-  | true, false =>
-    let expr ← mkAppS f ar.expr
-    let proof? := mkApp4 (← mkCongrPrefix ``congrArg) a ar.expr f (← ar.getProof)
-    return { expr, proof? }
-  | false, false =>
-    let expr ← mkAppS fr.expr ar.expr
-    let proof? := mkApp6 (← mkCongrPrefix ``congr) f fr.expr a ar.expr (← fr.getProof) (← ar.getProof)
-    return { expr, proof? }
+  match fr, ar with
+  | .rfl,  .rfl =>
+    return .rfl
+  | .step f' hf, .rfl =>
+    let e' ← mkAppS f' a
+    let h := mkApp4 (← mkCongrPrefix ``congrFun') f f' hf a
+    return .step e' h
+  | .rfl, .step a' ha =>
+    let e' ← mkAppS f a'
+    let h := mkApp4 (← mkCongrPrefix ``congrArg) a a' f ha
+    return .step e' h
+  | .step f' hf, .step a' ha =>
+    let e' ← mkAppS f' a'
+    let h := mkApp6 (← mkCongrPrefix ``congr) f f' a a' hf ha
+    return .step e' h
 
 /--
 Returns a proof using `congrFun`
@@ -69,16 +69,16 @@ Returns a proof using `congrFun`
 congrFun.{u, v} {α : Sort u} {β : α → Sort v} {f g : (x : α) → β x} (h : f = g) (a : α) : f a = g a
 ```
 -/
-def mkCongrFun (e : Expr) (f a : Expr) (fr : Result) (_ : e = .app f a) : SymM Result := do
+def mkCongrFun (e : Expr) (f a : Expr) (f' : Expr) (hf : Expr) (_ : e = .app f a) : SymM Result := do
   let .forallE x _ βx _ ← whnfD (← inferType f)
     | throwError "failed to build congruence proof, function expected{indentExpr f}"
   let α ← inferType a
   let u ← getLevel α
   let v ← getLevel (← inferType e)
   let β := Lean.mkLambda x .default α βx
-  let expr ← mkAppS fr.expr a
-  let proof? := mkApp6 (mkConst ``congrFun [u, v]) α β f fr.expr (← fr.getProof) a
-  return { expr, proof? }
+  let e' ← mkAppS f' a
+  let h := mkApp6 (mkConst ``congrFun [u, v]) α β f f' hf a
+  return .step e' h
 
 /--
 Simplify arguments of a function application with a fixed prefix structure.
@@ -89,16 +89,16 @@ def congrFixedPrefix (e : Expr) (prefixSize : Nat) (suffixSize : Nat) : SimpM Re
   let numArgs := e.getAppNumArgs
   if numArgs ≤ prefixSize then
     -- Nothing to be done
-    return { expr := e }
+    return .rfl
   else if numArgs > prefixSize + suffixSize then
     -- **TODO**: over-applied case
-    return { expr := e }
+    return .rfl
   else
     go numArgs e
 where
   go (i : Nat) (e : Expr) : SimpM Result := do
     if i == prefixSize then
-      return { expr := e }
+      return .rfl
     else
       match h : e with
       | .app f a => mkCongr e f a (← go (i - 1) f) (← simp a) h
@@ -114,35 +114,34 @@ def congrInterlaced (e : Expr) (rewritable : Array Bool) : SimpM Result := do
   let numArgs := e.getAppNumArgs
   if h : numArgs = 0 then
     -- Nothing to be done
-    return { expr := e }
+    return .rfl
   else if h : numArgs > rewritable.size then
     -- **TODO**: over-applied case
-    return { expr := e }
+    return .rfl
   else
     go numArgs e (by omega)
 where
   go (i : Nat) (e : Expr) (h : i ≤ rewritable.size) : SimpM Result := do
     if h : i = 0 then
-      return { expr := e }
+      return .rfl
     else
       match h : e with
       | .app f a =>
         let fr ← go (i - 1) f (by omega)
         if rewritable[i - 1] then
           mkCongr e f a fr (← simp a) h
-        else if isSameExpr fr.expr f then
-          return { expr := e }
-        else
-          mkCongrFun e f a fr h
+        else match fr with
+          | .rfl => return .rfl
+          | .step f' hf => mkCongrFun e f a f' hf h
       | _ => unreachable!
 
 /--
 Simplify arguments using a pre-generated congruence theorem.
 Used for functions with proof or `Decidable` arguments.
 -/
-def congrThm (e : Expr) (_ : CongrTheorem) : SimpM Result := do
+def congrThm (_e : Expr) (_ : CongrTheorem) : SimpM Result := do
   -- **TODO**
-  return { expr := e }
+  return .rfl
 
 /--
 Main entry point for simplifying function application arguments.
@@ -151,7 +150,7 @@ Dispatches to the appropriate strategy based on the function's `CongrInfo`.
 public def congrArgs (e : Expr) : SimpM Result := do
   let f := e.getAppFn
   match (← getCongrInfo f) with
-  | .none => return { expr := e }
+  | .none => return .rfl
   | .fixedPrefix prefixSize suffixSize => congrFixedPrefix e prefixSize suffixSize
   | .interlaced rewritable => congrInterlaced e rewritable
   | .congrTheorem thm => congrThm e thm
diff --git a/src/Lean/Meta/Sym/CongrInfo.lean b/src/Lean/Meta/Sym/Simp/CongrInfo.lean
similarity index 100%
rename from src/Lean/Meta/Sym/CongrInfo.lean
rename to src/Lean/Meta/Sym/Simp/CongrInfo.lean
diff --git a/src/Lean/Meta/Sym/DiscrTree.lean b/src/Lean/Meta/Sym/Simp/DiscrTree.lean
similarity index 100%
rename from src/Lean/Meta/Sym/DiscrTree.lean
rename to src/Lean/Meta/Sym/Simp/DiscrTree.lean
diff --git a/src/Lean/Meta/Sym/Simp/Main.lean b/src/Lean/Meta/Sym/Simp/Main.lean
new file mode 100644
index 0000000000..60f1788e83
--- /dev/null
+++ b/src/Lean/Meta/Sym/Simp/Main.lean
@@ -0,0 +1,75 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.Simp.SimpM
+import Lean.Meta.Tactic.Grind.AlphaShareBuilder
+import Lean.Meta.Sym.Simp.Result
+import Lean.Meta.Sym.Simp.Simproc
+import Lean.Meta.Sym.Simp.Congr
+namespace Lean.Meta.Sym.Simp
+open Grind
+
+def simpLambda (_ : Expr) : SimpM Result := do
+  -- **TODO**
+  return .rfl
+
+def simpForall (_ : Expr) : SimpM Result := do
+  -- **TODO**
+  return .rfl
+
+def simpLet (_ : Expr) : SimpM Result := do
+  -- **TODO**
+  return .rfl
+
+def simpFVar (_ : Expr) : SimpM Result := do
+  -- **TODO**
+  return .rfl
+
+def simpMVar (_ : Expr) : SimpM Result := do
+  -- **TODO**
+  return .rfl
+
+def simpApp (e : Expr) : SimpM Result := do
+  congrArgs e
+
+def simpStep : Simproc := fun e => do
+  match e with
+  | .lit _ | .sort _ | .bvar _ | .const .. => return .rfl
+  | .proj .. =>
+    reportIssue! "unexpected kernel projection term during simplification{indentExpr e}\npre-process and fold them as projection applications"
+    return .rfl
+  | .mdata m b =>
+    let r ← simp b
+    match r with
+    | .rfl => return .rfl
+    | .step b' h => return .step (← mkMDataS m b') h
+  | .lam .. => simpLambda e
+  | .forallE .. => simpForall e
+  | .letE .. => simpLet e
+  | .fvar .. => simpFVar e
+  | .mvar .. => simpMVar e
+  | .app .. => simpApp e
+
+abbrev cacheResult (e : Expr) (r : Result) : SimpM Result := do
+  modify fun s => { s with cache := s.cache.insert { expr := e } r }
+  return r
+
+@[export lean_sym_simp]
+def simpImpl (e₁ : Expr) : SimpM Result := do
+  if (← get).numSteps >= (← getConfig).maxSteps then
+    throwError "`simp` failed: maximum number of steps exceeded"
+  if let some result := (← getCache).find? { expr := e₁ } then
+    return result
+  match (← pre e₁) with
+  | .step e₂ h₁ => cacheResult e₁ (← mkEqTransResult e₁ e₂ h₁ (← simp e₂))
+  | .rfl =>
+  let r₁ ← (simpStep >> post) e₁
+  match r₁ with
+  | .rfl => cacheResult e₁ r₁
+  | .step e₂ h₁ => cacheResult e₁ (← mkEqTransResult e₁ e₂ h₁ (← simp e₂))
+
+end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/Simp/Result.lean b/src/Lean/Meta/Sym/Simp/Result.lean
new file mode 100644
index 0000000000..1b5a5a3131
--- /dev/null
+++ b/src/Lean/Meta/Sym/Simp/Result.lean
@@ -0,0 +1,25 @@
+/-
+Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.Simp.SimpM
+public import Lean.Meta.Sym.InferType
+namespace Lean.Meta.Sym.Simp
+
+public abbrev Result.isRfl (result : Result) : Bool :=
+  result matches .rfl
+
+public def mkEqTrans (e₁ : Expr) (e₂ : Expr) (h₁ : Expr) (e₃ : Expr) (h₂ : Expr) : SymM Expr := do
+  let α ← Sym.inferType e₁
+  let u ← Sym.getLevel α
+  return mkApp6 (mkConst ``Eq.trans [u]) α e₁ e₂ e₃ h₁ h₂
+
+public abbrev mkEqTransResult (e₁ : Expr) (e₂ : Expr) (h₁ : Expr) (r₂ : Result) : SymM Result :=
+  match r₂ with
+  | .rfl => return .step e₂ h₁
+  | .step e₃ h₂ => return .step e₃ (← mkEqTrans e₁ e₂ h₁ e₃ h₂)
+
+end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/Rewrite.lean b/src/Lean/Meta/Sym/Simp/Rewrite.lean
similarity index 51%
rename from src/Lean/Meta/Sym/Rewrite.lean
rename to src/Lean/Meta/Sym/Simp/Rewrite.lean
index 435b1641ae..58655db04f 100644
--- a/src/Lean/Meta/Sym/Rewrite.lean
+++ b/src/Lean/Meta/Sym/Simp/Rewrite.lean
@@ -5,17 +5,14 @@ Authors: Leonardo de Moura
 -/
 module
 prelude
-public import Lean.Meta.Sym.SimpM
-public import Lean.Meta.Sym.SimpFun
+public import Lean.Meta.Sym.Simp.SimpM
+public import Lean.Meta.Sym.Simp.Simproc
+public import Lean.Meta.Sym.Simp.Theorems
 import Lean.Meta.Sym.InstantiateS
-import Lean.Meta.Sym.DiscrTree
+import Lean.Meta.Sym.Simp.DiscrTree
 namespace Lean.Meta.Sym.Simp
 open Grind
 
-public def mkTheoremFromDecl (declName : Name) : MetaM Theorem := do
-  let (pattern, rhs) ← mkEqPatternFromDecl declName
-  return { expr := mkConst declName, pattern, rhs }
-
 /--
 Creates proof term for a rewriting step.
 Handles both constant expressions (common case, avoids `instantiateLevelParams`)
@@ -30,22 +27,22 @@ def mkValue (expr : Expr) (pattern : Pattern) (result : MatchUnifyResult) : Expr
 /--
 Tries to rewrite `e` using the given theorem.
 -/
--- **TODO**: Define `Step` result?
-public def Theorem.rewrite? (thm : Theorem) (e : Expr) : SimpM (Option Result) := do
+public def Theorem.rewrite (thm : Theorem) (e : Expr) : SimpM Result := do
   if let some result ← thm.pattern.match? e then
-    let proof? := some <| mkValue thm.expr thm.pattern result
-    let rhs    := thm.rhs.instantiateLevelParams thm.pattern.levelParams result.us
-    let rhs    ← shareCommonInc rhs
-    let expr   ← instantiateRevBetaS rhs result.args
-    return some { expr, proof?  }
+    let proof := mkValue thm.expr thm.pattern result
+    let rhs   := thm.rhs.instantiateLevelParams thm.pattern.levelParams result.us
+    let rhs   ← shareCommonInc rhs
+    let expr  ← instantiateRevBetaS rhs result.args
+    return .step expr proof
   else
-    return none
+    return .rfl
 
-public def rewrite : SimpFun := fun e => do
+public def Theorems.rewrite (thms : Theorems) : Simproc := fun e => do
   -- **TODO**: over-applied terms
-  for thm in (← read).thms.getMatch e do
-    if let some result ← thm.rewrite? e then
+  for thm in thms.getMatch e do
+    let result ← thm.rewrite e
+    if !result.isRfl then
       return result
-  return { expr := e }
+  return .rfl
 
 end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/SimpM.lean b/src/Lean/Meta/Sym/Simp/SimpM.lean
similarity index 74%
rename from src/Lean/Meta/Sym/SimpM.lean
rename to src/Lean/Meta/Sym/Simp/SimpM.lean
index ee92172546..76911557ab 100644
--- a/src/Lean/Meta/Sym/SimpM.lean
+++ b/src/Lean/Meta/Sym/Simp/SimpM.lean
@@ -7,7 +7,6 @@ module
 prelude
 public import Lean.Meta.Sym.SymM
 public import Lean.Meta.Sym.Pattern
-import Lean.Meta.Sym.DiscrTree
 public section
 namespace Lean.Meta.Sym.Simp
 
@@ -99,7 +98,6 @@ invalidating the cache and causing O(2^n) behavior on conditional trees.
 - Avoids entering control-flow binders
 -/
 
-
 /-- Configuration options for the structural simplifier. -/
 structure Config where
   /-- If `true`, unfold let-bindings (zeta reduction) during simplification. -/
@@ -109,44 +107,18 @@ structure Config where
   -- **TODO**: many are still missing
 
 /-- The result of simplifying some expression `e`. -/
-structure Result where
-  /-- The simplified version of `e` -/
-  expr           : Expr
-  /-- A proof that `e = expr`, where the simplified expression is on the RHS.
-  If `none`, the proof is assumed to be `refl`. -/
-  proof?         : Option Expr := none
+inductive Result where
+  | /-- No change -/
+    rfl
+  | /-- Simplified expression `e'` and a proof that `e = e'` -/
+    step (e' : Expr) (proof : Expr)
 
-/--
-A simplification theorem for the structural simplifier.
-
-Contains both the theorem expression and a precomputed pattern for efficient unification
-during rewriting.
--/
-structure Theorem where
-  /-- The theorem expression, typically `Expr.const declName` for a named theorem. -/
-  expr    : Expr
-  /-- Precomputed pattern extracted from the theorem's type for efficient matching. -/
-  pattern : Pattern
-  /-- Right-hand side of the equation. -/
-  rhs     : Expr
-
-instance : BEq Theorem where
-  beq thm₁ thm₂ := thm₁.expr == thm₂.expr
-
-/-- Collection of simplification theorems available to the simplifier. -/
-structure Theorems where
-  thms : DiscrTree Theorem := {}
-
-def Theorems.insert (thms : Theorems) (thm : Theorem) : Theorems :=
-  { thms with thms := insertPattern thms.thms thm.pattern thm }
-
-def Theorems.getMatch (thms : Theorems) (e : Expr) : Array Theorem :=
-  Sym.getMatch thms.thms e
+private opaque MethodsRefPointed : NonemptyType.{0}
+def MethodsRef : Type := MethodsRefPointed.type
+instance : Nonempty MethodsRef := by exact MethodsRefPointed.property
 
 /-- Read-only context for the simplifier. -/
 structure Context where
-  /-- Available simplification theorems. -/
-  thms   : Theorems := {}
   /-- Simplifier configuration options. -/
   config : Config := {}
   /-- Size of the local context when simplification started.
@@ -170,28 +142,64 @@ structure State where
   numSteps := 0
 
 /-- Monad for the structural simplifier, layered on top of `SymM`. -/
-abbrev SimpM := ReaderT Context StateRefT State SymM
+abbrev SimpM := ReaderT MethodsRef $ ReaderT Context StateRefT State SymM
 
 instance : Inhabited (SimpM α) where
   default := throwError "<default>"
 
+abbrev Simproc := Expr → SimpM Result
+
+structure Methods where
+  pre        : Simproc  := fun _ => return .rfl
+  post       : Simproc  := fun _ => return .rfl
+  discharge? : Expr → SimpM (Option Expr) := fun _ => return none
+  /--
+  `wellBehavedDischarge` must **not** be set to `true` IF `discharge?`
+  access local declarations with index >= `Context.lctxInitIndices` when
+  `contextual := false`.
+  Reason: it would prevent us from aggressively caching `simp` results.
+  -/
+  wellBehavedDischarge : Bool := true
+  deriving Inhabited
+
+unsafe def Methods.toMethodsRefImpl (m : Methods) : MethodsRef :=
+  unsafeCast m
+
+@[implemented_by Methods.toMethodsRefImpl]
+opaque Methods.toMethodsRef (m : Methods) : MethodsRef
+
+unsafe abbrev MethodsRef.toMethodsImpl (m : MethodsRef) : Methods :=
+  unsafeCast m
+
+@[implemented_by MethodsRef.toMethodsImpl]
+opaque MethodsRef.toMethods (m : MethodsRef) : Methods
+
+def getMethods : SimpM Methods :=
+  return MethodsRef.toMethods (← read)
+
 /-- Runs a `SimpM` computation with the given theorems and configuration. -/
-abbrev SimpM.run (x : SimpM α) (thms : Theorems := {}) (config : Config := {}) : SymM α := do
+def SimpM.run (x : SimpM α) (methods : Methods := {}) (config : Config := {}) : SymM α := do
   let initialLCtxSize := (← getLCtx).decls.size
-  x { initialLCtxSize, thms, config } |>.run' {}
+  x methods.toMethodsRef { initialLCtxSize, config } |>.run' {}
 
 @[extern "lean_sym_simp"] -- Forward declaration
-opaque simp (e : Expr) : SimpM Result
+opaque simp : Simproc
 
 def getConfig : SimpM Config :=
-  return (← read).config
+  return (← readThe Context).config
 
 abbrev getCache : SimpM Cache :=
   return (← get).cache
 
+abbrev pre : Simproc := fun e => do
+  (← getMethods).pre e
+
+abbrev post : Simproc := fun e => do
+  (← getMethods).post e
+
 end Simp
 
-def simp (e : Expr) (thms : Simp.Theorems := {}) (config : Simp.Config := {}) : SymM Simp.Result := do
-  Simp.SimpM.run (Simp.simp e) thms config
+abbrev simp (e : Expr) (methods :  Simp.Methods := {}) (config : Simp.Config := {}) : SymM Simp.Result := do
+  Simp.SimpM.run (Simp.simp e) methods config
 
 end Lean.Meta.Sym
diff --git a/src/Lean/Meta/Sym/Simp/Simproc.lean b/src/Lean/Meta/Sym/Simp/Simproc.lean
new file mode 100644
index 0000000000..8a9246edcf
--- /dev/null
+++ b/src/Lean/Meta/Sym/Simp/Simproc.lean
@@ -0,0 +1,22 @@
+/-
+Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.Simp.SimpM
+public import Lean.Meta.Sym.Simp.Result
+namespace Lean.Meta.Sym.Simp
+open Grind
+
+public abbrev Simproc.andThen (f g : Simproc) : Simproc := fun e₁ => do
+  let r ← f e₁
+  match r with
+  | .rfl => g e₁
+  | .step e₂ h₁ => mkEqTransResult e₁ e₂ h₁ (← g e₂)
+
+public instance : AndThen Simproc where
+  andThen f g := Simproc.andThen f (g ())
+
+end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/Simp/Theorems.lean b/src/Lean/Meta/Sym/Simp/Theorems.lean
new file mode 100644
index 0000000000..bb3228a6d1
--- /dev/null
+++ b/src/Lean/Meta/Sym/Simp/Theorems.lean
@@ -0,0 +1,44 @@
+/-
+Copyright (c) 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Leonardo de Moura
+-/
+module
+prelude
+public import Lean.Meta.Sym.Pattern
+import Lean.Meta.Sym.Simp.DiscrTree
+public section
+namespace Lean.Meta.Sym.Simp
+
+/--
+A simplification theorem for the structural simplifier.
+
+Contains both the theorem expression and a precomputed pattern for efficient unification
+during rewriting.
+-/
+structure Theorem where
+  /-- The theorem expression, typically `Expr.const declName` for a named theorem. -/
+  expr    : Expr
+  /-- Precomputed pattern extracted from the theorem's type for efficient matching. -/
+  pattern : Pattern
+  /-- Right-hand side of the equation. -/
+  rhs     : Expr
+
+instance : BEq Theorem where
+  beq thm₁ thm₂ := thm₁.expr == thm₂.expr
+
+/-- Collection of simplification theorems available to the simplifier. -/
+structure Theorems where
+  thms : DiscrTree Theorem := {}
+
+def Theorems.insert (thms : Theorems) (thm : Theorem) : Theorems :=
+  { thms with thms := insertPattern thms.thms thm.pattern thm }
+
+def Theorems.getMatch (thms : Theorems) (e : Expr) : Array Theorem :=
+  Sym.getMatch thms.thms e
+
+def mkTheoremFromDecl (declName : Name) : MetaM Theorem := do
+  let (pattern, rhs) ← mkEqPatternFromDecl declName
+  return { expr := mkConst declName, pattern, rhs }
+
+end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/SimpFun.lean b/src/Lean/Meta/Sym/SimpFun.lean
deleted file mode 100644
index a663a452c1..0000000000
--- a/src/Lean/Meta/Sym/SimpFun.lean
+++ /dev/null
@@ -1,29 +0,0 @@
-/-
-Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Leonardo de Moura
--/
-module
-prelude
-public import Lean.Meta.Sym.SimpM
-import Lean.Meta.Sym.EqTrans
-namespace Lean.Meta.Sym.Simp
-open Grind
-public def mkEqTrans (e : Expr) (r₁ : Result) (r₂ : Result) : SimpM Result := do
-  let proof? ← Sym.mkEqTrans e r₁.expr r₁.proof? r₂.expr r₂.proof?
-  return { r₂ with proof? }
-
-public abbrev SimpFun := Expr → SimpM Result
-
-public abbrev SimpFun.andThen (f g : SimpFun) : SimpFun := fun e => do
-  let r₁ ← f e
-  if isSameExpr e r₁.expr then
-    g e
-  else
-    let r₂ ← g r₁.expr
-    mkEqTrans e r₁ r₂
-
-public instance : AndThen SimpFun where
-  andThen f g := SimpFun.andThen f (g ())
-
-end Lean.Meta.Sym.Simp
diff --git a/src/Lean/Meta/Sym/SimpResult.lean b/src/Lean/Meta/Sym/SimpResult.lean
deleted file mode 100644
index b5d4550d12..0000000000
--- a/src/Lean/Meta/Sym/SimpResult.lean
+++ /dev/null
@@ -1,15 +0,0 @@
-/-
-Copyright (c) 2026 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-Released under Apache 2.0 license as described in the file LICENSE.
-Authors: Leonardo de Moura
--/
-module
-prelude
-public import Lean.Meta.Sym.SimpM
-namespace Lean.Meta.Sym.Simp
-
-public def Result.getProof (result : Result) : SymM Expr := do
-  let some proof := result.proof? | mkEqRefl result.expr
-  return proof
-
-end Lean.Meta.Sym.Simp
diff --git a/tests/bench/sym/simp_1.lean b/tests/bench/sym/simp_1.lean
index edc00984f8..b803e6fcc0 100644
--- a/tests/bench/sym/simp_1.lean
+++ b/tests/bench/sym/simp_1.lean
@@ -7,20 +7,27 @@ namespace SimpBench
 ## `SymM` Simplifier benchmarks
 -/
 
-def mkSimpTheorems : MetaM Sym.Simp.Theorems := do
-  let result : Sym.Simp.Theorems := {}
-  let result := result.insert (← Sym.Simp.mkTheoremFromDecl ``Nat.zero_add)
-  return result
+def getProofSize (r : Sym.Simp.Result) : MetaM Nat :=
+  match r with
+  | .rfl => return 0
+  | .step _ p => p.numObjs
+
+def mkSimpMethods : MetaM Sym.Simp.Methods := do
+  let thms : Sym.Simp.Theorems := {}
+  let thm ← Sym.Simp.mkTheoremFromDecl ``Nat.zero_add
+  let thms := thms.insert thm
+  return { post := thms.rewrite }
 
 def simp (e : Expr) : MetaM (Sym.Simp.Result × Float) := Sym.SymM.run' do
   let e ← Grind.shareCommon e
-  let thms ← mkSimpTheorems
+  let methods ← mkSimpMethods
   let startTime ← IO.monoNanosNow
-  let r ← Sym.simp e thms { maxSteps := 100000000 }
+  let r ← Sym.simp e methods { maxSteps := 100000000 }
   let endTime ← IO.monoNanosNow
   -- logInfo e
-  -- logInfo r.expr
-  -- check (← r.getProof)
+  -- match r with
+  -- | .rfl => logInfo "rfl"
+  -- | .step e' h => logInfo e'; logInfo h; check h
   let timeMs := (endTime - startTime).toFloat / 1000000.0
   return (r, timeMs)
 
@@ -37,7 +44,7 @@ def benchTransChain (n : Nat) : MetaM Unit := do
   let e ← mkTransitivityChain n
   forallTelescope e fun _ e => do
     let (r, timeMs) ← simp e
-    let proofSize ← r.proof?.get!.numObjs
+    let proofSize ← getProofSize r
     IO.println s!"trans_chain_{n}: {timeMs}ms, proof_size={proofSize}"
 
 def mkCongrArgStress (depth width : Nat) : MetaM Expr := do
@@ -72,7 +79,7 @@ def benchCongrArgExplosion (depth width : Nat) : MetaM Unit := do
   let e ← mkCongrArgStress depth width
   forallTelescope e fun _ e => do
     let (r, timeMs) ← simp e
-    let proofSize ← r.proof?.get!.numObjs
+    let proofSize ← getProofSize r
     IO.println s!"congr_arg_explosion_{depth}x{width}: {timeMs}ms, proof_size={proofSize}"
 
 -- We simulate this by having many structurally similar subterms
@@ -91,7 +98,7 @@ def mkManySubterms (n : Nat) : MetaM Expr := do
 def benchManyRewrites (n : Nat) : MetaM Unit := do
   let e ← mkManySubterms n
   let (r, timeMs) ← simp e
-  let proofSize ← r.proof?.get!.numObjs
+  let proofSize ← getProofSize r
   IO.println s!"many_rewrites_{n}: {timeMs}ms, proof_size={proofSize}"
 
 /-! ## Run all benchmarks -/
@@ -99,6 +106,11 @@ def runAllBenchmarks : MetaM Unit := do
   IO.println "=== Simplifier Stress Tests ==="
   IO.println ""
 
+  -- benchTransChain 3
+  -- benchCongrArgExplosion 4 3
+  -- benchManyRewrites 3
+  -- if true then return () -- #exit
+
   IO.println ""
   IO.println "--- Benchmark 1: Transitivity chain ---"
   for n in [5, 50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 2000, 3000, 4000, 5000] do