chore:

tests/bench/sym/simp_1.lean

@@ -12,6 +12,16 @@ def getProofSize (r : Sym.Simp.Result) : MetaM Nat :=
   | .rfl _ => return 0
   | .step _ p _ => p.numObjs
 
+def checkWithKernel (r : Sym.Simp.Result) : MetaM Float := do
+  match r with
+  | .rfl _ => return 0.0
+  | .step _ p _ =>
+    let p := ShareCommon.shareCommon' p
+    let startTime ← IO.monoNanosNow
+    Meta.checkWithKernel p
+    let endTime ← IO.monoNanosNow
+    return (endTime - startTime).toFloat / 1000000.0
+
 def mkSimpMethods : MetaM Sym.Simp.Methods := do
   let thms : Sym.Simp.Theorems := {}
   let thm ← Sym.Simp.mkTheoremFromDecl ``Nat.zero_add
@@ -31,6 +41,33 @@ def simp (e : Expr) : MetaM (Sym.Simp.Result × Float) := Sym.SymM.run do
   let timeMs := (endTime - startTime).toFloat / 1000000.0
   return (r, timeMs)
 
+def benchSimp (name : String) (e : Expr) (check := false) : MetaM Unit :=
+  forallTelescope e fun _ e => do
+  let (r, timeMs) ← simp e
+  let proofSize ← getProofSize r
+  if check then
+    let kMs ← checkWithKernel r
+    IO.println s!"{name}: {timeMs}ms, kernel: {kMs}ms, proof_size={proofSize}"
+  else
+    IO.println s!"{name}: {timeMs}ms, proof_size={proofSize}"
+
+def ppExample (e : Expr) (info := false) : MetaM Unit := do
+  forallTelescope e fun _ e => do
+  IO.println "Example:"
+  IO.println (← ppExpr e)
+  IO.println "====>"
+  match (← simp e).1 with
+  | .rfl _ => IO.println "<no change>"
+  | .step e' h _ =>
+    IO.println (← ppExpr e')
+    IO.println "Proof:"
+    if info then
+      logInfo h
+    else
+      IO.println (← ppExpr h)
+
+  IO.println ""
+
 def mkTransitivityChain (n : Nat) : MetaM Expr := do
   withLocalDeclD `x (mkConst ``Nat) fun x => do
     let zero := mkNatLit 0
@@ -40,12 +77,9 @@ def mkTransitivityChain (n : Nat) : MetaM Expr := do
     mkForallFVars #[x] e
 
 /-- Benchmark: transitivity chain construction -/
-def benchTransChain (n : Nat) : MetaM Unit := do
+def benchTransChain (n : Nat) (check := true) : MetaM Unit := do
   let e ← mkTransitivityChain n
-  forallTelescope e fun _ e => do
-    let (r, timeMs) ← simp e
-    let proofSize ← getProofSize r
-    IO.println s!"trans_chain_{n}: {timeMs}ms, proof_size={proofSize}"
+  benchSimp s!"trans_chain_{n}" e check
 
 def mkCongrArgStress (depth width : Nat) : MetaM Expr := do
   -- Create a term with `depth` layers of functions, each with `width` arguments
@@ -75,12 +109,9 @@ def mkCongrArgStress (depth width : Nat) : MetaM Expr := do
 
     mkForallFVars fs inner
 
-def benchCongrArgExplosion (depth width : Nat) : MetaM Unit := do
+def benchCongrArgExplosion (depth width : Nat) (check := true) : MetaM Unit := do
   let e ← mkCongrArgStress depth width
-  forallTelescope e fun _ e => do
-    let (r, timeMs) ← simp e
-    let proofSize ← getProofSize r
-    IO.println s!"congr_arg_explosion_{depth}x{width}: {timeMs}ms, proof_size={proofSize}"
+  benchSimp s!"congr_arg_explosion_{depth}x{width}" e check
 
 -- We simulate this by having many structurally similar subterms
 def mkManySubterms (n : Nat) : MetaM Expr := do
@@ -95,39 +126,58 @@ def mkManySubterms (n : Nat) : MetaM Expr := do
   return e
 
 /-- Benchmark: many rewrite candidates -/
-def benchManyRewrites (n : Nat) : MetaM Unit := do
+def benchManyRewrites (n : Nat) (check := true) : MetaM Unit := do
   let e ← mkManySubterms n
-  let (r, timeMs) ← simp e
-  let proofSize ← getProofSize r
-  IO.println s!"many_rewrites_{n}: {timeMs}ms, proof_size={proofSize}"
+  benchSimp s!"many_rewrites_{n}" e check
+
+def mkTermTree (n : Nat) : MetaM Expr := do
+  withLocalDeclD `x Nat.mkType fun x => do
+    mkForallFVars #[x] (← go x n 0)
+where
+  go (x : Expr) (n : Nat) (i : Nat) : MetaM Expr := do
+    if h : n = 0 then
+      return mkNatAdd (mkNatLit 0) (mkNatAdd (mkNatLit i) x)
+    else
+      let lhs ← go x (n/2) i
+      let rhs ← go x (n/2) (i + n/2)
+      mkAppM ``Prod.mk #[lhs, rhs]
+
+def benchTermTree (n : Nat) (check := true) : MetaM Unit := do
+  let e ← mkManySubterms n
+  benchSimp s!"term_tree_{n}" e check
+
+
+def run (k : Nat → MetaM Unit) : MetaM Unit := do
+  for n in [10, 20, 50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 2000, 3000, 4000, 5000] do
+    k n
 
 set_option maxRecDepth 100000
+set_option maxHeartbeats 1000000
 
 /-! ## Run all benchmarks -/
 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
-    benchTransChain n
+  ppExample (← mkTransitivityChain 5)
+  run benchTransChain
 
   IO.println ""
   IO.println "--- Benchmark 2: CongrArg explosion ---"
-  for n in [10, 20, 50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 2000, 3000, 4000, 5000] do
-    for w in [3] do
-      benchCongrArgExplosion n w
+  ppExample (← mkCongrArgStress 5 3)
+  run (benchCongrArgExplosion · 3)
 
   IO.println ""
   IO.println "--- Benchmark 3: Many rewrites ---"
-  for n in [10, 50, 100, 200, 300, 400, 500, 600, 700, 800, 900, 1000, 2000, 3000, 4000, 5000] do
-    benchManyRewrites n
+  ppExample (← mkManySubterms 5)
+  run benchManyRewrites
+
+  IO.println ""
+  IO.println "--- Benchmark 4: Term tree rewrites ---"
+  ppExample (← mkTermTree 8)
+  run benchTermTree
 
 #eval runAllBenchmarks