From a5c4d44bd775b0315320fb3309be6b2559c8d989 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 13:10:04 -0500
Subject: [PATCH 01/48] feat: withBenchmark combinator for fine-grained
 benchmarking

---
 Benchmarks/Command.lean | 38 +++++++++++++++++++++++++++++++++++---
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 7d903fc4..da54030c 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -7,8 +7,18 @@ import Lean
 
 open Lean Parser.Command Elab.Command
 
+initialize
+  registerTraceClass `Benchmark
+
+#check withHeartbeats
+
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
-  let stx ← `(command| example $declSig:optDeclSig $val:declVal )
+  logInfo m!"Running {id} benchmark\n"
+
+  let stx ← `(command|
+    set_option trace.benchmark true in
+    example $declSig:optDeclSig $val:declVal
+  )
 
   let n := 5
   let mut runTimes := #[]
@@ -39,12 +49,34 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
 NOTE: even if the actual default value changes at some point in the future,
 this value should *NOT* be updated, to ensure the percentages we've reported
 in previous versions remain comparable. -/
-def defaultMaxHeartbeats : Nat := 200000
+private def defaultMaxHeartbeats : Nat := 200000
+
+private def percentOfDefaultMaxHeartbeats (heartbeats : Nat) : Nat :=
+  heartbeats / (defaultMaxHeartbeats * 10)
 
 open Elab.Tactic in
 elab "logHeartbeats" tac:tactic : tactic => do
   let ((), heartbeats) ← withHeartbeats <|
     evalTactic tac
-  let percent := heartbeats / (defaultMaxHeartbeats * 10)
+  let percent := percentOfDefaultMaxHeartbeats heartbeats
 
   logInfo m!"used {heartbeats / 1000} heartbeats ({percent}% of the default maximum)"
+
+variable {m} [Monad m] [MonadLiftT BaseIO m] [MonadLiftT IO m]
+  [MonadOptions m] [MonadTrace m] [MonadRef m] [AddMessageContext m] in
+/-- `withBenchmark x` is a combinator that will trace the time and heartbeats
+used by `x` to the `benchmark` trace class, in a message like:
+  `{header} took {x}s and {y} heartbeats ({z}% of the default maximum)`
+
+NOTE: the maximum reffered to in the message is `defaultMaxHeartbeats`,
+deliberately *not* the currently confiugred `maxHeartbeats` option, to keep the
+numbers comparable across different values of that option. It's thus entirely
+possible to see more than 100% being reported here. -/
+def withBenchmark (header : String) (x : m α) : m α := do
+  let start ← IO.monoMsNow
+  let (a, heartbeats) ← withHeartbeats x
+  let t := ((← IO.monoMsNow) - start)
+  let percent := percentOfDefaultMaxHeartbeats heartbeats
+  trace[Benchmark] m!"{header} took {t}ms and {heartbeats} heartbeats \
+    ({percent}% of the default maximum)"
+  return a

From 8e44bd9598e8920ab0606e2545f0558334b2070a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:22:23 -0500
Subject: [PATCH 02/48] feat: aggregated benchmarks

---
 Benchmarks/Command.lean | 115 ++++++++++++++++++++++++++++++++++------
 1 file changed, 100 insertions(+), 15 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index da54030c..d6743896 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -8,15 +8,16 @@ import Lean
 open Lean Parser.Command Elab.Command
 
 initialize
-  registerTraceClass `Benchmark
-
-#check withHeartbeats
+  registerOption `benchmark {
+    defValue := false
+    descr := "enables/disables benchmarking in `withBenchmark` combinator"
+  }
 
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   logInfo m!"Running {id} benchmark\n"
 
   let stx ← `(command|
-    set_option trace.benchmark true in
+    set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
@@ -62,21 +63,105 @@ elab "logHeartbeats" tac:tactic : tactic => do
 
   logInfo m!"used {heartbeats / 1000} heartbeats ({percent}% of the default maximum)"
 
-variable {m} [Monad m] [MonadLiftT BaseIO m] [MonadLiftT IO m]
-  [MonadOptions m] [MonadTrace m] [MonadRef m] [AddMessageContext m] in
-/-- `withBenchmark x` is a combinator that will trace the time and heartbeats
-used by `x` to the `benchmark` trace class, in a message like:
+section withBenchmark
+variable {m} [Monad m] [MonadLiftT BaseIO m] [MonadOptions m] [MonadLog m]
+  [AddMessageContext m]
+
+/-- if the `benchmark` option is true, execute `x` and call `f` with the amount
+of heartbeats and milliseconds (in that order!) taken by `x`.
+
+Otherwise, just execute `x` without measurements. -/
+private def withBenchmarkAux (x : m α) (f : Nat → Nat → m Unit)  : m α := do
+  if (← getBoolOption `benchmark) = false then
+    x
+  else
+    let start ← IO.monoMsNow
+    let (a, heartbeats) ← withHeartbeats x
+    let t := ((← IO.monoMsNow) - start)
+    f heartbeats t
+    return a
+
+
+/-- `withBenchmark header x` is a combinator that will, if the `benchmark`
+option is set to `true`, log the time and heartbeats used by `x`,
+in a message like:
   `{header} took {x}s and {y} heartbeats ({z}% of the default maximum)`
 
+Otherwise, if `benchmark` is set to false, `x` is returned as-is.
+
 NOTE: the maximum reffered to in the message is `defaultMaxHeartbeats`,
 deliberately *not* the currently confiugred `maxHeartbeats` option, to keep the
 numbers comparable across different values of that option. It's thus entirely
 possible to see more than 100% being reported here. -/
 def withBenchmark (header : String) (x : m α) : m α := do
-  let start ← IO.monoMsNow
-  let (a, heartbeats) ← withHeartbeats x
-  let t := ((← IO.monoMsNow) - start)
-  let percent := percentOfDefaultMaxHeartbeats heartbeats
-  trace[Benchmark] m!"{header} took {t}ms and {heartbeats} heartbeats \
-    ({percent}% of the default maximum)"
-  return a
+  withBenchmarkAux x fun heartbeats t => do
+    let percent := percentOfDefaultMaxHeartbeats heartbeats
+    logInfo m!"{header} took: {t}ms and {heartbeats} heartbeats \
+      ({percent}% of the default maximum)"
+
+/-- Benchmark the time and heartbeats taken by a tactic, if the `benchmark`
+option is set to `true` -/
+elab "with_benchmark" t:tactic : tactic => do
+  withBenchmark "{t}" <| Elab.Tactic.evalTactic t
+
+end withBenchmark
+
+/-!
+## Aggregated benchmark statistics
+We define `withAggregatedBenchmark`, which functions like `withBenchmark`,
+except it will store a running average of the statistics in a `BenchmarkState`
+which will be reported in one go when `reportAggregatedBenchmarks` is called.
+-/
+section
+
+structure BenchmarkState.Stats where
+  totalHeartbeats : Nat := 0
+  totalTimeInMs : Nat := 0
+  samples : Nat := 0
+
+structure BenchmarkState where
+  stats : Std.HashMap String BenchmarkState.Stats := .empty
+
+variable {m} [Monad m] [MonadStateOf BenchmarkState m] [MonadLiftT BaseIO m]
+  [MonadOptions m]
+
+/-- `withAggregatedBenchmark header x` is a combinator that will,
+if the `benchmark` option is set to `true`,
+measure the time and heartbeats to the benchmark state in a way that aggregates
+different measurements with the same `header`.
+
+See `reportAggregatedBenchmarks` to log the collected data.
+
+Otherwise, if `benchmark` is set to false, `x` is returned as-is.
+-/
+def withAggregatedBenchmark (header : String) (x : m α) : m α := do
+  withBenchmarkAux x fun heartbeats t => do
+    modify fun state =>
+      let s := state.stats.getD header {}
+      { stats := state.stats.insert header {
+          totalHeartbeats := s.totalHeartbeats + heartbeats
+          totalTimeInMs   := s.totalTimeInMs + t
+          samples         := s.samples + 1
+      }}
+
+variable [MonadLog m] [AddMessageContext m] in
+/--
+if the `benchmark` option is set to `true`, report the data collected by
+`withAggregatedBenchmark`, and reset the state (so that the next call to
+`reportAggregatedBenchmarks` will report only new data).
+-/
+def reportAggregatedBenchmarks : m Unit := do
+  if (← getBoolOption `benchmark) = false then
+    return
+
+  for ⟨header, stats⟩ in (← get).stats do
+    let heartbeats := stats.totalHeartbeats
+    let percent := percentOfDefaultMaxHeartbeats heartbeats
+    let t := stats.totalTimeInMs
+    let n := stats.samples
+    logInfo m!"{header} took: {t}ms and {heartbeats} heartbeats \
+      ({percent}% of the default maximum) in total over {n} samples"
+
+  set ({} : BenchmarkState)
+
+end

From 0ab52f3cb3c90702ce80526d648e92a8343ea603 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:28:59 -0500
Subject: [PATCH 03/48] refactor: extract a `withHeartbeatsAndMs` combinator

---
 Benchmarks/Command.lean | 27 ++++++++++++++-------------
 1 file changed, 14 insertions(+), 13 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index d6743896..b936ba25 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -13,6 +13,13 @@ initialize
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
 
+variable {m} [Monad m] [MonadLiftT BaseIO m] in
+def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
+  let start ← IO.monoMsNow
+  let (a, heartbeats) ← withHeartbeats x
+  let endTime ← IO.monoMsNow
+  return ⟨a, heartbeats, endTime - start⟩
+
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   logInfo m!"Running {id} benchmark\n"
 
@@ -22,14 +29,13 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   )
 
   let n := 5
-  let mut runTimes := #[]
   let mut totalRunTime := 0
-  for _ in [0:n] do
-    let start ← IO.monoMsNow
-    elabCommand stx
-    let endTime ← IO.monoMsNow
-    let runTime := endTime - start
-    runTimes := runTimes.push runTime
+  for i in [0:n] do
+    logInfo m!"\n\nRun {i} (out of {n}):\n"
+    let ((), _, runTime) ← withHeartbeatsAndMs <|
+      elabCommand stx
+
+    logInfo m!"Proof took {runTime / 1000}s in total"
     totalRunTime := totalRunTime + runTime
 
   let avg := totalRunTime.toFloat / n.toFloat / 1000
@@ -40,9 +46,6 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     {avg}s
   geomean over {n} runs:
     {geomean}s
-
-  indidividual runtimes:
-    {runTimes}
 "
 
 /-- The default `maxHeartbeats` setting.
@@ -75,9 +78,7 @@ private def withBenchmarkAux (x : m α) (f : Nat → Nat → m Unit)  : m α :=
   if (← getBoolOption `benchmark) = false then
     x
   else
-    let start ← IO.monoMsNow
-    let (a, heartbeats) ← withHeartbeats x
-    let t := ((← IO.monoMsNow) - start)
+    let (a, heartbeats, t) ← withHeartbeatsAndMs x
     f heartbeats t
     return a
 

From 1e71b8d058fc1459a18ef343ebfeb276697cfd12 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:33:56 -0500
Subject: [PATCH 04/48] feat: withBenchmarksReport combinator

---
 Benchmarks/Command.lean | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index b936ba25..22383927 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -165,4 +165,15 @@ def reportAggregatedBenchmarks : m Unit := do
 
   set ({} : BenchmarkState)
 
+variable [MonadLog m] [AddMessageContext m] in
+/--
+Execute `x` with the default `BenchmarkState`, and report the benchmarks after
+(see `reportAggregatedBenchmarks`).
+-/
+def withBenchmarksReport (x : StateT BenchmarkState m α) : m α :=
+  (Prod.fst <$> ·) <| StateT.run (s := {}) do
+    let a ← x
+    reportAggregatedBenchmarks
+    return a
+
 end

From 59d1c2d17b1f9444c593815f6bd1f28c1c3afc3b Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:43:01 -0500
Subject: [PATCH 05/48] feat: benchmark a less trivial condition, using the
 numBlocks register

---
 Benchmarks/SHA512.lean     | 2 ++
 Benchmarks/SHA512_150.lean | 3 ++-
 Benchmarks/SHA512_225.lean | 3 ++-
 Benchmarks/SHA512_400.lean | 3 ++-
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/Benchmarks/SHA512.lean b/Benchmarks/SHA512.lean
index 1ef411bb..94d47eed 100644
--- a/Benchmarks/SHA512.lean
+++ b/Benchmarks/SHA512.lean
@@ -15,9 +15,11 @@ namespace Benchmarks
 
 def SHA512Bench (nSteps : Nat) : Prop :=
   ∀ (s0 sf : ArmState)
+    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10)
     (_h_s0_pc : read_pc s0 = 0x1264c4#64)
     (_h_s0_err : read_err s0 = StateError.None)
     (_h_s0_sp_aligned : CheckSPAlignment s0)
     (_h_s0_program : s0.program = SHA512.program)
     (_h_run : sf = run nSteps s0),
     r StateField.ERR sf = StateError.None
+    ∧ r (.GPR 2#5) sf = BitVec.ofNat _ (9 - (nSteps / 485))
diff --git a/Benchmarks/SHA512_150.lean b/Benchmarks/SHA512_150.lean
index b54d8f83..798e9078 100644
--- a/Benchmarks/SHA512_150.lean
+++ b/Benchmarks/SHA512_150.lean
@@ -9,7 +9,8 @@ import Benchmarks.SHA512
 
 open Benchmarks
 
-benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 => by
+theorem sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
+  simp only [h, bitvec_rules]
   done
diff --git a/Benchmarks/SHA512_225.lean b/Benchmarks/SHA512_225.lean
index 26310030..38c776fb 100644
--- a/Benchmarks/SHA512_225.lean
+++ b/Benchmarks/SHA512_225.lean
@@ -9,7 +9,8 @@ import Benchmarks.SHA512
 
 open Benchmarks
 
-benchmark sha512_225_instructions : SHA512Bench 225 := fun s0 => by
+benchmark sha512_225_instructions : SHA512Bench 225 := fun s0 _ h => by
   intros
   sym_n 225
+  simp only [h, bitvec_rules]
   done
diff --git a/Benchmarks/SHA512_400.lean b/Benchmarks/SHA512_400.lean
index bd725f6f..55958e5a 100644
--- a/Benchmarks/SHA512_400.lean
+++ b/Benchmarks/SHA512_400.lean
@@ -9,7 +9,8 @@ import Benchmarks.Command
 
 open Benchmarks
 
-benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 => by
+benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
+  simp only [h, bitvec_rules]
   done

From 50128e323391277789ad8eca758056a7a508f421 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:47:48 -0500
Subject: [PATCH 06/48] refactor: BenchT abbrev

---
 Benchmarks/Command.lean | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 22383927..7d67eb09 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -165,12 +165,14 @@ def reportAggregatedBenchmarks : m Unit := do
 
   set ({} : BenchmarkState)
 
+abbrev BenchT := StateT BenchmarkState
+
 variable [MonadLog m] [AddMessageContext m] in
 /--
 Execute `x` with the default `BenchmarkState`, and report the benchmarks after
 (see `reportAggregatedBenchmarks`).
 -/
-def withBenchmarksReport (x : StateT BenchmarkState m α) : m α :=
+def withBenchmarksReport (x : BenchT m α) : m α :=
   (Prod.fst <$> ·) <| StateT.run (s := {}) do
     let a ← x
     reportAggregatedBenchmarks

From 83781625c8bb1e17a65507d0b0669fd8fbd3f496 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:52:36 -0500
Subject: [PATCH 07/48] fix: pc of SHA512 benchmark

---
 Benchmarks/SHA512.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Benchmarks/SHA512.lean b/Benchmarks/SHA512.lean
index 94d47eed..b0901b0c 100644
--- a/Benchmarks/SHA512.lean
+++ b/Benchmarks/SHA512.lean
@@ -16,7 +16,7 @@ namespace Benchmarks
 def SHA512Bench (nSteps : Nat) : Prop :=
   ∀ (s0 sf : ArmState)
     (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10)
-    (_h_s0_pc : read_pc s0 = 0x1264c4#64)
+    (_h_s0_pc : read_pc s0 = 0x1264c0#64)
     (_h_s0_err : read_err s0 = StateError.None)
     (_h_s0_sp_aligned : CheckSPAlignment s0)
     (_h_s0_program : s0.program = SHA512.program)

From 7dc9282447f8ea11cba1b18af4b4cb5283c6deb1 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 14:54:19 -0500
Subject: [PATCH 08/48] fix: reinstate 150 step benchmark

---
 Benchmarks/SHA512_150.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Benchmarks/SHA512_150.lean b/Benchmarks/SHA512_150.lean
index 798e9078..abc176d0 100644
--- a/Benchmarks/SHA512_150.lean
+++ b/Benchmarks/SHA512_150.lean
@@ -9,7 +9,7 @@ import Benchmarks.SHA512
 
 open Benchmarks
 
-theorem sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
+benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
   simp only [h, bitvec_rules]

From 93a3f8180a2768665f7dde8afb6f6b002cb88243 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 15:38:54 -0500
Subject: [PATCH 09/48] feat: preserve insertion order of the first sample of
 an aggregated stat

---
 Benchmarks/Command.lean | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 7d67eb09..a25688e5 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -121,6 +121,7 @@ structure BenchmarkState.Stats where
   samples : Nat := 0
 
 structure BenchmarkState where
+  insertionOrder : List String := []
   stats : Std.HashMap String BenchmarkState.Stats := .empty
 
 variable {m} [Monad m] [MonadStateOf BenchmarkState m] [MonadLiftT BaseIO m]
@@ -139,7 +140,12 @@ def withAggregatedBenchmark (header : String) (x : m α) : m α := do
   withBenchmarkAux x fun heartbeats t => do
     modify fun state =>
       let s := state.stats.getD header {}
-      { stats := state.stats.insert header {
+      { insertionOrder :=
+          if s.samples = 0 then
+            header :: state.insertionOrder
+          else
+            state.insertionOrder
+        stats := state.stats.insert header {
           totalHeartbeats := s.totalHeartbeats + heartbeats
           totalTimeInMs   := s.totalTimeInMs + t
           samples         := s.samples + 1
@@ -155,7 +161,9 @@ def reportAggregatedBenchmarks : m Unit := do
   if (← getBoolOption `benchmark) = false then
     return
 
-  for ⟨header, stats⟩ in (← get).stats do
+  let { insertionOrder, stats } ← get
+  for header in insertionOrder do
+    let stats := stats.getD header {}
     let heartbeats := stats.totalHeartbeats
     let percent := percentOfDefaultMaxHeartbeats heartbeats
     let t := stats.totalTimeInMs

From 1c56f0e516705f4b92570da2f3fcec0dbee3fdae Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 15:58:08 -0500
Subject: [PATCH 10/48] fix: sorry-out alignment goals in benchmark

---
 Benchmarks/SHA512_150.lean | 1 +
 Benchmarks/SHA512_225.lean | 1 +
 Benchmarks/SHA512_400.lean | 1 +
 3 files changed, 3 insertions(+)

diff --git a/Benchmarks/SHA512_150.lean b/Benchmarks/SHA512_150.lean
index abc176d0..d035319d 100644
--- a/Benchmarks/SHA512_150.lean
+++ b/Benchmarks/SHA512_150.lean
@@ -13,4 +13,5 @@ benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
   simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_225.lean b/Benchmarks/SHA512_225.lean
index 38c776fb..a10f76fe 100644
--- a/Benchmarks/SHA512_225.lean
+++ b/Benchmarks/SHA512_225.lean
@@ -13,4 +13,5 @@ benchmark sha512_225_instructions : SHA512Bench 225 := fun s0 _ h => by
   intros
   sym_n 225
   simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_400.lean b/Benchmarks/SHA512_400.lean
index 55958e5a..ea1d8de0 100644
--- a/Benchmarks/SHA512_400.lean
+++ b/Benchmarks/SHA512_400.lean
@@ -13,4 +13,5 @@ benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
   simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
   done

From 9443720f1da7f598bbc564b1b84a8a20ad73b355 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 13:23:00 -0500
Subject: [PATCH 11/48] feat: macro to disable multiple linters

---
 Benchmarks/Command.lean | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index be1912a9..bf605c43 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -52,6 +52,20 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     {geomean}s
 "
 
+/-- Set various options to disable linters -/
+macro "disable_linters" "in" cmd:command : command => `(command|
+  set_option linter.constructorNameAsVariable false in
+  set_option linter.deprecated false in
+  set_option linter.missingDocs false in
+  set_option linter.omit false in
+  set_option linter.suspiciousUnexpanderPatterns false in
+  set_option linter.unnecessarySimpa false in
+  set_option linter.unusedRCasesPattern false in
+  set_option linter.unusedSectionVars false in
+  set_option linter.unusedVariables false in
+  $cmd
+)
+
 /-- The default `maxHeartbeats` setting.
 
 NOTE: even if the actual default value changes at some point in the future,

From e838675e376c8aa317884199141e6933df6e4fb0 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 13:32:36 -0500
Subject: [PATCH 12/48] feat: change the behaviour of `benchmark` when
 `profiler` is set to true.

In particular, run only once, and set other profiler-related options
---
 Benchmarks/Command.lean | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index bf605c43..74ff0c3c 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -12,6 +12,10 @@ initialize
     defValue := false
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
+  registerOption `benchmark.profilerDir {
+    defValue := "profiles/"
+    descr := "where to put profile output files"
+  }
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
 def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
@@ -20,15 +24,33 @@ def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
+/--
+Run a benchmark for a set number of times, and report the average runtime.
+
+If the `profiler` option is set true, we run the benchmark only once, with:
+- `trace.profiler` to true, and
+- `trace.profiler.output` set based on the `benchmark.profilerDir` and the
+    id of the benchmark
+-/
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   logInfo m!"Running {id} benchmark\n"
 
-  let stx ← `(command|
+  let mut n := 5
+  let mut stx ← `(command|
     set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
-  let n := 5
+  if (← getBoolOption `profiler) then
+    let outDir := (← getOptions).getString `benchmark.profilesDir
+    let out := Syntax.mkStrLit s!"{outDir}/{id.getId}"
+    stx ← `(command|
+      set_option trace.profiler true in
+      set_option trace.profiler.out $out in
+      $stx:command
+    )
+    n := 1 -- only run once, if `profiler` is set to true
+
   let mut totalRunTime := 0
   -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
   -- exp(1/n * (log a₁ + log a₂ + log aₙ)).

From 1123c546fe5b2a28fbceac6edcaf1bd2f83305d9 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 13:32:56 -0500
Subject: [PATCH 13/48] feat: add benchmarks with kernel and linters disabled

---
 Benchmarks.lean                            |  5 +++
 Benchmarks/Command.lean                    | 46 +++++++++++++---------
 Benchmarks/SHA512_150_noKernel_noLint.lean | 19 +++++++++
 Benchmarks/SHA512_225_noKernel_noLint.lean | 19 +++++++++
 Benchmarks/SHA512_400_noKernel_noLint.lean | 19 +++++++++
 Benchmarks/SHA512_75.lean                  | 18 +++++++++
 Benchmarks/SHA512_75_noKernel_noLint.lean  | 19 +++++++++
 7 files changed, 126 insertions(+), 19 deletions(-)
 create mode 100644 Benchmarks/SHA512_150_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_225_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_400_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_75.lean
 create mode 100644 Benchmarks/SHA512_75_noKernel_noLint.lean

diff --git a/Benchmarks.lean b/Benchmarks.lean
index 7b7fcc0a..cc1f1c6b 100644
--- a/Benchmarks.lean
+++ b/Benchmarks.lean
@@ -3,6 +3,11 @@ Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Author(s): Alex Keizer
 -/
+import Benchmarks.SHA512_75
+import Benchmarks.SHA512_75_noKernel_noLint
 import Benchmarks.SHA512_150
+import Benchmarks.SHA512_150_noKernel_noLint
 import Benchmarks.SHA512_225
+import Benchmarks.SHA512_225_noKernel_noLint
 import Benchmarks.SHA512_400
+import Benchmarks.SHA512_400_noKernel_noLint
diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 74ff0c3c..7439605c 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -12,6 +12,11 @@ initialize
     defValue := false
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
+  registerOption `benchmark.runs {
+    defValue := (5 : Nat)
+    descr := "controls how many runs the `benchmark` command does. \
+    NOTE: this value is ignored when the `profiler` option is set to true"
+  }
   registerOption `benchmark.profilerDir {
     defValue := "profiles/"
     descr := "where to put profile output files"
@@ -35,22 +40,22 @@ If the `profiler` option is set true, we run the benchmark only once, with:
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   logInfo m!"Running {id} benchmark\n"
 
-  let mut n := 5
+  let originalOpts ← getOptions
+  let mut n := originalOpts.getNat `benchmark.runs
+  let mut opts := originalOpts.setBool `benchmark true
   let mut stx ← `(command|
-    set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
   if (← getBoolOption `profiler) then
-    let outDir := (← getOptions).getString `benchmark.profilesDir
-    let out := Syntax.mkStrLit s!"{outDir}/{id.getId}"
-    stx ← `(command|
-      set_option trace.profiler true in
-      set_option trace.profiler.out $out in
-      $stx:command
-    )
+    let outDir := (← getOptions).getString `benchmark.profilerDir
+    opts := opts.setBool `trace.profiler true
+    opts := opts.setString `trace.profiler.output s!"{outDir}/{id.getId}"
     n := 1 -- only run once, if `profiler` is set to true
 
+  -- Set options
+  modifyScope fun scope => { scope with opts }
+
   let mut totalRunTime := 0
   -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
   -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
@@ -64,6 +69,9 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     totalRunTime := totalRunTime + runTime
     totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
 
+  -- Restore options
+  modifyScope fun scope => { scope with opts := originalOpts }
+
   let avg := totalRunTime.toFloat / n.toFloat / 1000
   let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
   logInfo m!"\
@@ -76,16 +84,16 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
 
 /-- Set various options to disable linters -/
 macro "disable_linters" "in" cmd:command : command => `(command|
-  set_option linter.constructorNameAsVariable false in
-  set_option linter.deprecated false in
-  set_option linter.missingDocs false in
-  set_option linter.omit false in
-  set_option linter.suspiciousUnexpanderPatterns false in
-  set_option linter.unnecessarySimpa false in
-  set_option linter.unusedRCasesPattern false in
-  set_option linter.unusedSectionVars false in
-  set_option linter.unusedVariables false in
-  $cmd
+set_option linter.constructorNameAsVariable false in
+set_option linter.deprecated false in
+set_option linter.missingDocs false in
+set_option linter.omit false in
+set_option linter.suspiciousUnexpanderPatterns false in
+set_option linter.unnecessarySimpa false in
+set_option linter.unusedRCasesPattern false in
+set_option linter.unusedSectionVars false in
+set_option linter.unusedVariables false in
+$cmd
 )
 
 /-- The default `maxHeartbeats` setting.
diff --git a/Benchmarks/SHA512_150_noKernel_noLint.lean b/Benchmarks/SHA512_150_noKernel_noLint.lean
new file mode 100644
index 00000000..a801bd79
--- /dev/null
+++ b/Benchmarks/SHA512_150_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_150_noKernel_noLint : SHA512Bench 150 := fun s0 _ h => by
+  intros
+  sym_n 150
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_225_noKernel_noLint.lean b/Benchmarks/SHA512_225_noKernel_noLint.lean
new file mode 100644
index 00000000..6bdc12b5
--- /dev/null
+++ b/Benchmarks/SHA512_225_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_225_noKernel_noLint : SHA512Bench 225 := fun s0 _ h => by
+  intros
+  sym_n 225
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_400_noKernel_noLint.lean b/Benchmarks/SHA512_400_noKernel_noLint.lean
new file mode 100644
index 00000000..06f90357
--- /dev/null
+++ b/Benchmarks/SHA512_400_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_400_noKernel_noLint : SHA512Bench 400 := fun s0 _ h => by
+  intros
+  sym_n 400
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_75.lean b/Benchmarks/SHA512_75.lean
new file mode 100644
index 00000000..75ffe5ca
--- /dev/null
+++ b/Benchmarks/SHA512_75.lean
@@ -0,0 +1,18 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+set_option profiler true in
+benchmark sha512_75 : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_75_noKernel_noLint.lean b/Benchmarks/SHA512_75_noKernel_noLint.lean
new file mode 100644
index 00000000..454c4e51
--- /dev/null
+++ b/Benchmarks/SHA512_75_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_75_noKernel_noLint : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done

From 430fd9e10d92617fe4976d8ccbce42c5858d1d9f Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 13:53:20 -0500
Subject: [PATCH 14/48] fix: remove accidental profiler option

---
 Benchmarks/SHA512_75.lean | 1 -
 1 file changed, 1 deletion(-)

diff --git a/Benchmarks/SHA512_75.lean b/Benchmarks/SHA512_75.lean
index 75ffe5ca..a0ab30db 100644
--- a/Benchmarks/SHA512_75.lean
+++ b/Benchmarks/SHA512_75.lean
@@ -9,7 +9,6 @@ import Benchmarks.SHA512
 
 open Benchmarks
 
-set_option profiler true in
 benchmark sha512_75 : SHA512Bench 75 := fun s0 _ h => by
   intros
   sym_n 75

From ff86cf135e2fe1cdb7869c1663e9bff24efde13a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 13:53:33 -0500
Subject: [PATCH 15/48] feat: bail early if there are no runs in a benchmark

---
 Benchmarks/Command.lean | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 7439605c..593c1cdd 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -53,6 +53,9 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     opts := opts.setString `trace.profiler.output s!"{outDir}/{id.getId}"
     n := 1 -- only run once, if `profiler` is set to true
 
+  if n = 0 then
+    return ()
+
   -- Set options
   modifyScope fun scope => { scope with opts }
 

From 718d8d659b04ecf7efc72a838e48f0102e40453b Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 14:20:17 -0500
Subject: [PATCH 16/48] refactor: prefer traces over log

Traces cause less clutter in the infoview, and still give a decent enough output in batch mode as well
---
 Benchmarks/Command.lean | 67 +++++++++++++++++++++++------------------
 1 file changed, 37 insertions(+), 30 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 593c1cdd..5517a874 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -21,6 +21,8 @@ initialize
     defValue := "profiles/"
     descr := "where to put profile output files"
   }
+  /- Shouldn't be set directly, instead, use the `benchmark` command -/
+  registerTraceClass `benchmark
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
 def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
@@ -29,6 +31,9 @@ def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
+def withBenchTraceNode (msg : MessageData) : CommandElabM α → CommandElabM α :=
+  withTraceNode `benchmark (fun _ => pure msg)
+
 /--
 Run a benchmark for a set number of times, and report the average runtime.
 
@@ -38,12 +43,11 @@ If the `profiler` option is set true, we run the benchmark only once, with:
     id of the benchmark
 -/
 elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
-  logInfo m!"Running {id} benchmark\n"
-
   let originalOpts ← getOptions
-  let mut n := originalOpts.getNat `benchmark.runs
-  let mut opts := originalOpts.setBool `benchmark true
-  let mut stx ← `(command|
+  let mut n := originalOpts.getNat `benchmark.runs 5
+  let mut opts := originalOpts
+  opts := opts.setBool `benchmark true
+  let stx ← `(command|
     example $declSig:optDeclSig $val:declVal
   )
 
@@ -52,6 +56,8 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     opts := opts.setBool `trace.profiler true
     opts := opts.setString `trace.profiler.output s!"{outDir}/{id.getId}"
     n := 1 -- only run once, if `profiler` is set to true
+  else
+    opts := opts.setBool `trace.benchmark true
 
   if n = 0 then
     return ()
@@ -59,31 +65,32 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   -- Set options
   modifyScope fun scope => { scope with opts }
 
-  let mut totalRunTime := 0
-  -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
-  -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
-  let mut totalRunTimeLog := 0
-  for i in [0:n] do
-    logInfo m!"\n\nRun {i} (out of {n}):\n"
-    let ((), _, runTime) ← withHeartbeatsAndMs <|
-      elabCommand stx
-
-    logInfo m!"Proof took {runTime / 1000}s in total"
-    totalRunTime := totalRunTime + runTime
-    totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
-
-  -- Restore options
-  modifyScope fun scope => { scope with opts := originalOpts }
-
-  let avg := totalRunTime.toFloat / n.toFloat / 1000
-  let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
-  logInfo m!"\
-{id}:
-  average runtime over {n} runs:
-    {avg}s
-  geomean over {n} runs:
-    {geomean}s
-"
+  withBenchTraceNode m!"Running {id} benchmark" <| do
+    let mut totalRunTime := 0
+    -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
+    -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
+    let mut totalRunTimeLog : Float := 0
+    for i in [0:n] do
+      let runTime ← withBenchTraceNode m!"Run {i+1} (out of {n}):" <| do
+        let ((), _, runTime) ← withHeartbeatsAndMs <|
+          elabCommand stx
+
+        trace[benchmark] m!"Proof took {runTime / 1000}s in total"
+        pure runTime
+      totalRunTime := totalRunTime + runTime
+      totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
+
+    let avg := totalRunTime.toFloat / n.toFloat / 1000
+    let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
+    trace[benchmark] m!"\
+  {id}:
+    average runtime over {n} runs:
+      {avg}s
+    geomean over {n} runs:
+      {geomean}s
+  "
+    -- Restore options
+    modifyScope fun scope => { scope with opts := originalOpts }
 
 /-- Set various options to disable linters -/
 macro "disable_linters" "in" cmd:command : command => `(command|

From a2afadbf976d4c95519b0c8786223e68bcd5248d Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 14:25:27 -0500
Subject: [PATCH 17/48] refactor: invoke lean directly to run benchmarks

This way, we don't get old data because of lake's caching.
We still incorporate all benchmark files in the `Benchmarks` data, but only to ensure we build all dependencies; we disable running the actual benchmark when run with `lake build Benchmarks`
---
 Makefile      | 12 ++++++++++++
 lakefile.lean |  1 +
 2 files changed, 13 insertions(+)

diff --git a/Makefile b/Makefile
index 00cf37c6..33d72264 100644
--- a/Makefile
+++ b/Makefile
@@ -5,6 +5,8 @@
 SHELL := /bin/bash
 
 LAKE = lake
+LEAN = $(LAKE) env lean
+GIT = git
 
 NUM_TESTS?=3
 VERBOSE?=--verbose
@@ -37,9 +39,19 @@ awslc_elf:
 cosim:
 	time -p lake exe lnsym $(VERBOSE) --num-tests $(NUM_TESTS)
 
+BENCH = $(LEAN) -Dweak.benchmark.runs=5
 .PHONY: benchmarks
 benchmarks:
+	echo "HEAD is on $($(GIT) rev-parse --short HEAD)"
 	$(LAKE) build Benchmarks
+	$(BENCH) Benchmarks/SHA512_75.lean
+	$(LEAN) Benchmarks/SHA512_75_noKernel_noLint.lean
+	$(LEAN) Benchmarks/SHA512_150.lean
+	$(LEAN) Benchmarks/SHA512_150_noKernel_noLint.lean
+	$(LEAN) Benchmarks/SHA512_225.lean
+	$(LEAN) Benchmarks/SHA512_225_noKernel_noLint.lean
+	$(LEAN) Benchmarks/SHA512_400.lean
+	$(LEAN) Benchmarks/SHA512_400_noKernel_noLint.lean
 
 .PHONY: clean clean_all
 clean:
diff --git a/lakefile.lean b/lakefile.lean
index 7175ad63..66000295 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -33,6 +33,7 @@ lean_lib «Doc» where
   -- add library configuration options here
 
 lean_lib «Benchmarks» where
+  leanOptions := #[⟨`weak.benchmark.runs, (0 : Nat)⟩]
   -- add library configuration options here
 
 @[default_target]

From 99f3f0f7b332b2108395cadfc14cfbeedcd48d04 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 14:38:22 -0500
Subject: [PATCH 18/48] fix: print hash of HEAD when doing benchmarks

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 33d72264..4d8d2d12 100644
--- a/Makefile
+++ b/Makefile
@@ -42,7 +42,7 @@ cosim:
 BENCH = $(LEAN) -Dweak.benchmark.runs=5
 .PHONY: benchmarks
 benchmarks:
-	echo "HEAD is on $($(GIT) rev-parse --short HEAD)"
+	echo "HEAD is on `$(GIT) rev-parse --short HEAD`"
 	$(LAKE) build Benchmarks
 	$(BENCH) Benchmarks/SHA512_75.lean
 	$(LEAN) Benchmarks/SHA512_75_noKernel_noLint.lean

From 998dbbc45bd27eb7e5bdc4eb31231cb42733ea70 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 14:39:19 -0500
Subject: [PATCH 19/48] fix: actually use `BENCH`, which explicitly sets the
 number of runs

---
 Makefile | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/Makefile b/Makefile
index 4d8d2d12..5b22a16d 100644
--- a/Makefile
+++ b/Makefile
@@ -45,13 +45,13 @@ benchmarks:
 	echo "HEAD is on `$(GIT) rev-parse --short HEAD`"
 	$(LAKE) build Benchmarks
 	$(BENCH) Benchmarks/SHA512_75.lean
-	$(LEAN) Benchmarks/SHA512_75_noKernel_noLint.lean
-	$(LEAN) Benchmarks/SHA512_150.lean
-	$(LEAN) Benchmarks/SHA512_150_noKernel_noLint.lean
-	$(LEAN) Benchmarks/SHA512_225.lean
-	$(LEAN) Benchmarks/SHA512_225_noKernel_noLint.lean
-	$(LEAN) Benchmarks/SHA512_400.lean
-	$(LEAN) Benchmarks/SHA512_400_noKernel_noLint.lean
+	$(BENCH) Benchmarks/SHA512_75_noKernel_noLint.lean
+	$(BENCH) Benchmarks/SHA512_150.lean
+	$(BENCH) Benchmarks/SHA512_150_noKernel_noLint.lean
+	$(BENCH) Benchmarks/SHA512_225.lean
+	$(BENCH) Benchmarks/SHA512_225_noKernel_noLint.lean
+	$(BENCH) Benchmarks/SHA512_400.lean
+	$(BENCH) Benchmarks/SHA512_400_noKernel_noLint.lean
 
 .PHONY: clean clean_all
 clean:

From 5253302d9d5c031fe89261a68cb11a37f886141f Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 14:40:55 -0500
Subject: [PATCH 20/48] feat: `profiler` make target

This target runs all the same benchmarks, but with the profiler option enabled
---
 Makefile  | 13 +++++++++++++
 README.md |  2 ++
 2 files changed, 15 insertions(+)

diff --git a/Makefile b/Makefile
index 5b22a16d..ced17c10 100644
--- a/Makefile
+++ b/Makefile
@@ -53,6 +53,19 @@ benchmarks:
 	$(BENCH) Benchmarks/SHA512_400.lean
 	$(BENCH) Benchmarks/SHA512_400_noKernel_noLint.lean
 
+PROF = $(LEAN) -Dprofiler=true
+.PHONY: profile
+profile:
+	$(LAKE) build Benchmarks
+	$(PROF) Benchmarks/SHA512_75.lean
+	$(PROF) Benchmarks/SHA512_75_noKernel_noLint.lean
+	$(PROF) Benchmarks/SHA512_150.lean
+	$(PROF) Benchmarks/SHA512_150_noKernel_noLint.lean
+	$(PROF) Benchmarks/SHA512_225.lean
+	$(PROF) Benchmarks/SHA512_225_noKernel_noLint.lean
+	$(PROF) Benchmarks/SHA512_400.lean
+	$(PROF) Benchmarks/SHA512_400_noKernel_noLint.lean
+
 .PHONY: clean clean_all
 clean:
 	$(LAKE) clean
diff --git a/README.md b/README.md
index 47939c18..2054fdf4 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,8 @@ native-code programs of interest.
 
 `benchmarks`: run benchmarks for the symbolic simulator.
 
+`profiler`: run a single round of each benchmark, with the profiler enabled
+
 ### Makefile variables that can be passed in at the command line
 
 `VERBOSE`: Verbose mode; prints disassembly of the instructions being

From 8c0940aaefdfa7082b2d93e87ebe646850e82798 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 27 Sep 2024 16:26:05 -0500
Subject: [PATCH 21/48] feat: suppress benchmark trace nodes in the profiler

---
 Benchmarks/Command.lean | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 5517a874..b83681ba 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -31,8 +31,18 @@ def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
-def withBenchTraceNode (msg : MessageData) : CommandElabM α → CommandElabM α :=
-  withTraceNode `benchmark (fun _ => pure msg)
+/-- Adds a trace node with the `benchmark` class, but only if the profiler
+option is *not* set.
+
+We deliberately suppress benchmarking nodes when profiling, since it generally
+only adds noise
+-/
+def withBenchTraceNode (msg : MessageData) (x : CommandElabM α )
+    : CommandElabM α := do
+  if (← getBoolOption `profiler) then
+    x
+  else
+    withTraceNode `benchmark (fun _ => pure msg) x
 
 /--
 Run a benchmark for a set number of times, and report the average runtime.

From ff43937aca2f7bc7dd938ee1e6c1f294f6b72352 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 30 Sep 2024 20:12:10 -0500
Subject: [PATCH 22/48] refactor: drop a use of `h_sp?`

---
 Tactics/Sym.lean | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index b0d5621f..2b2e5414 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -279,19 +279,18 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
     -- `simp` here
     withMainContext' <| do
       let hStep ← SymContext.findFromUserName h_step.getId
-      let lctx ← getLCtx
-      let decls := (← getThe SymContext).h_sp?.bind lctx.findFromUserName?
+      let decls := (← getThe AxEffects).stackAlignmentProof?
       let decls := decls.toArray
       -- If we know SP is aligned, `simp` with that fact
 
       if !decls.isEmpty then
         trace[Tactic.sym] "simplifying {hStep.toExpr} \
-          with {decls.map (·.toExpr)}"
+          with {decls}"
         -- If `decls` is empty, we have no more knowledge than before, so
         -- everything that could've been `simp`ed, already should have been
         let some goal ← do
             let (ctx, simprocs) ← LNSymSimpContext
-              (config := {decide := false}) (decls := decls)
+              (config := {decide := false}) (exprs := decls)
             let goal ← getMainGoal
             LNSymSimp goal ctx simprocs hStep.fvarId
           | throwError "internal error: simp closed goal unexpectedly"

From 79bc06d1781514cab1b5c01cb902a6f2e708924e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 30 Sep 2024 20:15:32 -0500
Subject: [PATCH 23/48] refactor: drop the other use of hSp

---
 Tactics/Sym.lean | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 2b2e5414..abc6e14f 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -182,10 +182,8 @@ def explodeStep (hStep : Expr) : SymM Unit :=
     eff ← eff.withProgramEq c.effects.programProof
     eff ← eff.withField (← c.effects.getField .ERR).proof
 
-    if let some h_sp := c.h_sp? then
-      let hSp ← SymContext.findFromUserName h_sp
-      -- let effWithSp?
-      eff ← match ← eff.withStackAlignment? hSp.toExpr with
+    if let some hSp := c.effects.stackAlignmentProof? then
+      eff ← match ← eff.withStackAlignment? hSp with
         | some newEff => pure newEff
         | none => do
             trace[Tactic.sym] "failed to show stack alignment"
@@ -210,7 +208,7 @@ def explodeStep (hStep : Expr) : SymM Unit :=
               let (ctx, simprocs) ←
                 LNSymSimpContext
                   (config := {failIfUnchanged := false, decide := true})
-                  (decls := #[hSp])
+                  (exprs := #[hSp])
               LNSymSimp subGoal ctx simprocs
 
             if let some subGoal := subGoal? then

From e887c14e2ce4a2f7d93c667d0fbebe4fc2254686 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 30 Sep 2024 20:16:56 -0500
Subject: [PATCH 24/48] refactor: remove `SymContext.h_sp?`

---
 Tactics/Sym/Context.lean | 11 -----------
 1 file changed, 11 deletions(-)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index e38c1ee0..5acdd000 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -77,9 +77,6 @@ structure SymContext where
   and we assume that no overflow happens
   (i.e., `base - x` can never be equal to `base + y`) -/
   pc : BitVec 64
-  /-- `h_sp?`, if present, is a local hypothesis of the form
-  `CheckSPAlignment state` -/
-  h_sp?  : Option Name
 
   /-- The `simp` context used for effect aggregation.
   This collects references to all (non-)effect hypotheses of the intermediate
@@ -197,14 +194,12 @@ end
 This is not a `ToMessageData` instance because we need access to `MetaM` -/
 def toMessageData (c : SymContext) : MetaM MessageData := do
   let h_run ← userNameToMessageData c.h_run
-  let h_sp?  ← c.h_sp?.mapM userNameToMessageData
 
   return m!"\{ finalState := {c.finalState},
   runSteps? := {c.runSteps?},
   h_run := {h_run},
   program := {c.program},
   pc := {c.pc},
-  h_sp? := {h_sp?},
   state_prefix := {c.state_prefix},
   curr_state_number := {c.currentStateNumber},
   effects := {c.effects} }"
@@ -265,7 +260,6 @@ private def initial (state : Expr) : MetaM SymContext := do
       instructions := ∅
     }
     pc := 0
-    h_sp? := none
     aggregateSimpCtx,
     aggregateSimprocs,
     effects := AxEffects.initial state
@@ -401,9 +395,6 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
       modifyThe AxEffects ({ · with
         stackAlignmentProof? := some decl.toExpr
       })
-      modifyThe SymContext ({· with
-        h_sp? := decl.userName
-      })
     )
 
   /- TODO(@alexkeizer): search for any other hypotheses of the form
@@ -458,7 +449,6 @@ evaluation:
   * the `currentStateNumber` is incremented
 -/
 def prepareForNextStep : SymM Unit := do
-  let s ← getNextStateName
   let pc ← do
     let { value, ..} ← AxEffects.getFieldM .PC
     try
@@ -469,7 +459,6 @@ def prepareForNextStep : SymM Unit := do
 
   modifyThe SymContext (fun c => { c with
     pc
-    h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")
     runSteps?   := (· - 1) <$> c.runSteps?
     currentStateNumber := c.currentStateNumber + 1
   })

From 42142efb29f9cfbab637664c998af782233d6b2a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 30 Sep 2024 20:18:42 -0500
Subject: [PATCH 25/48] refactor: simplify `sym1` simplification step

---
 Tactics/Sym.lean | 12 ++++--------
 1 file changed, 4 insertions(+), 8 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index abc6e14f..e72b97c1 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -277,18 +277,14 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
     -- `simp` here
     withMainContext' <| do
       let hStep ← SymContext.findFromUserName h_step.getId
-      let decls := (← getThe AxEffects).stackAlignmentProof?
-      let decls := decls.toArray
-      -- If we know SP is aligned, `simp` with that fact
 
-      if !decls.isEmpty then
+      -- If we know SP is aligned, `simp` with that fact
+      if let some hSp := (← getThe AxEffects).stackAlignmentProof? then
         trace[Tactic.sym] "simplifying {hStep.toExpr} \
-          with {decls}"
-        -- If `decls` is empty, we have no more knowledge than before, so
-        -- everything that could've been `simp`ed, already should have been
+          with {hSp}"
         let some goal ← do
             let (ctx, simprocs) ← LNSymSimpContext
-              (config := {decide := false}) (exprs := decls)
+              (config := {decide := false}) (exprs := #[hSp])
             let goal ← getMainGoal
             LNSymSimp goal ctx simprocs hStep.fvarId
           | throwError "internal error: simp closed goal unexpectedly"

From b48f168f7be884a7e718312acd26e02ccb314584 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 1 Oct 2024 16:51:55 -0500
Subject: [PATCH 26/48] feat: attempt to preserve stack alignment roof in
 AxEffect update

When an write to memory or a write to a register that is not SP is made, we update the proof of stack alignment. When a write to SP is made, we create a new mvar with the proof obligation that the new value is aligned, and store that mvar in a new `sideConditions` field.
---
 Arm/Insts/Common.lean      | 16 +++++++--
 Tactics/Sym.lean           | 58 ++++++++++--------------------
 Tactics/Sym/AxEffects.lean | 74 +++++++++++++++++++++-----------------
 3 files changed, 75 insertions(+), 73 deletions(-)

diff --git a/Arm/Insts/Common.lean b/Arm/Insts/Common.lean
index d15fff48..8668568a 100644
--- a/Arm/Insts/Common.lean
+++ b/Arm/Insts/Common.lean
@@ -221,20 +221,32 @@ def CheckSPAlignment (s : ArmState) : Prop :=
 instance : Decidable (CheckSPAlignment s) := by unfold CheckSPAlignment; infer_instance
 
 @[state_simp_rules]
-theorem CheckSPAligment_of_w_different (h : StateField.GPR 31#5 ≠ fld) :
+theorem CheckSPAligment_w_different_eq (h : StateField.GPR 31#5 ≠ fld) :
   CheckSPAlignment (w fld v s) = CheckSPAlignment s := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_w_of_ne_sp_of (h : StateField.GPR 31#5 ≠ fld) :
+    CheckSPAlignment s → CheckSPAlignment (w fld v s) := by
+  simp only [CheckSPAligment_w_different_eq h, imp_self]
+
 @[state_simp_rules]
 theorem CheckSPAligment_of_w_sp :
   CheckSPAlignment (w (StateField.GPR 31#5) v s) = (Aligned v 4) := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_w_sp_of (h : Aligned v 4) :
+    CheckSPAlignment (w (StateField.GPR 31#5) v s) := by
+  simp only [CheckSPAlignment, read_gpr, r_of_w_same, zeroExtend_eq, h]
+
 @[state_simp_rules]
-theorem CheckSPAligment_of_write_mem_bytes :
+theorem CheckSPAligment_write_mem_bytes_eq :
   CheckSPAlignment (write_mem_bytes n addr v s) = CheckSPAlignment s := by
   simp_all only [CheckSPAlignment, state_simp_rules, minimal_theory, bitvec_rules]
 
+theorem CheckSPAligment_write_mem_bytes_of :
+  CheckSPAlignment s → CheckSPAlignment (write_mem_bytes n addr v s) := by
+  simp only [CheckSPAligment_write_mem_bytes_eq, imp_self]
+
 @[state_simp_rules]
 theorem CheckSPAlignment_AddWithCarry_64_4 (st : ArmState) (y : BitVec 64) (carry_in : BitVec 1)
   (x_aligned : CheckSPAlignment st)
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index e72b97c1..ce9c9db4 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -168,7 +168,7 @@ store an `AxEffects` object with the newly added variables in the monad state
 def explodeStep (hStep : Expr) : SymM Unit :=
   withMainContext' do
     let c ← getThe SymContext
-    let mut eff ← AxEffects.fromEq hStep
+    let mut eff ← AxEffects.fromEq hStep c.effects.stackAlignmentProof?
 
     let stateExpr ← getCurrentState
     /- Assert that the initial state of the obtained `AxEffects` is equal to
@@ -183,44 +183,24 @@ def explodeStep (hStep : Expr) : SymM Unit :=
     eff ← eff.withField (← c.effects.getField .ERR).proof
 
     if let some hSp := c.effects.stackAlignmentProof? then
-      eff ← match ← eff.withStackAlignment? hSp with
-        | some newEff => pure newEff
-        | none => do
-            trace[Tactic.sym] "failed to show stack alignment"
-            -- FIXME: in future, we'd like to detect when the `sp_aligned`
-            -- hypothesis is actually necessary, and add the proof obligation
-            -- on-demand. For now, however, we over-approximate, and say that
-            -- if the original state was known to be aligned, and something
-            -- writes to the SP, then we eagerly add the obligation to proof
-            -- that the result is aligned as well.
-            -- If you don't want this obligation, simply remove the hypothesis
-            -- that the original state is aligned
-            let spEff ← eff.getField .SP
-            let subGoal ← mkFreshMVarId
-            -- subGoal.setTag <|
-            let hAligned ← do
-              let name := Name.mkSimple s!"h_{← getNextStateName}_sp_aligned"
-              mkFreshExprMVarWithId subGoal (userName := name) <|
-                mkAppN (mkConst ``Aligned) #[toExpr 64, spEff.value, toExpr 4]
-
-            trace[Tactic.sym] "created subgoal to show alignment:\n{subGoal}"
-            let subGoal? ← do
-              let (ctx, simprocs) ←
-                LNSymSimpContext
-                  (config := {failIfUnchanged := false, decide := true})
-                  (exprs := #[hSp])
-              LNSymSimp subGoal ctx simprocs
-
-            if let some subGoal := subGoal? then
-              trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
-              appendGoals [subGoal]
-            else
-              trace[Tactic.sym] "subgoal got closed by simplification"
-
-            let stackAlignmentProof? := some <|
-              mkAppN (mkConst ``CheckSPAlignment_of_r_sp_aligned)
-                #[eff.currentState, spEff.value, spEff.proof, hAligned]
-            pure { eff with stackAlignmentProof? }
+      for subGoal in eff.sideConditions do
+        trace[Tactic.sym] "attempting to discharge side-condition:\n  {subGoal}"
+        let subGoal? ← do
+          let (ctx, simprocs) ←
+            LNSymSimpContext
+              (config := {failIfUnchanged := false, decide := true})
+              (exprs := #[hSp])
+          LNSymSimp subGoal ctx simprocs
+
+        if let some subGoal := subGoal? then
+          trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
+          subGoal.setTag (.mkSimple s!"h_{← getNextStateName}_sp_aligned")
+          appendGoals [subGoal]
+        else
+          trace[Tactic.sym] "subgoal got closed by simplification"
+    else
+      appendGoals eff.sideConditions
+    eff := { eff with sideConditions := [] }
 
     -- Add new (non-)effect hyps to the context, and to the aggregation simpset
     withMainContext' <| do
diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index 8c8a6aa5..09501796 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -78,6 +78,10 @@ structure AxEffects where
   However, if SP is written to, no effort is made to prove alignment of the
   new value; the field will be set to `none` -/
   stackAlignmentProof? : Option Expr
+
+  /-- `sideContitions` are proof obligations that come up during effect
+  characterization. -/
+  sideConditions : List MVarId
   deriving Repr
 
 namespace AxEffects
@@ -141,6 +145,7 @@ def initial (state : Expr) : AxEffects where
       mkConst ``Program,
       mkApp (mkConst ``ArmState.program) state]
   stackAlignmentProof? := none
+  sideConditions := []
 
 /-! ## ToMessageData -/
 
@@ -304,6 +309,11 @@ private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
         #[eff.currentState, n, addr, val])
       eff.programProof
 
+  -- Update the stack alignment proof
+  let stackAlignmentProof? := eff.stackAlignmentProof?.map fun proof =>
+    mkAppN (mkConst ``CheckSPAligment_write_mem_bytes_of)
+      #[eff.currentState, n, addr, val, proof]
+
   -- Assemble the result
   let addWrite (e : Expr) :=
     -- `@write_mem_bytes <n> <addr> <val> <e>`
@@ -315,6 +325,7 @@ private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
     memoryEffect    := addWrite eff.memoryEffect
     memoryEffectProof
     programProof
+    stackAlignmentProof?
   }
   withTraceNode `Tactic.sym (fun _ => pure "new state") <| do
       trace[Tactic.sym] "{eff}"
@@ -398,6 +409,24 @@ private def update_w (eff : AxEffects) (fld val : Expr) :
       (mkAppN (mkConst ``w_program) #[fld, val, eff.currentState])
       eff.programProof
 
+  -- Update the stack alignment proof
+  let mut sideConditions := eff.sideConditions
+  let mut stackAlignmentProof? := eff.stackAlignmentProof?
+  if let some proof := stackAlignmentProof? then
+    if rField ≠ StateField.SP then
+      let hNeq ← mkDecideProof <|
+        mkApp3 (.const ``Ne [1])
+          (mkConst ``StateField) (toExpr StateField.SP) fld
+      stackAlignmentProof? := mkAppN (mkConst ``CheckSPAligment_w_of_ne_sp_of)
+        #[fld, eff.currentState, val, hNeq, proof]
+    else
+      let hAligned ← mkFreshExprMVar (some <|
+        mkApp3 (mkConst ``Aligned) (toExpr 64) val (toExpr 4)
+      )
+      sideConditions := hAligned.mvarId! :: sideConditions
+      stackAlignmentProof? := mkAppN (mkConst ``CheckSPAligment_w_sp_of)
+        #[val, eff.currentState, hAligned]
+
   -- Assemble the result
   let eff := { eff with
     currentState    := mkApp3 (mkConst ``w) fld val eff.currentState
@@ -406,6 +435,8 @@ private def update_w (eff : AxEffects) (fld val : Expr) :
     -- memory effects are unchanged
     memoryEffectProof
     programProof
+    stackAlignmentProof?
+    sideConditions
   }
   eff.traceCurrentState "new state"
   return eff
@@ -477,9 +508,12 @@ def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
     -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
     --    Presumably, we would *not* want to encapsulate `memoryEffect` here
     let programProof ← rewriteType eff.programProof eq
+    let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM
+      (rewriteType · eq)
 
     return { eff with
-      currentState, fields, nonEffectProof, memoryEffectProof, programProof
+      currentState, fields, nonEffectProof, memoryEffectProof, programProof,
+      stackAlignmentProof?
     }
 
 /-- Given a proof `eq : ?s = <sequence of w/write_mem to eff.currentState>`,
@@ -503,10 +537,14 @@ def updateWithEq (eff : AxEffects) (eq : Expr) : MetaM AxEffects :=
 where `?s` and `?s0` are arbitrary `ArmState`s,
 return an `AxEffect` with `?s0` as the initial state,
 the rhs of the equality as the current state,
-and the (non-)effects updated accordingly -/
-def fromEq (eq : Expr) : MetaM AxEffects := do
+and the (non-)effects updated accordingly
+
+One can optionally pass in a proof that `?s0` has a well-aligned stack pointer.
+-/
+def fromEq (eq : Expr) (stackAlignmentProof? : Option Expr := none) :
+    MetaM AxEffects := do
   let s0 ← mkFreshExprMVar mkArmState
-  let eff := initial s0
+  let eff := { initial s0 with stackAlignmentProof? }
   let eff ← eff.updateWithEq eq
   return { eff with initialState := ← instantiateMVars eff.initialState}
 
@@ -570,34 +608,6 @@ def withField (eff : AxEffects) (eq : Expr) : MetaM AxEffects := do
       let fields := eff.fields.insert field { value, proof }
       return { eff with fields }
 
-/-- Given a proof of `CheckSPAlignment <initialState>`,
-attempt to transport it to a proof of `CheckSPAlignment <currentState>`
-and store that proof in `stackAlignmentProof?`.
-
-Returns `none` if the proof failed to be transported,
-i.e., if SP was written to. -/
-def withStackAlignment? (eff : AxEffects) (spAlignment : Expr) :
-    MetaM (Option AxEffects) := do
-  let msg := m!"withInitialStackAlignment? {spAlignment}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
-    eff.traceCurrentState
-
-    let { value, proof } ← eff.getField StateField.SP
-    let expected :=
-      mkApp2 (mkConst ``r) (toExpr <| StateField.SP) eff.initialState
-    trace[Tactic.sym] "checking whether value:\n  {value}\n\
-      is syntactically equal to expected value\n  {expected}"
-    if value != expected then
-      trace[Tactic.sym] "failed to transport proof:
-        expected value to be {expected}, but found {value}"
-      return none
-
-    let stackAlignmentProof? := some <|
-      mkAppN (mkConst ``CheckSPAlignment_of_r_sp_eq)
-        #[eff.initialState, eff.currentState, proof, spAlignment]
-    trace[Tactic.sym] "constructed stackAlignmentProof: {stackAlignmentProof?}"
-    return some { eff with stackAlignmentProof? }
-
 /-! ## Composition -/
 
 /- TODO: write a function that combines two effects `left` and `right`,

From 7e08a4b431208b2709feb9e312b80a679a002d04 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 15:23:00 -0500
Subject: [PATCH 27/48] revert bad merge

---
 Benchmarks/Command.lean | 122 ++++++++++++++++++++++++++++------------
 1 file changed, 86 insertions(+), 36 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index bf605c43..44c2efd4 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -12,6 +12,17 @@ initialize
     defValue := false
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
+  registerOption `benchmark.runs {
+    defValue := (5 : Nat)
+    descr := "controls how many runs the `benchmark` command does. \
+    NOTE: this value is ignored when the `profiler` option is set to true"
+  }
+  registerOption `benchmark.profilerDir {
+    defValue := "profiles/"
+    descr := "where to put profile output files"
+  }
+  /- Shouldn't be set directly, instead, use the `benchmark` command -/
+  registerTraceClass `benchmark
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
 def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
@@ -20,50 +31,89 @@ def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
-elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
-  logInfo m!"Running {id} benchmark\n"
+/-- Adds a trace node with the `benchmark` class, but only if the profiler
+option is *not* set.
+
+We deliberately suppress benchmarking nodes when profiling, since it generally
+only adds noise
+-/
+def withBenchTraceNode (msg : MessageData) (x : CommandElabM α )
+    : CommandElabM α := do
+  if (← getBoolOption `profiler) then
+    x
+  else
+    withTraceNode `benchmark (fun _ => pure msg) x (collapsed := false)
+
+/--
+Run a benchmark for a set number of times, and report the average runtime.
 
+If the `profiler` option is set true, we run the benchmark only once, with:
+- `trace.profiler` to true, and
+- `trace.profiler.output` set based on the `benchmark.profilerDir` and the
+    id of the benchmark
+-/
+elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
+  let originalOpts ← getOptions
+  let mut n := originalOpts.getNat `benchmark.runs 5
+  let mut opts := originalOpts
+  opts := opts.setBool `benchmark true
   let stx ← `(command|
-    set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
-  let n := 5
-  let mut totalRunTime := 0
-  -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
-  -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
-  let mut totalRunTimeLog := 0
-  for i in [0:n] do
-    logInfo m!"\n\nRun {i} (out of {n}):\n"
-    let ((), _, runTime) ← withHeartbeatsAndMs <|
-      elabCommand stx
-
-    logInfo m!"Proof took {runTime / 1000}s in total"
-    totalRunTime := totalRunTime + runTime
-    totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
-
-  let avg := totalRunTime.toFloat / n.toFloat / 1000
-  let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
-  logInfo m!"\
-{id}:
-  average runtime over {n} runs:
-    {avg}s
-  geomean over {n} runs:
-    {geomean}s
-"
+  if (← getBoolOption `profiler) then
+    let outDir := (← getOptions).getString `benchmark.profilerDir
+    opts := opts.setBool `trace.profiler true
+    opts := opts.setString `trace.profiler.output s!"{outDir}/{id.getId}"
+    n := 1 -- only run once, if `profiler` is set to true
+  else
+    opts := opts.setBool `trace.benchmark true
+
+  if n = 0 then
+    return ()
+
+  -- Set options
+  modifyScope fun scope => { scope with opts }
+
+  withBenchTraceNode m!"Running {id} benchmark" <| do
+    let mut totalRunTime := 0
+    -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
+    -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
+    let mut totalRunTimeLog : Float := 0
+    for i in [0:n] do
+      let runTime ← withBenchTraceNode m!"Run {i+1} (out of {n}):" <| do
+        let ((), _, runTime) ← withHeartbeatsAndMs <|
+          elabCommand stx
+
+        trace[benchmark] m!"Proof took {runTime / 1000}s in total"
+        pure runTime
+      totalRunTime := totalRunTime + runTime
+      totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
+
+    let avg := totalRunTime.toFloat / n.toFloat / 1000
+    let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
+    trace[benchmark] m!"\
+  {id}:
+    average runtime over {n} runs:
+      {avg}s
+    geomean over {n} runs:
+      {geomean}s
+  "
+    -- Restore options
+    modifyScope fun scope => { scope with opts := originalOpts }
 
 /-- Set various options to disable linters -/
 macro "disable_linters" "in" cmd:command : command => `(command|
-  set_option linter.constructorNameAsVariable false in
-  set_option linter.deprecated false in
-  set_option linter.missingDocs false in
-  set_option linter.omit false in
-  set_option linter.suspiciousUnexpanderPatterns false in
-  set_option linter.unnecessarySimpa false in
-  set_option linter.unusedRCasesPattern false in
-  set_option linter.unusedSectionVars false in
-  set_option linter.unusedVariables false in
-  $cmd
+set_option linter.constructorNameAsVariable false in
+set_option linter.deprecated false in
+set_option linter.missingDocs false in
+set_option linter.omit false in
+set_option linter.suspiciousUnexpanderPatterns false in
+set_option linter.unnecessarySimpa false in
+set_option linter.unusedRCasesPattern false in
+set_option linter.unusedSectionVars false in
+set_option linter.unusedVariables false in
+$cmd
 )
 
 /-- The default `maxHeartbeats` setting.

From 1ef20d7295f080183cfaaf74a64eb99588e42557 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 15:24:33 -0500
Subject: [PATCH 28/48] chore: rename `withHeartbeatsAndMs` ->
 `...AndMilliseconds`

---
 Benchmarks/Command.lean | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 44c2efd4..43cd8703 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -25,7 +25,8 @@ initialize
   registerTraceClass `benchmark
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
-def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
+/-- Measure the heartbeats and time (in milliseconds) taken by `x` -/
+def withHeartbeatsAndMilliseconds (x : m α) : m (α × Nat × Nat) := do
   let start ← IO.monoMsNow
   let (a, heartbeats) ← withHeartbeats x
   let endTime ← IO.monoMsNow
@@ -82,7 +83,7 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
     let mut totalRunTimeLog : Float := 0
     for i in [0:n] do
       let runTime ← withBenchTraceNode m!"Run {i+1} (out of {n}):" <| do
-        let ((), _, runTime) ← withHeartbeatsAndMs <|
+        let ((), _, runTime) ← withHeartbeatsAndMilliseconds <|
           elabCommand stx
 
         trace[benchmark] m!"Proof took {runTime / 1000}s in total"
@@ -146,7 +147,7 @@ private def withBenchmarkAux (x : m α) (f : Nat → Nat → m Unit)  : m α :=
   if (← getBoolOption `benchmark) = false then
     x
   else
-    let (a, heartbeats, t) ← withHeartbeatsAndMs x
+    let (a, heartbeats, t) ← withHeartbeatsAndMilliseconds x
     f heartbeats t
     return a
 

From 193ba2dacd3ee35ff1f6c2590905a3343a98dac8 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 15:26:50 -0500
Subject: [PATCH 29/48] remove `profilerDir` option

The `trace.profiler.output` option has to be set when the frontend finishes the entire file. Therefore, it doesn't make sense to control it from within the benchmark elaborator: it's likely to do unexpected things if we have multiple benchmarks in the same file.

Instead, we should be setting the option from the makefile that runs the benchmark to begin with
---
 Benchmarks/Command.lean | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index 43cd8703..d27fca74 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -17,10 +17,6 @@ initialize
     descr := "controls how many runs the `benchmark` command does. \
     NOTE: this value is ignored when the `profiler` option is set to true"
   }
-  registerOption `benchmark.profilerDir {
-    defValue := "profiles/"
-    descr := "where to put profile output files"
-  }
   /- Shouldn't be set directly, instead, use the `benchmark` command -/
   registerTraceClass `benchmark
 
@@ -63,9 +59,7 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
   )
 
   if (← getBoolOption `profiler) then
-    let outDir := (← getOptions).getString `benchmark.profilerDir
     opts := opts.setBool `trace.profiler true
-    opts := opts.setString `trace.profiler.output s!"{outDir}/{id.getId}"
     n := 1 -- only run once, if `profiler` is set to true
   else
     opts := opts.setBool `trace.benchmark true

From ca101bcf4f2be679d5210ee4f762ad17e596773b Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 16:09:58 -0500
Subject: [PATCH 30/48] feat: write shell-scripts to automatically log
 benchmark/profile data to a file

---
 .gitignore           |  1 +
 Makefile             | 32 +++++++++++---------------------
 scripts/benchmark.sh | 20 ++++++++++++++++++++
 scripts/profile.sh   | 20 ++++++++++++++++++++
 4 files changed, 52 insertions(+), 21 deletions(-)
 create mode 100755 scripts/benchmark.sh
 create mode 100755 scripts/profile.sh

diff --git a/.gitignore b/.gitignore
index 6c2506f9..1e3e3f4d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@
 /lake-packages/*
 /.lake/*
 *.log
+/data/*
diff --git a/Makefile b/Makefile
index ced17c10..e3900c20 100644
--- a/Makefile
+++ b/Makefile
@@ -39,32 +39,22 @@ awslc_elf:
 cosim:
 	time -p lake exe lnsym $(VERBOSE) --num-tests $(NUM_TESTS)
 
-BENCH = $(LEAN) -Dweak.benchmark.runs=5
+BENCHMARKS = Benchmarks/SHA512_75.lean \
+	Benchmarks/SHA512_75_noKernel_noLint.lean \
+	Benchmarks/SHA512_150.lean \
+	Benchmarks/SHA512_150_noKernel_noLint.lean \
+	Benchmarks/SHA512_225.lean \
+	Benchmarks/SHA512_225_noKernel_noLint.lean \
+	Benchmarks/SHA512_400.lean \
+	Benchmarks/SHA512_400_noKernel_noLint.lean
+
 .PHONY: benchmarks
 benchmarks:
-	echo "HEAD is on `$(GIT) rev-parse --short HEAD`"
-	$(LAKE) build Benchmarks
-	$(BENCH) Benchmarks/SHA512_75.lean
-	$(BENCH) Benchmarks/SHA512_75_noKernel_noLint.lean
-	$(BENCH) Benchmarks/SHA512_150.lean
-	$(BENCH) Benchmarks/SHA512_150_noKernel_noLint.lean
-	$(BENCH) Benchmarks/SHA512_225.lean
-	$(BENCH) Benchmarks/SHA512_225_noKernel_noLint.lean
-	$(BENCH) Benchmarks/SHA512_400.lean
-	$(BENCH) Benchmarks/SHA512_400_noKernel_noLint.lean
+	./scripts/benchmark.sh $(BENCHMARKS)
 
-PROF = $(LEAN) -Dprofiler=true
 .PHONY: profile
 profile:
-	$(LAKE) build Benchmarks
-	$(PROF) Benchmarks/SHA512_75.lean
-	$(PROF) Benchmarks/SHA512_75_noKernel_noLint.lean
-	$(PROF) Benchmarks/SHA512_150.lean
-	$(PROF) Benchmarks/SHA512_150_noKernel_noLint.lean
-	$(PROF) Benchmarks/SHA512_225.lean
-	$(PROF) Benchmarks/SHA512_225_noKernel_noLint.lean
-	$(PROF) Benchmarks/SHA512_400.lean
-	$(PROF) Benchmarks/SHA512_400_noKernel_noLint.lean
+	./scripts/profile.sh $(BENCHMARKS)
 
 .PHONY: clean clean_all
 clean:
diff --git a/scripts/benchmark.sh b/scripts/benchmark.sh
new file mode 100755
index 00000000..028b2bf4
--- /dev/null
+++ b/scripts/benchmark.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+LAKE=lake
+BENCH="$LAKE env lean -Dweak.benchmark.runs=5"
+OUT="data/benchmarks"
+
+timestamp=$(date +"%Y-%m-%d_%H%M%S")
+rev=$(git rev-parse --short HEAD)
+echo "HEAD is on $rev"
+out="$OUT/${timestamp}_${rev}"
+mkdir -p "$out"
+
+$LAKE build Benchmarks
+for file in "$@"; do
+    echo
+    echo + $file
+    echo
+    base="$(basename "$file" ".lean")"
+    $BENCH $file | tee "$out/$base"
+done
diff --git a/scripts/profile.sh b/scripts/profile.sh
new file mode 100755
index 00000000..95eb5501
--- /dev/null
+++ b/scripts/profile.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+LAKE=lake
+PROF="$LAKE env lean -Dprofiler=true"
+OUT="data/profiles"
+
+timestamp=$(date +"%Y-%m-%d_%H%M%S")
+rev=$(git rev-parse --short HEAD)
+echo "HEAD is on $rev"
+out="$OUT/${timestamp}_${rev}"
+mkdir -p "$out"
+
+$LAKE build Benchmarks
+for file in "$@"; do
+    echo
+    echo + $file
+    echo
+    base="$(basename "$file" ".lean")"
+    $PROF -Dtrace.profiler.output="$out/$base" "$file.prof" | tee "$file.log"
+done

From 5c94d3705dde5ce84646a6e498f6d6087ca1251e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 16:16:06 -0500
Subject: [PATCH 31/48] fix: set profile file extension to `.json`

---
 scripts/profile.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/scripts/profile.sh b/scripts/profile.sh
index 95eb5501..ccc1433d 100755
--- a/scripts/profile.sh
+++ b/scripts/profile.sh
@@ -16,5 +16,5 @@ for file in "$@"; do
     echo + $file
     echo
     base="$(basename "$file" ".lean")"
-    $PROF -Dtrace.profiler.output="$out/$base" "$file.prof" | tee "$file.log"
+    $PROF -Dtrace.profiler.output="$out/$base.json" "$file" | tee "$base.log"
 done

From e1860adfe54d2902c1382a0f87eeed043bf85fe3 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 17:07:08 -0500
Subject: [PATCH 32/48]  Squashed commit of benchmarks-v3 / PR #204

---
 .gitignore                                 |   1 +
 Benchmarks.lean                            |   5 +
 Benchmarks/Command.lean                    | 121 ++++++++++++++-------
 Benchmarks/SHA512_150_noKernel_noLint.lean |  19 ++++
 Benchmarks/SHA512_225_noKernel_noLint.lean |  19 ++++
 Benchmarks/SHA512_400_noKernel_noLint.lean |  19 ++++
 Benchmarks/SHA512_75.lean                  |  17 +++
 Benchmarks/SHA512_75_noKernel_noLint.lean  |  19 ++++
 Makefile                                   |  17 ++-
 README.md                                  |   2 +
 lakefile.lean                              |   1 +
 scripts/benchmark.sh                       |  20 ++++
 scripts/profile.sh                         |  20 ++++
 13 files changed, 241 insertions(+), 39 deletions(-)
 create mode 100644 Benchmarks/SHA512_150_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_225_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_400_noKernel_noLint.lean
 create mode 100644 Benchmarks/SHA512_75.lean
 create mode 100644 Benchmarks/SHA512_75_noKernel_noLint.lean
 create mode 100755 scripts/benchmark.sh
 create mode 100755 scripts/profile.sh

diff --git a/.gitignore b/.gitignore
index 6c2506f9..1e3e3f4d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@
 /lake-packages/*
 /.lake/*
 *.log
+/data/*
diff --git a/Benchmarks.lean b/Benchmarks.lean
index 7b7fcc0a..cc1f1c6b 100644
--- a/Benchmarks.lean
+++ b/Benchmarks.lean
@@ -3,6 +3,11 @@ Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 Released under Apache 2.0 license as described in the file LICENSE.
 Author(s): Alex Keizer
 -/
+import Benchmarks.SHA512_75
+import Benchmarks.SHA512_75_noKernel_noLint
 import Benchmarks.SHA512_150
+import Benchmarks.SHA512_150_noKernel_noLint
 import Benchmarks.SHA512_225
+import Benchmarks.SHA512_225_noKernel_noLint
 import Benchmarks.SHA512_400
+import Benchmarks.SHA512_400_noKernel_noLint
diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index bf605c43..d27fca74 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -12,58 +12,103 @@ initialize
     defValue := false
     descr := "enables/disables benchmarking in `withBenchmark` combinator"
   }
+  registerOption `benchmark.runs {
+    defValue := (5 : Nat)
+    descr := "controls how many runs the `benchmark` command does. \
+    NOTE: this value is ignored when the `profiler` option is set to true"
+  }
+  /- Shouldn't be set directly, instead, use the `benchmark` command -/
+  registerTraceClass `benchmark
 
 variable {m} [Monad m] [MonadLiftT BaseIO m] in
-def withHeartbeatsAndMs (x : m α) : m (α × Nat × Nat) := do
+/-- Measure the heartbeats and time (in milliseconds) taken by `x` -/
+def withHeartbeatsAndMilliseconds (x : m α) : m (α × Nat × Nat) := do
   let start ← IO.monoMsNow
   let (a, heartbeats) ← withHeartbeats x
   let endTime ← IO.monoMsNow
   return ⟨a, heartbeats, endTime - start⟩
 
-elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
-  logInfo m!"Running {id} benchmark\n"
+/-- Adds a trace node with the `benchmark` class, but only if the profiler
+option is *not* set.
+
+We deliberately suppress benchmarking nodes when profiling, since it generally
+only adds noise
+-/
+def withBenchTraceNode (msg : MessageData) (x : CommandElabM α )
+    : CommandElabM α := do
+  if (← getBoolOption `profiler) then
+    x
+  else
+    withTraceNode `benchmark (fun _ => pure msg) x (collapsed := false)
 
+/--
+Run a benchmark for a set number of times, and report the average runtime.
+
+If the `profiler` option is set true, we run the benchmark only once, with:
+- `trace.profiler` to true, and
+- `trace.profiler.output` set based on the `benchmark.profilerDir` and the
+    id of the benchmark
+-/
+elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
+  let originalOpts ← getOptions
+  let mut n := originalOpts.getNat `benchmark.runs 5
+  let mut opts := originalOpts
+  opts := opts.setBool `benchmark true
   let stx ← `(command|
-    set_option benchmark true in
     example $declSig:optDeclSig $val:declVal
   )
 
-  let n := 5
-  let mut totalRunTime := 0
-  -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
-  -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
-  let mut totalRunTimeLog := 0
-  for i in [0:n] do
-    logInfo m!"\n\nRun {i} (out of {n}):\n"
-    let ((), _, runTime) ← withHeartbeatsAndMs <|
-      elabCommand stx
-
-    logInfo m!"Proof took {runTime / 1000}s in total"
-    totalRunTime := totalRunTime + runTime
-    totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
-
-  let avg := totalRunTime.toFloat / n.toFloat / 1000
-  let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
-  logInfo m!"\
-{id}:
-  average runtime over {n} runs:
-    {avg}s
-  geomean over {n} runs:
-    {geomean}s
-"
+  if (← getBoolOption `profiler) then
+    opts := opts.setBool `trace.profiler true
+    n := 1 -- only run once, if `profiler` is set to true
+  else
+    opts := opts.setBool `trace.benchmark true
+
+  if n = 0 then
+    return ()
+
+  -- Set options
+  modifyScope fun scope => { scope with opts }
+
+  withBenchTraceNode m!"Running {id} benchmark" <| do
+    let mut totalRunTime := 0
+    -- geomean = exp(log((a₁ a₂ ... aₙ)^1/n)) =
+    -- exp(1/n * (log a₁ + log a₂ + log aₙ)).
+    let mut totalRunTimeLog : Float := 0
+    for i in [0:n] do
+      let runTime ← withBenchTraceNode m!"Run {i+1} (out of {n}):" <| do
+        let ((), _, runTime) ← withHeartbeatsAndMilliseconds <|
+          elabCommand stx
+
+        trace[benchmark] m!"Proof took {runTime / 1000}s in total"
+        pure runTime
+      totalRunTime := totalRunTime + runTime
+      totalRunTimeLog := totalRunTimeLog + Float.log runTime.toFloat
+
+    let avg := totalRunTime.toFloat / n.toFloat / 1000
+    let geomean := (Float.exp (totalRunTimeLog / n.toFloat)) / 1000.0
+    trace[benchmark] m!"\
+  {id}:
+    average runtime over {n} runs:
+      {avg}s
+    geomean over {n} runs:
+      {geomean}s
+  "
+    -- Restore options
+    modifyScope fun scope => { scope with opts := originalOpts }
 
 /-- Set various options to disable linters -/
 macro "disable_linters" "in" cmd:command : command => `(command|
-  set_option linter.constructorNameAsVariable false in
-  set_option linter.deprecated false in
-  set_option linter.missingDocs false in
-  set_option linter.omit false in
-  set_option linter.suspiciousUnexpanderPatterns false in
-  set_option linter.unnecessarySimpa false in
-  set_option linter.unusedRCasesPattern false in
-  set_option linter.unusedSectionVars false in
-  set_option linter.unusedVariables false in
-  $cmd
+set_option linter.constructorNameAsVariable false in
+set_option linter.deprecated false in
+set_option linter.missingDocs false in
+set_option linter.omit false in
+set_option linter.suspiciousUnexpanderPatterns false in
+set_option linter.unnecessarySimpa false in
+set_option linter.unusedRCasesPattern false in
+set_option linter.unusedSectionVars false in
+set_option linter.unusedVariables false in
+$cmd
 )
 
 /-- The default `maxHeartbeats` setting.
@@ -96,7 +141,7 @@ private def withBenchmarkAux (x : m α) (f : Nat → Nat → m Unit)  : m α :=
   if (← getBoolOption `benchmark) = false then
     x
   else
-    let (a, heartbeats, t) ← withHeartbeatsAndMs x
+    let (a, heartbeats, t) ← withHeartbeatsAndMilliseconds x
     f heartbeats t
     return a
 
diff --git a/Benchmarks/SHA512_150_noKernel_noLint.lean b/Benchmarks/SHA512_150_noKernel_noLint.lean
new file mode 100644
index 00000000..a801bd79
--- /dev/null
+++ b/Benchmarks/SHA512_150_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_150_noKernel_noLint : SHA512Bench 150 := fun s0 _ h => by
+  intros
+  sym_n 150
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_225_noKernel_noLint.lean b/Benchmarks/SHA512_225_noKernel_noLint.lean
new file mode 100644
index 00000000..6bdc12b5
--- /dev/null
+++ b/Benchmarks/SHA512_225_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_225_noKernel_noLint : SHA512Bench 225 := fun s0 _ h => by
+  intros
+  sym_n 225
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_400_noKernel_noLint.lean b/Benchmarks/SHA512_400_noKernel_noLint.lean
new file mode 100644
index 00000000..06f90357
--- /dev/null
+++ b/Benchmarks/SHA512_400_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_400_noKernel_noLint : SHA512Bench 400 := fun s0 _ h => by
+  intros
+  sym_n 400
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_75.lean b/Benchmarks/SHA512_75.lean
new file mode 100644
index 00000000..a0ab30db
--- /dev/null
+++ b/Benchmarks/SHA512_75.lean
@@ -0,0 +1,17 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+benchmark sha512_75 : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_75_noKernel_noLint.lean b/Benchmarks/SHA512_75_noKernel_noLint.lean
new file mode 100644
index 00000000..454c4e51
--- /dev/null
+++ b/Benchmarks/SHA512_75_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_75_noKernel_noLint : SHA512Bench 75 := fun s0 _ h => by
+  intros
+  sym_n 75
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Makefile b/Makefile
index 00cf37c6..e3900c20 100644
--- a/Makefile
+++ b/Makefile
@@ -5,6 +5,8 @@
 SHELL := /bin/bash
 
 LAKE = lake
+LEAN = $(LAKE) env lean
+GIT = git
 
 NUM_TESTS?=3
 VERBOSE?=--verbose
@@ -37,9 +39,22 @@ awslc_elf:
 cosim:
 	time -p lake exe lnsym $(VERBOSE) --num-tests $(NUM_TESTS)
 
+BENCHMARKS = Benchmarks/SHA512_75.lean \
+	Benchmarks/SHA512_75_noKernel_noLint.lean \
+	Benchmarks/SHA512_150.lean \
+	Benchmarks/SHA512_150_noKernel_noLint.lean \
+	Benchmarks/SHA512_225.lean \
+	Benchmarks/SHA512_225_noKernel_noLint.lean \
+	Benchmarks/SHA512_400.lean \
+	Benchmarks/SHA512_400_noKernel_noLint.lean
+
 .PHONY: benchmarks
 benchmarks:
-	$(LAKE) build Benchmarks
+	./scripts/benchmark.sh $(BENCHMARKS)
+
+.PHONY: profile
+profile:
+	./scripts/profile.sh $(BENCHMARKS)
 
 .PHONY: clean clean_all
 clean:
diff --git a/README.md b/README.md
index 47939c18..2054fdf4 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,8 @@ native-code programs of interest.
 
 `benchmarks`: run benchmarks for the symbolic simulator.
 
+`profiler`: run a single round of each benchmark, with the profiler enabled
+
 ### Makefile variables that can be passed in at the command line
 
 `VERBOSE`: Verbose mode; prints disassembly of the instructions being
diff --git a/lakefile.lean b/lakefile.lean
index 7175ad63..66000295 100644
--- a/lakefile.lean
+++ b/lakefile.lean
@@ -33,6 +33,7 @@ lean_lib «Doc» where
   -- add library configuration options here
 
 lean_lib «Benchmarks» where
+  leanOptions := #[⟨`weak.benchmark.runs, (0 : Nat)⟩]
   -- add library configuration options here
 
 @[default_target]
diff --git a/scripts/benchmark.sh b/scripts/benchmark.sh
new file mode 100755
index 00000000..028b2bf4
--- /dev/null
+++ b/scripts/benchmark.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+LAKE=lake
+BENCH="$LAKE env lean -Dweak.benchmark.runs=5"
+OUT="data/benchmarks"
+
+timestamp=$(date +"%Y-%m-%d_%H%M%S")
+rev=$(git rev-parse --short HEAD)
+echo "HEAD is on $rev"
+out="$OUT/${timestamp}_${rev}"
+mkdir -p "$out"
+
+$LAKE build Benchmarks
+for file in "$@"; do
+    echo
+    echo + $file
+    echo
+    base="$(basename "$file" ".lean")"
+    $BENCH $file | tee "$out/$base"
+done
diff --git a/scripts/profile.sh b/scripts/profile.sh
new file mode 100755
index 00000000..ccc1433d
--- /dev/null
+++ b/scripts/profile.sh
@@ -0,0 +1,20 @@
+#!/bin/bash
+
+LAKE=lake
+PROF="$LAKE env lean -Dprofiler=true"
+OUT="data/profiles"
+
+timestamp=$(date +"%Y-%m-%d_%H%M%S")
+rev=$(git rev-parse --short HEAD)
+echo "HEAD is on $rev"
+out="$OUT/${timestamp}_${rev}"
+mkdir -p "$out"
+
+$LAKE build Benchmarks
+for file in "$@"; do
+    echo
+    echo + $file
+    echo
+    base="$(basename "$file" ".lean")"
+    $PROF -Dtrace.profiler.output="$out/$base.json" "$file" | tee "$base.log"
+done

From 32614c3b978c6dda501a160993151e02086672b5 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 17:00:44 -0500
Subject: [PATCH 33/48] chore: add more trace nodes, for better profiling data

---
 Tactics/Attr.lean        |  3 ++
 Tactics/Sym.lean         | 62 ++++++++++++++++++++--------------------
 Tactics/Sym/Context.lean | 16 +++++++----
 3 files changed, 45 insertions(+), 36 deletions(-)

diff --git a/Tactics/Attr.lean b/Tactics/Attr.lean
index e6b682e8..37026ee9 100644
--- a/Tactics/Attr.lean
+++ b/Tactics/Attr.lean
@@ -9,8 +9,11 @@ open Lean
 initialize
   -- CSE tactic's non-verbose summary logging.
   registerTraceClass `Tactic.cse.summary
+
   -- enable tracing for `sym_n` tactic and related components
   registerTraceClass `Tactic.sym
+  -- enable verbose tracing
+  registerTraceClass `Tactic.sym.debug
 
   -- enable tracing for heartbeat usage of `sym_n`
   registerTraceClass `Tactic.sym.heartbeats
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index ce9c9db4..dc4ab695 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -11,15 +11,15 @@ import Tactics.Sym.Context
 import Lean
 
 open BitVec
-open Lean Meta
-open Lean.Elab.Tactic
+open Lean
+open Lean.Meta Lean.Elab.Tactic
 
 open AxEffects SymContext
 
 /-- A wrapper around `evalTactic` that traces the passed tactic script,
 executes those tactics, and then traces the new goal state -/
 private def evalTacticAndTrace (tactic : TSyntax `tactic) : TacticM Unit :=
-  withTraceNode `Tactic.sym (fun _ => pure m!"running: {tactic}") <| do
+  withTraceNode m!"running: {tactic}" <| do
     evalTactic tactic
     trace[Tactic.sym] "new goal state:\n{← getGoals}"
 
@@ -50,7 +50,8 @@ to add a new local hypothesis in terms of `w` and `write_mem`
   `h_step : ?s' = w _ _ (w _ _ (... ?s))`
 -/
 def stepiTac (stepiEq : Expr) (hStep : Name) : SymReaderM Unit := fun ctx =>
-  withMainContext' do
+  withMainContext' <|
+  withVerboseTraceNode m!"stepiTac: {stepiEq}" <| do
     let pc := (Nat.toDigits 16 ctx.pc.toNat).asString
     --  ^^ The PC in hex
     let stepLemma := Name.str ctx.program s!"stepi_eq_0x{pc}"
@@ -89,8 +90,7 @@ Finally, we use this proof to change the type of `h_run` accordingly.
 -/
 def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
   let c ← readThe SymContext
-  let msg := m!"unfoldRun (runSteps? := {c.runSteps?})"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <|
+  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" <|
   match c.runSteps? with
     | some (_ + 1) => do
         trace[Tactic.sym] "runSteps is statically known to be non-zero, \
@@ -124,9 +124,9 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
         let subGoal ← mkFreshMVarId
         let _ ← mkFreshExprMVarWithId subGoal subGoalTy
 
-        let msg := m!"runSteps is not statically known, so attempt to prove:\
-          {subGoal}"
-        withTraceNode `Tactic.sym (fun _ => pure msg) <| subGoal.withContext <| do
+        withTraceNode m!"runSteps is not statically known, so attempt to prove:\
+          {subGoal}" <|
+        subGoal.withContext <| do
           setGoals [subGoal]
           whileTac () -- run `whileTac` to attempt to close `subGoal`
 
@@ -166,7 +166,8 @@ add the relevant hypotheses to the local context, and
 store an `AxEffects` object with the newly added variables in the monad state
 -/
 def explodeStep (hStep : Expr) : SymM Unit :=
-  withMainContext' do
+  withMainContext' <|
+  withTraceNode m!"explodeStep {hStep}" <| do
     let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep c.effects.stackAlignmentProof?
 
@@ -183,21 +184,22 @@ def explodeStep (hStep : Expr) : SymM Unit :=
     eff ← eff.withField (← c.effects.getField .ERR).proof
 
     if let some hSp := c.effects.stackAlignmentProof? then
-      for subGoal in eff.sideConditions do
-        trace[Tactic.sym] "attempting to discharge side-condition:\n  {subGoal}"
-        let subGoal? ← do
-          let (ctx, simprocs) ←
-            LNSymSimpContext
-              (config := {failIfUnchanged := false, decide := true})
-              (exprs := #[hSp])
-          LNSymSimp subGoal ctx simprocs
-
-        if let some subGoal := subGoal? then
-          trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
-          subGoal.setTag (.mkSimple s!"h_{← getNextStateName}_sp_aligned")
-          appendGoals [subGoal]
-        else
-          trace[Tactic.sym] "subgoal got closed by simplification"
+      withVerboseTraceNode m!"discharging side condiitions" <| do
+        for subGoal in eff.sideConditions do
+          trace[Tactic.sym] "attempting to discharge side-condition:\n  {subGoal}"
+          let subGoal? ← do
+            let (ctx, simprocs) ←
+              LNSymSimpContext
+                (config := {failIfUnchanged := false, decide := true})
+                (exprs := #[hSp])
+            LNSymSimp subGoal ctx simprocs
+
+          if let some subGoal := subGoal? then
+            trace[Tactic.sym] "subgoal got simplified to:\n{subGoal}"
+            subGoal.setTag (.mkSimple s!"h_{← getNextStateName}_sp_aligned")
+            appendGoals [subGoal]
+          else
+            trace[Tactic.sym] "subgoal got closed by simplification"
     else
       appendGoals eff.sideConditions
     eff := { eff with sideConditions := [] }
@@ -230,9 +232,9 @@ Symbolically simulate a single step, according the the symbolic simulation
 context `c`, returning the context for the next step in simulation. -/
 def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
   let stateNumber ← getCurrentStateNumber
-  let msg := m!"(sym1): simulating step {stateNumber}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
-    withTraceNode `Tactic.sym (fun _ => pure "verbose context") <| do
+  withTraceNode m!"(sym1): simulating step {stateNumber}" <|
+  withMainContext' do
+    withTraceNode "verbose context" <| do
       traceSymContext
       trace[Tactic.sym] "Goal state:\n {← getMainGoal}"
 
@@ -408,13 +410,11 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
 
         let goal ← subst goal hEqId
         trace[Tactic.sym] "performed subsitutition in:\n{goal}"
-
         replaceMainGoal [goal]
 
     -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
     -- added while symbolically evaluating to `simp`
-    let msg := m!"aggregating (non-)effects"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
+    withTraceNode m!"aggregating (non-)effects" <| withMainContext' do
       let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
       replaceMainGoal goal?.toList
 
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 5acdd000..92f3d8a5 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -204,17 +204,23 @@ def toMessageData (c : SymContext) : MetaM MessageData := do
   curr_state_number := {c.currentStateNumber},
   effects := {c.effects} }"
 
+section Tracing
 variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
   [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
-  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m] in
-def withSymTraceNode (msg : MessageData) (k : m α) : m α := do
-  withTraceNode `Tactic.sym (fun _ => pure msg) k
+  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
+
+def withTraceNode (msg : MessageData) (k : m α) : m α := do
+  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k
+def withVerboseTraceNode (msg : MessageData) (k : m α) : m α := do
+  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k
 
 def traceSymContext : SymM Unit :=
-  withTraceNode `Tactic.sym (fun _ => pure m!"SymContext: ") <| do
+  withTraceNode m!"SymContext: " <| do
     let m ← (← getThe SymContext).toMessageData
     trace[Tactic.sym] m
 
+end Tracing
+
 /-! ## Adding new simp theorems for aggregation -/
 
 /-- Add a set of new simp-theorems to the simp-theorems used
@@ -416,7 +422,7 @@ we create a new subgoal of this type.
 -/
 def fromMainContext (state? : Option Name) : TacticM SymContext := do
   let msg := m!"Building a `SymContext` from the local context"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
+  withTraceNode msg <| withMainContext' do
   trace[Tactic.Sym] "state? := {state?}"
   let lctx ← getLCtx
 

From 6aca3230eeff9b94dc72d879d8ec446f20e7f23c Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 17:31:17 -0500
Subject: [PATCH 34/48] chore: add tags to trace nodes

These tags show up in the profile (unlike the message itself)
---
 Tactics/Sym.lean         | 165 ++++++++++++++++++++-------------------
 Tactics/Sym/Context.lean |  70 ++++++++++-------
 2 files changed, 127 insertions(+), 108 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index dc4ab695..f4c01e0f 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -90,7 +90,7 @@ Finally, we use this proof to change the type of `h_run` accordingly.
 -/
 def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
   let c ← readThe SymContext
-  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" <|
+  withTraceNode m!"unfoldRun (runSteps? := {c.runSteps?})" (tag := "unfoldRun") <|
   match c.runSteps? with
     | some (_ + 1) => do
         trace[Tactic.sym] "runSteps is statically known to be non-zero, \
@@ -167,7 +167,7 @@ store an `AxEffects` object with the newly added variables in the monad state
 -/
 def explodeStep (hStep : Expr) : SymM Unit :=
   withMainContext' <|
-  withTraceNode m!"explodeStep {hStep}" <| do
+  withTraceNode m!"explodeStep {hStep}" (tag := "explodeStep") <| do
     let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep c.effects.stackAlignmentProof?
 
@@ -232,9 +232,9 @@ Symbolically simulate a single step, according the the symbolic simulation
 context `c`, returning the context for the next step in simulation. -/
 def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
   let stateNumber ← getCurrentStateNumber
-  withTraceNode m!"(sym1): simulating step {stateNumber}" <|
+  withTraceNode m!"(sym1): simulating step {stateNumber}" (tag:="sym1") <|
   withMainContext' do
-    withTraceNode "verbose context" <| do
+    withVerboseTraceNode "verbose context" (tag := "infoDump") <| do
       traceSymContext
       trace[Tactic.sym] "Goal state:\n {← getMainGoal}"
 
@@ -242,12 +242,14 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
     let h_step   := Lean.mkIdent (.mkSimple s!"h_step_{stateNumber + 1}")
 
     unfoldRun (fun _ => evalTacticAndTrace whileTac)
+
     -- Add new state to local context
-    let hRunId      := mkIdent <|← getHRunName
-    let nextStateId := mkIdent <|← getNextStateName
-    evalTacticAndTrace <|← `(tactic|
-      init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
-    )
+    withTraceNode "initNextStep" (tag := "initNextStep") <| do
+      let hRunId      := mkIdent <|← getHRunName
+      let nextStateId := mkIdent <|← getNextStateName
+      evalTacticAndTrace <|← `(tactic|
+        init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
+      )
 
     -- Apply relevant pre-generated `stepi` lemma
     withMainContext' <| do
@@ -262,15 +264,15 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 
       -- If we know SP is aligned, `simp` with that fact
       if let some hSp := (← getThe AxEffects).stackAlignmentProof? then
-        trace[Tactic.sym] "simplifying {hStep.toExpr} \
-          with {hSp}"
-        let some goal ← do
-            let (ctx, simprocs) ← LNSymSimpContext
-              (config := {decide := false}) (exprs := #[hSp])
-            let goal ← getMainGoal
-            LNSymSimp goal ctx simprocs hStep.fvarId
-          | throwError "internal error: simp closed goal unexpectedly"
-        replaceMainGoal [goal]
+        let msg := m!"simplifying {hStep.toExpr} with {hSp}"
+        withTraceNode msg (tag := "simplifyHStep") <| do
+          let some goal ← do
+              let (ctx, simprocs) ← LNSymSimpContext
+                (config := {decide := false}) (exprs := #[hSp])
+              let goal ← getMainGoal
+              LNSymSimp goal ctx simprocs hStep.fvarId
+            | throwError "internal error: simp closed goal unexpectedly"
+          replaceMainGoal [goal]
       else
         trace[Tactic.sym] "we have no relevant local hypotheses, \
           skipping simplification step"
@@ -293,44 +295,46 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 - log a warning and return `m`, if `runSteps? = some m` and `m < n`, or
 - return `n` unchanged, otherwise  -/
 def ensureAtMostRunSteps (n : Nat) : SymM Nat := do
-  let ctx ← getThe SymContext
-  match ctx.runSteps? with
-  | none => pure n
-  | some runSteps =>
-      if n ≤ runSteps then
-        pure n
-      else
-        withMainContext <| do
-          let hRun ← ctx.hRunDecl
-          logWarning m!"Symbolic simulation is limited to at most {runSteps} \
-            steps, because {hRun.toExpr} is of type:\n  {hRun.type}"
-          pure runSteps
-  return n
+  withVerboseTraceNode "" (tag := "ensureAtMostRunSteps") <| do
+    let ctx ← getThe SymContext
+    match ctx.runSteps? with
+    | none => pure n
+    | some runSteps =>
+        if n ≤ runSteps then
+          pure n
+        else
+          withMainContext <| do
+            let hRun ← ctx.hRunDecl
+            logWarning m!"Symbolic simulation is limited to at most {runSteps} \
+              steps, because {hRun.toExpr} is of type:\n  {hRun.type}"
+            pure runSteps
+    return n
 
 /-- Check that the step-thoerem corresponding to the current PC value exists,
 and throw a user-friendly error, pointing to `#genStepEqTheorems`,
 if it does not. -/
-def assertStepTheoremsGenerated : SymM Unit := do
-  let c ← getThe SymContext
-  let pc := c.pc.toHexWithoutLeadingZeroes
-  if !c.programInfo.instructions.contains c.pc then
-    let pcEff ← AxEffects.getFieldM .PC
-    throwError "\
-      Program {c.program} has no instruction at address {c.pc}.
-
-      We inferred this address as the program-counter from {pcEff.proof}, \
-      which has type:
-        {← inferType pcEff.proof}"
-
-  let step_thm := Name.str c.program ("stepi_eq_0x" ++ pc)
-  try
-    let _ ← getConstInfo step_thm
-  catch err =>
-    throwErrorAt err.getRef "{err.toMessageData}\n
-Did you remember to generate step theorems with:
-  #genStepEqTheorems {c.program}"
--- TODO: can we make this error ^^ into a `Try this:` suggestion that
---       automatically adds the right command just before the theorem?
+def assertStepTheoremsGenerated : SymM Unit :=
+  withVerboseTraceNode "" (tag := "assertStepTheoremsGenerated") <| do
+    let c ← getThe SymContext
+    let pc := c.pc.toHexWithoutLeadingZeroes
+    if !c.programInfo.instructions.contains c.pc then
+      let pcEff ← AxEffects.getFieldM .PC
+      throwError "\
+        Program {c.program} has no instruction at address {c.pc}.
+
+        We inferred this address as the program-counter from {pcEff.proof}, \
+        which has type:
+          {← inferType pcEff.proof}"
+
+    let step_thm := Name.str c.program ("stepi_eq_0x" ++ pc)
+    try
+      let _ ← getConstInfo step_thm
+    catch err =>
+      throwErrorAt err.getRef "{err.toMessageData}\n
+  Did you remember to generate step theorems with:
+    #genStepEqTheorems {c.program}"
+  -- TODO: can we make this error ^^ into a `Try this:` suggestion that
+  --       automatically adds the right command just before the theorem?
 
 /- used in `sym_n` tactic to specify an initial state -/
 syntax sym_at := "at" ident
@@ -390,32 +394,35 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         sym1 whileTac
 
     traceHeartbeats "symbolic simulation total"
-    let c ← getThe SymContext
-    -- Check if we can substitute the final state
-    if c.runSteps? = some 0 then
-      let msg := do
-        let hRun ← userNameToMessageData c.h_run
-        pure m!"runSteps := 0, substituting along {hRun}"
-      withTraceNode `Tactic.sym (fun _ => msg) <| withMainContext' do
-        let sfEq ← mkEq (← getCurrentState) c.finalState
-
-        let goal ← getMainGoal
-        trace[Tactic.sym] "original goal:\n{goal}"
-        let ⟨hEqId, goal⟩ ← do
-          let hRun ← SymContext.findFromUserName c.h_run
-          goal.note `this (← mkEqSymm hRun.toExpr) sfEq
-        goal.withContext <| do
-          trace[Tactic.sym] "added {← userNameToMessageData `this} of type \
-            {sfEq} in:\n{goal}"
-
-        let goal ← subst goal hEqId
-        trace[Tactic.sym] "performed subsitutition in:\n{goal}"
-        replaceMainGoal [goal]
-
-    -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
-    -- added while symbolically evaluating to `simp`
-    withTraceNode m!"aggregating (non-)effects" <| withMainContext' do
-      let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
-      replaceMainGoal goal?.toList
+    withTraceNode "Post processing" (tag := "postProccessing") <| do
+      let c ← getThe SymContext
+      -- Check if we can substitute the final state
+      if c.runSteps? = some 0 then
+        let msg := do
+          let hRun ← userNameToMessageData c.h_run
+          pure m!"runSteps := 0, substituting along {hRun}"
+        withMainContext' <|
+        withTraceNode `Tactic.sym (fun _ => msg) <| do
+          let sfEq ← mkEq (← getCurrentState) c.finalState
+
+          let goal ← getMainGoal
+          trace[Tactic.sym] "original goal:\n{goal}"
+          let ⟨hEqId, goal⟩ ← do
+            let hRun ← SymContext.findFromUserName c.h_run
+            goal.note `this (← mkEqSymm hRun.toExpr) sfEq
+          goal.withContext <| do
+            trace[Tactic.sym] "added {← userNameToMessageData `this} of type \
+              {sfEq} in:\n{goal}"
+
+          let goal ← subst goal hEqId
+          trace[Tactic.sym] "performed subsitutition in:\n{goal}"
+          replaceMainGoal [goal]
+
+      -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
+      -- added while symbolically evaluating to `simp`
+      withMainContext' <|
+      withTraceNode m!"aggregating (non-)effects" (tag := "aggregateEffects") <| do
+        let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
+        replaceMainGoal goal?.toList
 
     traceHeartbeats "final usage"
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 92f3d8a5..0d0ac09c 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -146,6 +146,26 @@ end SymM
 
 namespace SymContext
 
+/-! ## Trace Nodes -/
+section Tracing
+variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
+  [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
+  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
+
+def withTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k collapsed tag
+
+def withVerboseTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k collapsed tag
+
+end Tracing
+
 /-! ## Simple projections -/
 section
 open Lean (Ident mkIdent)
@@ -156,10 +176,11 @@ def program : Name := c.programInfo.name
 
 /-- Find the local declaration that corresponds to a given name,
 or throw an error if no local variable of that name exists -/
-def findFromUserName (name : Name) : MetaM LocalDecl := do
-  let some decl := (← getLCtx).findFromUserName? name
-    | throwError "Unknown local variable `{name}`"
-  return decl
+def findFromUserName (name : Name) : MetaM LocalDecl :=
+  withVerboseTraceNode m!"[findFromUserName] {name}" <| do
+    let some decl := (← getLCtx).findFromUserName? name
+      | throwError "Unknown local variable `{name}`"
+    return decl
 
 /-- Find the local declaration that corresponds to `c.h_run`,
 or throw an error if no local variable of that name exists -/
@@ -204,22 +225,12 @@ def toMessageData (c : SymContext) : MetaM MessageData := do
   curr_state_number := {c.currentStateNumber},
   effects := {c.effects} }"
 
-section Tracing
-variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
-  [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
-  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
-
-def withTraceNode (msg : MessageData) (k : m α) : m α := do
-  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k
-def withVerboseTraceNode (msg : MessageData) (k : m α) : m α := do
-  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k
-
 def traceSymContext : SymM Unit :=
   withTraceNode m!"SymContext: " <| do
     let m ← (← getThe SymContext).toMessageData
     trace[Tactic.sym] m
 
-end Tracing
+
 
 /-! ## Adding new simp theorems for aggregation -/
 
@@ -422,7 +433,7 @@ we create a new subgoal of this type.
 -/
 def fromMainContext (state? : Option Name) : TacticM SymContext := do
   let msg := m!"Building a `SymContext` from the local context"
-  withTraceNode msg <| withMainContext' do
+  withTraceNode msg (tag := "fromMainContext") <| withMainContext' do
   trace[Tactic.Sym] "state? := {state?}"
   let lctx ← getLCtx
 
@@ -455,16 +466,17 @@ evaluation:
   * the `currentStateNumber` is incremented
 -/
 def prepareForNextStep : SymM Unit := do
-  let pc ← do
-    let { value, ..} ← AxEffects.getFieldM .PC
-    try
-      reflectBitVecLiteral 64 value
-    catch err =>
-      trace[Tactic.sym] "failed to reflect PC: {err.toMessageData}"
-      pure <| (← getThe SymContext).pc + 4
-
-  modifyThe SymContext (fun c => { c with
-    pc
-    runSteps?   := (· - 1) <$> c.runSteps?
-    currentStateNumber := c.currentStateNumber + 1
-  })
+  withVerboseTraceNode "prepareForNextStep" (tag := "prepareForNextStep") <| do
+    let pc ← do
+      let { value, ..} ← AxEffects.getFieldM .PC
+      try
+        reflectBitVecLiteral 64 value
+      catch err =>
+        trace[Tactic.sym] "failed to reflect PC: {err.toMessageData}"
+        pure <| (← getThe SymContext).pc + 4
+
+    modifyThe SymContext (fun c => { c with
+      pc
+      runSteps?   := (· - 1) <$> c.runSteps?
+      currentStateNumber := c.currentStateNumber + 1
+    })

From df3ab2c7fb565db64ce640e0b0ad4f13453f6749 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 17:42:31 -0500
Subject: [PATCH 35/48] chore: reduce threshold in profiler

This help when profiling `sym_n`, because the threshold is applied to each step individually
---
 Benchmarks/Command.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index d27fca74..3762f788 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -60,6 +60,7 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
 
   if (← getBoolOption `profiler) then
     opts := opts.setBool `trace.profiler true
+    opts := opts.setNat `trace.profiler.threshold 1
     n := 1 -- only run once, if `profiler` is set to true
   else
     opts := opts.setBool `trace.benchmark true

From e24b5db320e33a7434313655fb75e89f723cbda0 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 2 Oct 2024 17:42:31 -0500
Subject: [PATCH 36/48] chore: reduce threshold in profiler

This help when profiling `sym_n`, because the threshold is applied to each step individually
---
 Benchmarks/Command.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Benchmarks/Command.lean b/Benchmarks/Command.lean
index d27fca74..3762f788 100644
--- a/Benchmarks/Command.lean
+++ b/Benchmarks/Command.lean
@@ -60,6 +60,7 @@ elab "benchmark" id:ident declSig:optDeclSig val:declVal : command => do
 
   if (← getBoolOption `profiler) then
     opts := opts.setBool `trace.profiler true
+    opts := opts.setNat `trace.profiler.threshold 1
     n := 1 -- only run once, if `profiler` is set to true
   else
     opts := opts.setBool `trace.benchmark true

From e5f62fedae80bb8b8ef99c698e91169dcdc7b3a0 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 14:14:05 -0500
Subject: [PATCH 37/48] feat: use `withCurrHeartbeats` to reset the heartbeat
 usage for each step in `sym_n`.

This unfortunately does not buy us much, as aggregation for 500 already seems to hit both the recursion limit and heartbeat budget, by itself.
---
 Tactics/Sym.lean | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index f4c01e0f..e89db2c0 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -231,6 +231,10 @@ elab "explode_step" h_step:term " at " state:term : tactic => withMainContext do
 Symbolically simulate a single step, according the the symbolic simulation
 context `c`, returning the context for the next step in simulation. -/
 def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
+  /- `withCurHeartbeats` sets the initial heartbeats to the current heartbeats,
+  effectively resetting our heartbeat budget back to the maximum. -/
+  withCurrHeartbeats <| do
+
   let stateNumber ← getCurrentStateNumber
   withTraceNode m!"(sym1): simulating step {stateNumber}" (tag:="sym1") <|
   withMainContext' do
@@ -394,6 +398,7 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         sym1 whileTac
 
     traceHeartbeats "symbolic simulation total"
+    withCurrHeartbeats <|
     withTraceNode "Post processing" (tag := "postProccessing") <| do
       let c ← getThe SymContext
       -- Check if we can substitute the final state
@@ -425,4 +430,4 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
         replaceMainGoal goal?.toList
 
-    traceHeartbeats "final usage"
+      traceHeartbeats "aggregation"

From fb314f41b41b60d56e33acaae3af690f0226aefe Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 14:18:13 -0500
Subject: [PATCH 38/48] fix SHA512 benchmark to compute the correct expected
 value for all number of steps

---
 Benchmarks/SHA512.lean | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/Benchmarks/SHA512.lean b/Benchmarks/SHA512.lean
index b0901b0c..685a18b7 100644
--- a/Benchmarks/SHA512.lean
+++ b/Benchmarks/SHA512.lean
@@ -15,11 +15,15 @@ namespace Benchmarks
 
 def SHA512Bench (nSteps : Nat) : Prop :=
   ∀ (s0 sf : ArmState)
-    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10)
+    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10#64)
     (_h_s0_pc : read_pc s0 = 0x1264c0#64)
     (_h_s0_err : read_err s0 = StateError.None)
     (_h_s0_sp_aligned : CheckSPAlignment s0)
     (_h_s0_program : s0.program = SHA512.program)
     (_h_run : sf = run nSteps s0),
     r StateField.ERR sf = StateError.None
-    ∧ r (.GPR 2#5) sf = BitVec.ofNat _ (9 - (nSteps / 485))
+    ∧ r (.GPR 2#5) sf = BitVec.ofNat 64 (10 - ((nSteps + 467) / 485))
+    -- / -------------------------------^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    -- | This computes the expected value of x2, taking into account that
+    -- | the loop body is 485 instructions long, and that x2 is first
+    -- | decremented after 18 instructions (485 - 18 = 467).

From 2fc7ece5d857d9c9ec81d8c3319f8463902e3987 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 14:18:13 -0500
Subject: [PATCH 39/48] fix SHA512 benchmark to compute the correct expected
 value for all number of steps

---
 Benchmarks/SHA512.lean | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/Benchmarks/SHA512.lean b/Benchmarks/SHA512.lean
index b0901b0c..685a18b7 100644
--- a/Benchmarks/SHA512.lean
+++ b/Benchmarks/SHA512.lean
@@ -15,11 +15,15 @@ namespace Benchmarks
 
 def SHA512Bench (nSteps : Nat) : Prop :=
   ∀ (s0 sf : ArmState)
-    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10)
+    (_h_s0_num_blocks : r (.GPR 2#5) s0 = 10#64)
     (_h_s0_pc : read_pc s0 = 0x1264c0#64)
     (_h_s0_err : read_err s0 = StateError.None)
     (_h_s0_sp_aligned : CheckSPAlignment s0)
     (_h_s0_program : s0.program = SHA512.program)
     (_h_run : sf = run nSteps s0),
     r StateField.ERR sf = StateError.None
-    ∧ r (.GPR 2#5) sf = BitVec.ofNat _ (9 - (nSteps / 485))
+    ∧ r (.GPR 2#5) sf = BitVec.ofNat 64 (10 - ((nSteps + 467) / 485))
+    -- / -------------------------------^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+    -- | This computes the expected value of x2, taking into account that
+    -- | the loop body is 485 instructions long, and that x2 is first
+    -- | decremented after 18 instructions (485 - 18 = 467).

From f13d7f0fcc7306c8a07a0706359eba4c027979e8 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 14:25:43 -0500
Subject: [PATCH 40/48] feat: add 50 step benchmark

Useful to profile extreme slowdowns, which cause even the 75 step benchmark to timeout
---
 Benchmarks/SHA512_50.lean                 | 17 +++++++++++++++++
 Benchmarks/SHA512_50_noKernel_noLint.lean | 19 +++++++++++++++++++
 Makefile                                  |  5 ++++-
 3 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 Benchmarks/SHA512_50.lean
 create mode 100644 Benchmarks/SHA512_50_noKernel_noLint.lean

diff --git a/Benchmarks/SHA512_50.lean b/Benchmarks/SHA512_50.lean
new file mode 100644
index 00000000..6788ec0d
--- /dev/null
+++ b/Benchmarks/SHA512_50.lean
@@ -0,0 +1,17 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+benchmark sha512_50 : SHA512Bench 50 := fun s0 _ h => by
+  intros
+  sym_n 50
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Benchmarks/SHA512_50_noKernel_noLint.lean b/Benchmarks/SHA512_50_noKernel_noLint.lean
new file mode 100644
index 00000000..5bf79044
--- /dev/null
+++ b/Benchmarks/SHA512_50_noKernel_noLint.lean
@@ -0,0 +1,19 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Tactics.Sym
+import Benchmarks.Command
+import Benchmarks.SHA512
+
+open Benchmarks
+
+disable_linters in
+set_option debug.skipKernelTC true in
+benchmark sha512_50_noKernel_noLint : SHA512Bench 50 := fun s0 _ h => by
+  intros
+  sym_n 50
+  simp only [h, bitvec_rules]
+  · exact (sorry : Aligned ..)
+  done
diff --git a/Makefile b/Makefile
index e3900c20..a06a2d1c 100644
--- a/Makefile
+++ b/Makefile
@@ -39,7 +39,10 @@ awslc_elf:
 cosim:
 	time -p lake exe lnsym $(VERBOSE) --num-tests $(NUM_TESTS)
 
-BENCHMARKS = Benchmarks/SHA512_75.lean \
+BENCHMARKS = \
+	Benchmarks/SHA512_50.lean \
+	Benchmarks/SHA512_50_noKernel_noLint.lean \
+	Benchmarks/SHA512_75.lean \
 	Benchmarks/SHA512_75_noKernel_noLint.lean \
 	Benchmarks/SHA512_150.lean \
 	Benchmarks/SHA512_150_noKernel_noLint.lean \

From 74b3b9c4bbfe9af1022bc9d158628cb606e1d439 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 14:31:51 -0500
Subject: [PATCH 41/48] feat: make benchmarks more robust against changes in
 `sym_n`

---
 Benchmarks/SHA512_150.lean                 | 3 ++-
 Benchmarks/SHA512_150_noKernel_noLint.lean | 4 ++--
 Benchmarks/SHA512_225_noKernel_noLint.lean | 4 ++--
 Benchmarks/SHA512_400.lean                 | 3 ++-
 Benchmarks/SHA512_400_noKernel_noLint.lean | 4 ++--
 Benchmarks/SHA512_50.lean                  | 4 ++--
 Benchmarks/SHA512_50_noKernel_noLint.lean  | 4 ++--
 Benchmarks/SHA512_75.lean                  | 4 ++--
 Benchmarks/SHA512_75_noKernel_noLint.lean  | 4 ++--
 9 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/Benchmarks/SHA512_150.lean b/Benchmarks/SHA512_150.lean
index da549001..0cf74cff 100644
--- a/Benchmarks/SHA512_150.lean
+++ b/Benchmarks/SHA512_150.lean
@@ -12,5 +12,6 @@ open Benchmarks
 benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_150_noKernel_noLint.lean b/Benchmarks/SHA512_150_noKernel_noLint.lean
index a801bd79..faec9e9f 100644
--- a/Benchmarks/SHA512_150_noKernel_noLint.lean
+++ b/Benchmarks/SHA512_150_noKernel_noLint.lean
@@ -14,6 +14,6 @@ set_option debug.skipKernelTC true in
 benchmark sha512_150_noKernel_noLint : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_225_noKernel_noLint.lean b/Benchmarks/SHA512_225_noKernel_noLint.lean
index 6bdc12b5..df2b7a88 100644
--- a/Benchmarks/SHA512_225_noKernel_noLint.lean
+++ b/Benchmarks/SHA512_225_noKernel_noLint.lean
@@ -14,6 +14,6 @@ set_option debug.skipKernelTC true in
 benchmark sha512_225_noKernel_noLint : SHA512Bench 225 := fun s0 _ h => by
   intros
   sym_n 225
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_400.lean b/Benchmarks/SHA512_400.lean
index 3be508ee..ae26c4e9 100644
--- a/Benchmarks/SHA512_400.lean
+++ b/Benchmarks/SHA512_400.lean
@@ -12,5 +12,6 @@ open Benchmarks
 benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_400_noKernel_noLint.lean b/Benchmarks/SHA512_400_noKernel_noLint.lean
index 06f90357..cefce28c 100644
--- a/Benchmarks/SHA512_400_noKernel_noLint.lean
+++ b/Benchmarks/SHA512_400_noKernel_noLint.lean
@@ -14,6 +14,6 @@ set_option debug.skipKernelTC true in
 benchmark sha512_400_noKernel_noLint : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_50.lean b/Benchmarks/SHA512_50.lean
index 6788ec0d..7e388869 100644
--- a/Benchmarks/SHA512_50.lean
+++ b/Benchmarks/SHA512_50.lean
@@ -12,6 +12,6 @@ open Benchmarks
 benchmark sha512_50 : SHA512Bench 50 := fun s0 _ h => by
   intros
   sym_n 50
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_50_noKernel_noLint.lean b/Benchmarks/SHA512_50_noKernel_noLint.lean
index 5bf79044..f08a6868 100644
--- a/Benchmarks/SHA512_50_noKernel_noLint.lean
+++ b/Benchmarks/SHA512_50_noKernel_noLint.lean
@@ -14,6 +14,6 @@ set_option debug.skipKernelTC true in
 benchmark sha512_50_noKernel_noLint : SHA512Bench 50 := fun s0 _ h => by
   intros
   sym_n 50
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_75.lean b/Benchmarks/SHA512_75.lean
index a0ab30db..068b06b3 100644
--- a/Benchmarks/SHA512_75.lean
+++ b/Benchmarks/SHA512_75.lean
@@ -12,6 +12,6 @@ open Benchmarks
 benchmark sha512_75 : SHA512Bench 75 := fun s0 _ h => by
   intros
   sym_n 75
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done
diff --git a/Benchmarks/SHA512_75_noKernel_noLint.lean b/Benchmarks/SHA512_75_noKernel_noLint.lean
index 454c4e51..516880c8 100644
--- a/Benchmarks/SHA512_75_noKernel_noLint.lean
+++ b/Benchmarks/SHA512_75_noKernel_noLint.lean
@@ -14,6 +14,6 @@ set_option debug.skipKernelTC true in
 benchmark sha512_75_noKernel_noLint : SHA512Bench 75 := fun s0 _ h => by
   intros
   sym_n 75
-  simp only [h, bitvec_rules]
-  · exact (sorry : Aligned ..)
+  simp (config := {failIfUnchanged := false}) only [h, bitvec_rules]
+  all_goals exact (sorry : Aligned ..)
   done

From b8fa9b53aea62cff75707a6c4bb06117ee1cfddb Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 15:11:24 -0500
Subject: [PATCH 42/48] refactor: move tracing helpers to new `Sym/Common` file

---
 Tactics/Sym.lean         |  1 +
 Tactics/Sym/Common.lean  | 30 ++++++++++++++++++++++++++++++
 Tactics/Sym/Context.lean | 24 ++----------------------
 3 files changed, 33 insertions(+), 22 deletions(-)
 create mode 100644 Tactics/Sym/Common.lean

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index e89db2c0..dd6aadb7 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -15,6 +15,7 @@ open Lean
 open Lean.Meta Lean.Elab.Tactic
 
 open AxEffects SymContext
+open Sym (withTraceNode withVerboseTraceNode)
 
 /-- A wrapper around `evalTactic` that traces the passed tactic script,
 executes those tactics, and then traces the new goal state -/
diff --git a/Tactics/Sym/Common.lean b/Tactics/Sym/Common.lean
new file mode 100644
index 00000000..78e7b823
--- /dev/null
+++ b/Tactics/Sym/Common.lean
@@ -0,0 +1,30 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer
+-/
+import Lean
+
+open Lean
+
+namespace Sym
+
+/-! ## Trace Nodes -/
+section Tracing
+variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
+  [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
+  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
+
+def withTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k collapsed tag
+
+def withVerboseTraceNode (msg : MessageData) (k : m α)
+    (collapsed : Bool := true)
+    (tag : String := "")
+    : m α := do
+  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k collapsed tag
+
+end Tracing
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index bd91a02a..f0a66184 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -9,6 +9,7 @@ import Lean.Meta
 import Arm.Exec
 import Tactics.Common
 import Tactics.Attr
+import Tactics.Sym.Common
 import Tactics.Sym.ProgramInfo
 import Tactics.Sym.AxEffects
 import Tactics.Sym.LCtxSearch
@@ -33,6 +34,7 @@ and is likely to be deprecated and removed in the near future. -/
 
 open Lean Meta Elab.Tactic
 open BitVec
+open Sym (withTraceNode withVerboseTraceNode)
 
 /-- A `SymContext` collects the names of various variables/hypotheses in
 the local context required for symbolic evaluation -/
@@ -146,26 +148,6 @@ end SymM
 
 namespace SymContext
 
-/-! ## Trace Nodes -/
-section Tracing
-variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
-  [MonadRef m] [AddMessageContext m] [MonadOptions m] {ε : Type}
-  [MonadAlwaysExcept ε m] [MonadLiftT BaseIO m]
-
-def withTraceNode (msg : MessageData) (k : m α)
-    (collapsed : Bool := true)
-    (tag : String := "")
-    : m α := do
-  Lean.withTraceNode `Tactic.sym (fun _ => pure msg) k collapsed tag
-
-def withVerboseTraceNode (msg : MessageData) (k : m α)
-    (collapsed : Bool := true)
-    (tag : String := "")
-    : m α := do
-  Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k collapsed tag
-
-end Tracing
-
 /-! ## Simple projections -/
 section
 open Lean (Ident mkIdent)
@@ -230,8 +212,6 @@ def traceSymContext : SymM Unit :=
     let m ← (← getThe SymContext).toMessageData
     trace[Tactic.sym] m
 
-
-
 /-! ## Adding new simp theorems for aggregation -/
 
 /-- Add a set of new simp-theorems to the simp-theorems used

From 366d48280eed0a5792309a96c31694fa719cc9b0 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 3 Oct 2024 15:11:34 -0500
Subject: [PATCH 43/48] feat: improve tracing for `AxEffects`

---
 Tactics/Sym/AxEffects.lean | 113 +++++++++++++++++++------------------
 1 file changed, 58 insertions(+), 55 deletions(-)

diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index 09501796..f43c6055 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -8,10 +8,12 @@ import Arm.State
 import Tactics.Common
 import Tactics.Attr
 import Tactics.Simp
+import Tactics.Sym.Common
 
 import Std.Data.HashMap
 
 open Lean Meta
+open Sym (withTraceNode withVerboseTraceNode)
 
 /-- A reflected `ArmState` field, see `AxEffects.fields` for more context -/
 structure AxEffects.FieldEffect where
@@ -170,9 +172,8 @@ instance : ToMessageData AxEffects where
     }"
 
 private def traceCurrentState (eff : AxEffects)
-    (header : MessageData := "current state") :
-    MetaM Unit :=
-  withTraceNode `Tactic.sym (fun _ => pure header) do
+    (header : MessageData := "current state") : MetaM Unit :=
+  withTraceNode header <| do
     trace[Tactic.sym] "{eff}"
 
 /-! ## Helpers -/
@@ -204,7 +205,7 @@ private def rewriteType (e eq : Expr) : MetaM Expr := do
 by constructing an application of `eff.nonEffectProof` -/
 partial def mkAppNonEffect (eff : AxEffects) (field : Expr) : MetaM Expr := do
   let msg := m!"constructing application of non-effects proof"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode msg (tag := "mkAppNonEffect") <| do
     trace[Tactic.sym] "nonEffectProof: {eff.nonEffectProof}"
 
     let nonEffectProof := mkApp eff.nonEffectProof field
@@ -223,8 +224,7 @@ partial def mkAppNonEffect (eff : AxEffects) (field : Expr) : MetaM Expr := do
 /-- Get the value for a field, if one is stored in `eff.fields`,
 or assemble an instantiation of the non-effects proof otherwise -/
 def getField (eff : AxEffects) (fld : StateField) : MetaM FieldEffect :=
-  let msg := m!"getField {fld}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode m!"getField {fld}" (tag := "getField") <| do
     eff.traceCurrentState
 
     if let some val := eff.fields.get? fld then
@@ -275,9 +275,8 @@ and all other fields are updated accordingly.
 Note that no effort is made to preserve `currentStateEq`; it is set to `none`!
 -/
 private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
-    MetaM AxEffects := do
-  trace[Tactic.sym] "adding write of {n} bytes of value {val} \
-    to memory address {addr}"
+    MetaM AxEffects :=
+  withTraceNode m!"processing: write_mem {n} {addr} {val} …" (tag := "updateWriteMem") <| do
 
   -- Update each field
   let fields ← eff.fields.toList.mapM fun ⟨fld, {value, proof}⟩ => do
@@ -327,8 +326,7 @@ private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
     programProof
     stackAlignmentProof?
   }
-  withTraceNode `Tactic.sym (fun _ => pure "new state") <| do
-      trace[Tactic.sym] "{eff}"
+  eff.traceCurrentState
   return eff
 
 /-- Execute `w <fld> <val>` against the state stored in `eff`
@@ -339,8 +337,8 @@ Note that no effort is made to preserve `currentStateEq`; it is set to `none`!
 -/
 private def update_w (eff : AxEffects) (fld val : Expr) :
     MetaM AxEffects := do
+  withTraceNode m!"processing: w {fld} {val} …" (tag := "updateWrite") <| do
   let rField ← reflectStateField fld
-  trace[Tactic.sym] "adding write of value {val} to register {rField}"
 
   -- Update all other fields
   let fields ←
@@ -454,20 +452,27 @@ private def assertIsDefEq (e expected : Expr) : MetaM Unit := do
 /-- Given an expression `e : ArmState`,
 which is a sequence of `w`/`write_mem`s to `eff.currentState`,
 return an `AxEffects` where `e` is the new `currentState`. -/
-partial def updateWithExpr (eff : AxEffects) (e : Expr) : MetaM AxEffects := do
-  let msg := m!"Updating effects with writes from: {e}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do match_expr e with
-    | write_mem_bytes n addr val e =>
-        let eff ← eff.updateWithExpr e
-        eff.update_write_mem n addr val
+private partial def updateWithExprAux (eff : AxEffects) (e : Expr) : MetaM AxEffects := do
+  match_expr e with
+  | write_mem_bytes n addr val e =>
+      let eff ← eff.updateWithExprAux e
+      eff.update_write_mem n addr val
 
-    | w field value e =>
-        let eff ← eff.updateWithExpr e
-        eff.update_w field value
+  | w field value e =>
+      let eff ← eff.updateWithExprAux e
+      eff.update_w field value
 
-    | _ =>
-        assertIsDefEq e eff.currentState
-        return eff
+  | _ =>
+      assertIsDefEq e eff.currentState
+      return eff
+
+/-- Given an expression `e : ArmState`,
+which is a sequence of `w`/`write_mem`s to `eff.currentState`,
+return an `AxEffects` where `e` is the new `currentState`. -/
+partial def updateWithExpr (eff : AxEffects) (e : Expr) : MetaM AxEffects := do
+  let msg := m!"Updating effects with writes from: {e}"
+  withTraceNode msg (tag := "updateWithExpr") <|
+    updateWithExprAux eff e
 
 /-- Given an expression `e : ArmState`,
 which is a sequence of `w`/`write_mem`s to the some state `s`,
@@ -482,55 +487,55 @@ def fromExpr (e : Expr) : MetaM AxEffects := do
   let eff ← eff.updateWithExpr e
   return { eff with initialState := ← instantiateMVars eff.initialState}
 
-
 /-- Given a proof `eq : s = <currentState>`,
 set `s` to be the new `currentState`, and update all proofs accordingly -/
 def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
     MetaM AxEffects := do
-  withTraceNode `Tactic.sym (fun _ => pure "adjusting `currenstState`") do
-    eff.traceCurrentState
+  withTraceNode m!"adjustCurrentStateWithEq" (tag := "adjustCurrentStateWithEq") do
     trace[Tactic.sym] "rewriting along {eq}"
+    eff.traceCurrentState
+
     assertHasType eq <| mkEqArmState s eff.currentState
     let eq ← mkEqSymm eq
 
     let currentState := s
 
     let fields ← eff.fields.toList.mapM fun (field, fieldEff) => do
-      withTraceNode `Tactic.sym (fun _ => pure m!"rewriting field {field}") do
+      withTraceNode m!"rewriting field {field}" (tag := "rewriteField") do
         trace[Tactic.sym] "original proof: {fieldEff.proof}"
         let proof ← rewriteType fieldEff.proof eq
         trace[Tactic.sym] "new proof: {proof}"
         pure (field, {fieldEff with proof})
     let fields := .ofList fields
 
-    let nonEffectProof ← rewriteType eff.nonEffectProof eq
-    let memoryEffectProof ← rewriteType eff.memoryEffectProof eq
-    -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
-    --    Presumably, we would *not* want to encapsulate `memoryEffect` here
-    let programProof ← rewriteType eff.programProof eq
-    let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM
-      (rewriteType · eq)
+    withTraceNode m!"rewriting other proofs" (tag := "rewriteMisc") <| do
+      let nonEffectProof ← rewriteType eff.nonEffectProof eq
+      let memoryEffectProof ← rewriteType eff.memoryEffectProof eq
+      -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
+      --    Presumably, we would *not* want to encapsulate `memoryEffect` here
+      let programProof ← rewriteType eff.programProof eq
+      let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM
+        (rewriteType · eq)
 
-    return { eff with
-      currentState, fields, nonEffectProof, memoryEffectProof, programProof,
-      stackAlignmentProof?
-    }
+      return { eff with
+        currentState, fields, nonEffectProof, memoryEffectProof, programProof,
+        stackAlignmentProof?
+      }
 
 /-- Given a proof `eq : ?s = <sequence of w/write_mem to eff.currentState>`,
 where `?s` and `?s0` are arbitrary `ArmState`s,
 return an `AxEffect` with the rhs of the equality as the current state,
 and the (non-)effects updated accordingly -/
 def updateWithEq (eff : AxEffects) (eq : Expr) : MetaM AxEffects :=
-  let msg := m!"Building effects with equality: {eq}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode m!"Building effects with equality: {eq}"
+                (tag := "updateWithEq") <| do
     let s ← mkFreshExprMVar mkArmState
     let rhs ← mkFreshExprMVar mkArmState
     assertHasType eq <| mkEqArmState s rhs
 
     let eff ← eff.updateWithExpr (← instantiateMVars rhs)
     let eff ← eff.adjustCurrentStateWithEq s eq
-    withTraceNode `Tactic.sym (fun _ => pure "new state") do
-      trace[Tactic.sym] "{eff}"
+    eff.traceCurrentState "new state"
     return eff
 
 /-- Given a proof `eq : ?s = <sequence of w/write_mem to ?s0>`,
@@ -572,8 +577,7 @@ Note: throws an error when `initialState = currentState` *and*
 the field already has a value stored, as the rewrite might produce expressions
 of unexpected types. -/
 def withField (eff : AxEffects) (eq : Expr) : MetaM AxEffects := do
-  let msg := m!"withField {eq}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode m!"withField {eq}" (tag := "withField") <| do
     eff.traceCurrentState
     let fieldE ← mkFreshExprMVar (mkConst ``StateField)
     let value ← mkFreshExprMVar none
@@ -631,8 +635,8 @@ NOTE: does not necessarily validate *which* type an expression has,
 validation will still pass if types are different to those we claim in the
 docstrings -/
 def validate (eff : AxEffects) : MetaM Unit := do
-  let msg := "validating that the axiomatic effects are well-formed"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode "validating that the axiomatic effects are well-formed"
+                (tag := "validate") <| do
     eff.traceCurrentState
 
     assertHasType eff.initialState mkArmState
@@ -667,8 +671,8 @@ that was just added to the local context -/
 def addHypothesesToLContext (eff : AxEffects) (hypPrefix : String := "h_")
     (mvar : Option MVarId := none) :
     TacticM AxEffects :=
-  let msg := m!"adding hypotheses to local context"
-  withTraceNode `Tactic.sym (fun _ => pure msg) do
+  withTraceNode m!"adding hypotheses to local context"
+                (tag := "addHypothesesToLContext") do
     eff.traceCurrentState
     let mut goal ← mvar.getDM getMainGoal
 
@@ -731,8 +735,8 @@ where
   replaceOrNote (goal : MVarId)
       (h : Name) (v : Expr) (t? : Option Expr := none) :
       MetaM (FVarId × MVarId) :=
-    let msg := m!"adding {h} to the local context"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+    withTraceNode m!"adding {h} to the local context"
+                  (tag := "replaceOrNote") <| do
       trace[Tactic.sym] "with value {v} and type {t?}"
       if let some decl := (← getLCtx).findFromUserName? h then
         let ⟨fvar, goal, _⟩ ← goal.replace decl.fvarId v t?
@@ -744,8 +748,8 @@ where
 /-- Return an array of `SimpTheorem`s of the proofs contained in
 the given `AxEffects` -/
 def toSimpTheorems (eff : AxEffects) : MetaM (Array SimpTheorem) := do
-  let msg := m!"computing SimpTheorems for (non-)effect hypotheses"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+  withTraceNode m!"computing SimpTheorems for (non-)effect hypotheses"
+                (tag := "toSimpTheorems") <| do
     let lctx ← getLCtx
     let baseName? :=
       if eff.currentState.isFVar then
@@ -757,8 +761,7 @@ def toSimpTheorems (eff : AxEffects) : MetaM (Array SimpTheorem) := do
 
     let add (thms : Array SimpTheorem) (e : Expr) (name : String)
         (prio : Nat := 1000) :=
-      let msg := m!"adding {e} with name {name}"
-      withTraceNode `Tactic.sym (fun _ => pure msg) <| do
+      withTraceNode m!"adding {e} with name {name}" <| do
         let origin : Origin :=
           if e.isFVar then
             .fvar e.fvarId!

From fa4030e181b22906fe4558d1171959f17e0bb068 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 4 Oct 2024 17:42:48 -0500
Subject: [PATCH 44/48] chore: make section names more meaningful

---
 Tactics/Sym/AxEffects.lean | 10 +++++-----
 Tactics/Sym/Context.lean   |  4 ++--
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index f43c6055..02ac980f 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -88,9 +88,9 @@ structure AxEffects where
 
 namespace AxEffects
 
-/-! ## Monad getters -/
+/-! ## Monadic getters -/
 
-section Monad
+section MonadicGetters
 variable {m} [Monad m] [MonadReaderOf AxEffects m]
 
 def getCurrentState       : m Expr := do return (← read).currentState
@@ -117,7 +117,7 @@ def getCurrentStateName : m Name := do
       | throwError "error: unknown fvar: {state}"
     return decl.userName
 
-end Monad
+end MonadicGetters
 
 /-! ## Initial Reflected State -/
 
@@ -236,7 +236,7 @@ def getField (eff : AxEffects) (fld : StateField) : MetaM FieldEffect :=
       let proof  ← eff.mkAppNonEffect (toExpr fld)
       pure { value, proof }
 
-section Monad
+section MonadicGettersAndSetters
 variable {m} [Monad m] [MonadLiftT MetaM m]
 
 variable [MonadReaderOf AxEffects m] in
@@ -264,7 +264,7 @@ This is a specialization of `setFieldEffect`. -/
 def setErrorProof (proof : Expr) : m Unit :=
   setFieldEffect .ERR { value := mkConst ``StateError.None, proof }
 
-end Monad
+end MonadicGettersAndSetters
 
 /-! ## Update a Reflected State -/
 
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index f0a66184..16f00785 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -169,7 +169,7 @@ or throw an error if no local variable of that name exists -/
 def hRunDecl : MetaM LocalDecl := do
   findFromUserName c.h_run
 
-section Monad
+section MonadicGetters
 variable {m} [Monad m] [MonadReaderOf SymContext m]
 
 def getCurrentStateNumber : m Nat := do return (← read).currentStateNumber
@@ -187,7 +187,7 @@ def getNextStateName : m Name := do
   let c ← read
   return Name.mkSimple s!"{c.state_prefix}{c.currentStateNumber + 1}"
 
-end Monad
+end MonadicGetters
 
 end
 

From 66ad4fdf969cdb4a2c2982c1b4cc9cc6c866d74a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 7 Oct 2024 15:53:12 -0500
Subject: [PATCH 45/48] refactor: extract out `MemoryEffects` structure

Co-authored-by @bollu<siddu.druid@gmail.com>
---
 Tactics/Common.lean            |  36 ++++++++++
 Tactics/Sym/AxEffects.lean     |  96 ++++++++++----------------
 Tactics/Sym/Common.lean        |   7 ++
 Tactics/Sym/MemoryEffects.lean | 122 +++++++++++++++++++++++++++++++++
 4 files changed, 201 insertions(+), 60 deletions(-)
 create mode 100644 Tactics/Sym/MemoryEffects.lean

diff --git a/Tactics/Common.lean b/Tactics/Common.lean
index 50212ce9..9763af33 100644
--- a/Tactics/Common.lean
+++ b/Tactics/Common.lean
@@ -277,6 +277,42 @@ def Lean.Expr.eqReadField? (e : Expr) : Option (Expr × Expr × Expr) := do
     | none
   some (field, state, value)
 
+/-- Return the expression for `Memory` -/
+def mkMemory : Expr := mkConst ``Memory
+
+/-! ## Expr Helpers -/
+
+/-- Throw an error if `e` is not of type `expectedType` -/
+def assertHasType (e expectedType : Expr) : MetaM Unit := do
+  let eType ← inferType e
+  if !(←isDefEq eType expectedType) then
+    throwError "{e} {← mkHasTypeButIsExpectedMsg eType expectedType}"
+
+/-- Throw an error if `e` is not def-eq to `expected` -/
+def assertIsDefEq (e expected : Expr) : MetaM Unit := do
+  if !(←isDefEq e expected) then
+    throwError "expected:\n  {expected}\nbut found:\n  {e}"
+
+/--
+Rewrites `e` via some `eq`, producing a proof `e = e'` for some `e'`.
+Rewrites with a fresh metavariable as the ambient goal.
+Fails if the rewrite produces any subgoals.
+-/
+-- source: https://github.com/leanprover-community/mathlib4/blob/b35703fe5a80f1fa74b82a2adc22f3631316a5b3/Mathlib/Lean/Expr/Basic.lean#L476-L477
+def rewrite (e eq : Expr) : MetaM Expr := do
+  let ⟨_, eq', []⟩ ← (← mkFreshExprMVar none).mvarId!.rewrite e eq
+    | throwError "Expr.rewrite may not produce subgoals."
+  return eq'
+
+/--
+Rewrites the type of `e` via some `eq`, then moves `e` into the new type via `Eq.mp`.
+Rewrites with a fresh metavariable as the ambient goal.
+Fails if the rewrite produces any subgoals.
+-/
+-- source: https://github.com/leanprover-community/mathlib4/blob/b35703fe5a80f1fa74b82a2adc22f3631316a5b3/Mathlib/Lean/Expr/Basic.lean#L476-L477
+def rewriteType (e eq : Expr) : MetaM Expr := do
+  mkEqMP (← rewrite (← inferType e) eq) e
+
 /-! ## Tracing helpers -/
 
 def traceHeartbeats (cls : Name) (header : Option String := none) :
diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index 02ac980f..4687a68f 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -9,6 +9,7 @@ import Tactics.Common
 import Tactics.Attr
 import Tactics.Simp
 import Tactics.Sym.Common
+import Tactics.Sym.MemoryEffects
 
 import Std.Data.HashMap
 
@@ -59,17 +60,8 @@ structure AxEffects where
     where `f₁, ⋯, fₙ` are the keys of `fields`
   -/
   nonEffectProof : Expr
-  /-- An expression of a (potentially empty) sequence of `write_mem`s
-  to the initial state, which describes the effects on memory.
-  See `memoryEffectProof` for more detail -/
-  memoryEffect : Expr
-  /-- An expression that contains the proof of:
-    ```lean
-    ∀ n addr,
-      read_mem_bytes n addr <currentState>
-      = read_mem_bytes n addr <memoryEffect>
-    ``` -/
-  memoryEffectProof : Expr
+  /-- The memory effects -/
+  memoryEffects : MemoryEffects
   /-- A proof that `<currentState>.program = <initialState>.program` -/
   programProof : Expr
   /-- An optional proof of `CheckSPAlignment <currentState>`.
@@ -96,8 +88,8 @@ variable {m} [Monad m] [MonadReaderOf AxEffects m]
 def getCurrentState       : m Expr := do return (← read).currentState
 def getInitialState       : m Expr := do return (← read).initialState
 def getNonEffectProof     : m Expr := do return (← read).nonEffectProof
-def getMemoryEffect       : m Expr := do return (← read).memoryEffect
-def getMemoryEffectProof  : m Expr := do return (← read).memoryEffectProof
+def getMemoryEffect       : m Expr := do return (← read).memoryEffects.effects
+def getMemoryEffectProof  : m Expr := do return (← read).memoryEffects.proof
 def getProgramProof       : m Expr := do return (← read).programProof
 
 def getStackAlignmentProof? : m (Option Expr) := do
@@ -132,15 +124,7 @@ def initial (state : Expr) : AxEffects where
     -- `fun f => rfl`
     mkLambda `f .default (mkConst ``StateField) <|
       mkEqReflArmState <| mkApp2 (mkConst ``r) (.bvar 0) state
-  memoryEffect      := state
-  memoryEffectProof :=
-    -- `fun n addr => rfl`
-    mkLambda `n .default (mkConst ``Nat) <|
-      let bv64 := mkApp (mkConst ``BitVec) (toExpr 64)
-      mkLambda `addr .default bv64 <|
-        mkApp2 (.const ``Eq.refl [1])
-          (mkApp (mkConst ``BitVec) <| mkNatMul (.bvar 1) (toExpr 8))
-          (mkApp3 (mkConst ``read_mem_bytes) (.bvar 1) (.bvar 0) state)
+  memoryEffects     := .initial state
   programProof      :=
     -- `rfl`
     mkAppN (.const ``Eq.refl [1]) #[
@@ -166,8 +150,7 @@ instance : ToMessageData AxEffects where
       currentState := {eff.currentState},
       fields := {eff.fields},
       nonEffectProof := {eff.nonEffectProof},
-      memoryEffect := {eff.memoryEffect},
-      memoryEffectProof := {eff.memoryEffectProof},
+      memoryEffects := {eff.memoryEffects},
       programProof := {eff.programProof}
     }"
 
@@ -276,7 +259,7 @@ Note that no effort is made to preserve `currentStateEq`; it is set to `none`!
 -/
 private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
     MetaM AxEffects :=
-  withTraceNode m!"processing: write_mem {n} {addr} {val} …" (tag := "updateWriteMem") <| do
+  Sym.withTraceNode m!"processing: write_mem {n} {addr} {val} …" (tag := "updateWriteMem") <| do
 
   -- Update each field
   let fields ← eff.fields.toList.mapM fun ⟨fld, {value, proof}⟩ => do
@@ -294,11 +277,10 @@ private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
     mkLambdaFVars args proof
     -- ^^ `fun f ... => Eq.trans (@r_of_write_mem_bytes f ...) <proof>`
 
-  -- Update the memory effects proof
-  let memoryEffectProof :=
-    -- `read_mem_bytes_write_mem_bytes_of_read_mem_eq <memoryEffectProof> ...`
-    mkAppN (mkConst ``read_mem_bytes_write_mem_bytes_of_read_mem_eq)
-      #[eff.currentState, eff.memoryEffect, eff.memoryEffectProof, n, addr, val]
+  -- Update the memory effects
+  let memoryEffects ←
+    eff.memoryEffects.updateWriteMem eff.currentState  n addr val
+
 
   -- Update the program proof
   let programProof ←
@@ -314,15 +296,13 @@ private def update_write_mem (eff : AxEffects) (n addr val : Expr) :
       #[eff.currentState, n, addr, val, proof]
 
   -- Assemble the result
-  let addWrite (e : Expr) :=
-    -- `@write_mem_bytes <n> <addr> <val> <e>`
-    mkApp4 (mkConst ``write_mem_bytes) n addr val e
+  let currentState := -- `@write_mem_bytes <n> <addr> <val> <currentState>`
+    mkApp4 (mkConst ``write_mem_bytes) n addr val eff.currentState
   let eff := { eff with
-    currentState    := addWrite eff.currentState
+    currentState
     fields          := .ofList fields
     nonEffectProof
-    memoryEffect    := addWrite eff.memoryEffect
-    memoryEffectProof
+    memoryEffects
     programProof
     stackAlignmentProof?
   }
@@ -337,7 +317,7 @@ Note that no effort is made to preserve `currentStateEq`; it is set to `none`!
 -/
 private def update_w (eff : AxEffects) (fld val : Expr) :
     MetaM AxEffects := do
-  withTraceNode m!"processing: w {fld} {val} …" (tag := "updateWrite") <| do
+  Sym.withTraceNode m!"processing: w {fld} {val} …" (tag := "updateWrite") <| do
   let rField ← reflectStateField fld
 
   -- Update all other fields
@@ -394,11 +374,8 @@ private def update_w (eff : AxEffects) (fld val : Expr) :
         withLocalDeclD name h_neq_type fun h_neq =>
           k (args.push h_neq) h_neq
 
-  -- Update the memory effect proof
-  let memoryEffectProof :=
-    -- `read_mem_bytes_w_of_read_mem_eq ...`
-    mkAppN (mkConst ``read_mem_bytes_w_of_read_mem_eq)
-      #[eff.currentState, eff.memoryEffect, eff.memoryEffectProof, fld, val]
+  -- Update the memory effects
+  let memoryEffects ← eff.memoryEffects.updateWrite eff.currentState fld val
 
   -- Update the program proof
   let programProof ←
@@ -430,8 +407,7 @@ private def update_w (eff : AxEffects) (fld val : Expr) :
     currentState    := mkApp3 (mkConst ``w) fld val eff.currentState
     fields := Std.HashMap.ofList fields
     nonEffectProof
-    -- memory effects are unchanged
-    memoryEffectProof
+    memoryEffects
     programProof
     stackAlignmentProof?
     sideConditions
@@ -491,7 +467,7 @@ def fromExpr (e : Expr) : MetaM AxEffects := do
 set `s` to be the new `currentState`, and update all proofs accordingly -/
 def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
     MetaM AxEffects := do
-  withTraceNode m!"adjustCurrentStateWithEq" (tag := "adjustCurrentStateWithEq") do
+  Sym.withTraceNode m!"adjustCurrentStateWithEq" (tag := "adjustCurrentStateWithEq") do
     trace[Tactic.sym] "rewriting along {eq}"
     eff.traceCurrentState
 
@@ -508,17 +484,15 @@ def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
         pure (field, {fieldEff with proof})
     let fields := .ofList fields
 
-    withTraceNode m!"rewriting other proofs" (tag := "rewriteMisc") <| do
+    Sym.withTraceNode m!"rewriting other proofs" (tag := "rewriteMisc") <| do
       let nonEffectProof ← rewriteType eff.nonEffectProof eq
-      let memoryEffectProof ← rewriteType eff.memoryEffectProof eq
-      -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
-      --    Presumably, we would *not* want to encapsulate `memoryEffect` here
+      let memoryEffects ← eff.memoryEffects.adjustCurrentStateWithEq eq
       let programProof ← rewriteType eff.programProof eq
       let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM
         (rewriteType · eq)
 
       return { eff with
-        currentState, fields, nonEffectProof, memoryEffectProof, programProof,
+        currentState, fields, nonEffectProof, memoryEffects, programProof,
         stackAlignmentProof?
       }
 
@@ -635,7 +609,7 @@ NOTE: does not necessarily validate *which* type an expression has,
 validation will still pass if types are different to those we claim in the
 docstrings -/
 def validate (eff : AxEffects) : MetaM Unit := do
-  withTraceNode "validating that the axiomatic effects are well-formed"
+  Sym.withTraceNode "validating that the axiomatic effects are well-formed"
                 (tag := "validate") <| do
     eff.traceCurrentState
 
@@ -646,13 +620,13 @@ def validate (eff : AxEffects) : MetaM Unit := do
       check fieldEff.value
       check fieldEff.proof
 
+    eff.memoryEffects.validate
     check eff.nonEffectProof
-    check eff.memoryEffect
-    check eff.memoryEffectProof
     check eff.programProof
     if let some h := eff.stackAlignmentProof? then
       check h
 
+
 /-! ## Tactic Environment -/
 section Tactic
 open Elab.Tactic
@@ -671,7 +645,7 @@ that was just added to the local context -/
 def addHypothesesToLContext (eff : AxEffects) (hypPrefix : String := "h_")
     (mvar : Option MVarId := none) :
     TacticM AxEffects :=
-  withTraceNode m!"adding hypotheses to local context"
+  Sym.withTraceNode m!"adding hypotheses to local context"
                 (tag := "addHypothesesToLContext") do
     eff.traceCurrentState
     let mut goal ← mvar.getDM getMainGoal
@@ -697,12 +671,14 @@ def addHypothesesToLContext (eff : AxEffects) (hypPrefix : String := "h_")
     let nonEffectProof := Expr.fvar nonEffectProof
     goal := goal'
 
-    trace[Tactic.sym] "adding memory effects with {eff.memoryEffectProof}"
+    trace[Tactic.sym] "adding memory effects with {eff.memoryEffects.proof}"
     let ⟨memoryEffectProof, goal'⟩ ← goal.withContext do
       let name := .mkSimple s!"{hypPrefix}memory_effects"
-      let proof := eff.memoryEffectProof
+      let proof := eff.memoryEffects.proof
       replaceOrNote goal name proof
-    let memoryEffectProof := Expr.fvar memoryEffectProof
+    let memoryEffects := { eff.memoryEffects with
+      proof := Expr.fvar memoryEffectProof
+    }
     goal := goal'
 
     trace[Tactic.sym] "adding program hypothesis with {eff.programProof}"
@@ -728,7 +704,7 @@ def addHypothesesToLContext (eff : AxEffects) (hypPrefix : String := "h_")
 
     replaceMainGoal [goal]
     return {eff with
-      fields, nonEffectProof, memoryEffectProof, programProof,
+      fields, nonEffectProof, memoryEffects, programProof,
       stackAlignmentProof?
     }
 where
@@ -748,7 +724,7 @@ where
 /-- Return an array of `SimpTheorem`s of the proofs contained in
 the given `AxEffects` -/
 def toSimpTheorems (eff : AxEffects) : MetaM (Array SimpTheorem) := do
-  withTraceNode m!"computing SimpTheorems for (non-)effect hypotheses"
+  Sym.withTraceNode m!"computing SimpTheorems for (non-)effect hypotheses"
                 (tag := "toSimpTheorems") <| do
     let lctx ← getLCtx
     let baseName? :=
@@ -782,7 +758,7 @@ def toSimpTheorems (eff : AxEffects) : MetaM (Array SimpTheorem) := do
       thms ← add thms proof s!"field_{field}" (prio := 1500)
 
     thms ← add thms eff.nonEffectProof "nonEffectProof"
-    thms ← add thms eff.memoryEffectProof "memoryEffectProof"
+    thms ← add thms eff.memoryEffects.proof "memoryEffectProof"
     thms ← add thms eff.programProof "programProof"
     if let some stackAlignmentProof := eff.stackAlignmentProof? then
       thms ← add thms stackAlignmentProof "stackAlignmentProof"
diff --git a/Tactics/Sym/Common.lean b/Tactics/Sym/Common.lean
index 78e7b823..1c571cb7 100644
--- a/Tactics/Sym/Common.lean
+++ b/Tactics/Sym/Common.lean
@@ -27,4 +27,11 @@ def withVerboseTraceNode (msg : MessageData) (k : m α)
     : m α := do
   Lean.withTraceNode `Tactic.sym.verbose (fun _ => pure msg) k collapsed tag
 
+/-- Create a trace note that folds `header` with `(NOTE: can be large)`,
+and prints `msg` under such a trace node.
+-/
+def traceLargeMsg (header : MessageData) (msg : MessageData) : MetaM Unit :=
+    withTraceNode m!"{header} (NOTE: can be large)" do
+      trace[Tactic.sym] msg
+
 end Tracing
diff --git a/Tactics/Sym/MemoryEffects.lean b/Tactics/Sym/MemoryEffects.lean
new file mode 100644
index 00000000..f71e7bc4
--- /dev/null
+++ b/Tactics/Sym/MemoryEffects.lean
@@ -0,0 +1,122 @@
+/-
+Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Author(s): Alex Keizer, Siddharth Bhat
+-/
+
+import Arm.State
+import Tactics.Common
+import Tactics.Attr
+import Tactics.Simp
+import Tactics.Sym.Common
+
+import Std.Data.HashMap
+
+open Lean Meta
+
+structure MemoryEffects where
+  /-- An expression of a (potentially empty) sequence of `write_mem`s
+  to the initial state, which describes the effects on memory.
+  See `memoryEffectProof` for more detail -/
+  effects : Expr
+  /-- An expression that contains the proof of:
+    ```lean
+    ∀ n addr,
+      read_mem_bytes n addr <currentState>
+      = read_mem_bytes n addr <all memory effects so far, e.g. (write_mem_bytes n addr val s₀)>
+    ``` -/
+  proof : Expr
+deriving Repr
+
+instance : ToMessageData MemoryEffects where
+  toMessageData eff :=
+    m!"\
+    \{ effects := {eff.effects},
+      proof := {eff.proof
+    }"
+
+namespace MemoryEffects
+
+/-! ## Initial Reflected State -/
+
+/-- An initial `MemoryEffects`, representing no memory changes to the
+initial `state` -/
+def initial (state : Expr) : MemoryEffects where
+  effects := state
+  proof :=
+    -- `fun n addr => rfl`
+    mkLambda `n .default (mkConst ``Nat) <|
+      let bv64 := mkApp (mkConst ``BitVec) (toExpr 64)
+      mkLambda `addr .default bv64 <|
+        mkApp2 (.const ``Eq.refl [1])
+          (mkApp (mkConst ``BitVec) <| mkNatMul (.bvar 1) (toExpr 8))
+          (mkApp3 (mkConst ``read_mem_bytes) (.bvar 1) (.bvar 0) state)
+
+/-- Update the memory effects with a memory write -/
+def updateWriteMem (eff : MemoryEffects) (currentState : Expr)
+    (n addr val : Expr) :
+    MetaM MemoryEffects := do
+  let effects := mkApp4 (mkConst ``write_mem_bytes) n addr val eff.effects
+  let proof :=
+    -- `read_mem_bytes_write_mem_bytes_of_read_mem_eq <memoryEffectProof> ...`
+    mkAppN (mkConst ``read_mem_bytes_write_mem_bytes_of_read_mem_eq)
+      #[currentState, eff.effects, eff.proof, n, addr, val]
+  return { effects, proof }
+
+/-- Update the memory effects with a register write.
+
+This doesn't change the actual effect, but since the `currentState` has changed,
+we need to update proofs -/
+def updateWrite (eff : MemoryEffects) (currentState : Expr)
+    (fld val : Expr) :
+    MetaM MemoryEffects := do
+  let proof := -- `read_mem_bytes_w_of_read_mem_eq ...`
+    mkAppN (mkConst ``read_mem_bytes_w_of_read_mem_eq)
+      #[currentState, eff.effects, eff.proof, fld, val]
+  return { eff with proof }
+
+/-- Transport all proofs along an equality `eq : <currentState> = s`,
+so that `s` is the new `currentState` -/
+def adjustCurrentStateWithEq (eff : MemoryEffects) (eq : Expr) :
+    MetaM MemoryEffects := do
+  let proof ← rewriteType eff.proof eq
+    -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
+    --    Presumably, we would *not* want to encapsulate `memoryEffect` here
+  return { eff with proof }
+
+/-- Convert a `MemoryEffects` into a `MessageData` for logging. -/
+def toMessageData (eff : MemoryEffects) : MetaM MessageData := do
+  let out := m!"effects: {eff.effects}"
+  return out
+
+/-- Trace the current state of `MemoryEffects`. -/
+def traceCurrentState (eff : MemoryEffects) : MetaM Unit := do
+  Sym.traceLargeMsg "memoryEffects" (← eff.toMessageData)
+
+
+
+/-- type check all expressions stored in `eff`,
+throwing an error if one is not type-correct.
+
+In principle, the various `MemoryEffects` definitions should return only
+well-formed expressions, making `validate` a no-op.
+In practice, however, running `validate` is helpful for catching bugs in those
+definitions. Do note that typechecking might be a bit expensive, so we generally
+only call `validate` while debugging, not during normal execution.
+See also the `Tactic.sym.debug` option, which controls whether `validate` is
+called for each step of the `sym_n` tactic.
+
+NOTE: does not necessarily validate *which* type an expression has,
+validation will still pass if types are different to those we claim in the
+docstrings
+-/
+def validate (eff : MemoryEffects) : MetaM Unit := do
+  let msg := "validating that the axiomatic effects are well-formed"
+  Sym.withTraceNode msg do
+    eff.traceCurrentState
+    check eff.effects
+    assertHasType eff.effects mkMemory
+
+    check eff.proof
+
+end MemoryEffects

From 077b5dac713680266f7e183604bedb2344041142 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 7 Oct 2024 15:58:12 -0500
Subject: [PATCH 46/48] fix `MemoryEffects.validate` to expect the right type

---
 Tactics/Sym/MemoryEffects.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tactics/Sym/MemoryEffects.lean b/Tactics/Sym/MemoryEffects.lean
index f71e7bc4..cbebc4c6 100644
--- a/Tactics/Sym/MemoryEffects.lean
+++ b/Tactics/Sym/MemoryEffects.lean
@@ -115,7 +115,7 @@ def validate (eff : MemoryEffects) : MetaM Unit := do
   Sym.withTraceNode msg do
     eff.traceCurrentState
     check eff.effects
-    assertHasType eff.effects mkMemory
+    assertHasType eff.effects mkArmState
 
     check eff.proof
 

From 40cb6f2e1e5cc9bdef4571245333e745c6ff3c4e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 7 Oct 2024 16:01:06 -0500
Subject: [PATCH 47/48] feat: integrate simp_mem with sym_n

Co-authored-by @bollu<siddu.druid@gmail.com>
---
 Tactics/Sym.lean | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index dd6aadb7..fb31ee54 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -5,6 +5,7 @@ Author(s): Shilpi Goel, Alex Keizer
 -/
 import Arm.Exec
 import Arm.Memory.MemoryProofs
+import Arm.Memory.SeparateAutomation
 import Tactics.FetchAndDecode
 import Tactics.Sym.Context
 
@@ -296,6 +297,41 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 
       traceHeartbeats
 
+/-! ## simp_mem -/
+
+/-- TODO(@bollu): implement `simpMem` to simplify the
+reads from memory in the given target.
+We *do not* want to iterate over all the hypotheses at this point,
+since that would be way too expensive.
+-/
+def simpMem (_c : SymContext) (mvar : MVarId) : SymM MVarId := do
+  /-
+  This should actually allow us to not need the follow processing in `Proofs/Popcount32`:
+  simp only [←Memory.mem_eq_iff_read_mem_bytes_eq] at *
+  simp only [memory_rules] at *
+  intro n addr h_separate
+  -- (TODO @alex/@bollu) Can we hope to make this shorter after the marriage
+  -- of `sym_n` and `simp_mem`?
+  simp_mem
+  sym_aggregate
+  simp only [*, bitvec_rules]
+  sym_aggregate
+  -/
+  let gs ← getGoals
+  setGoals [mvar]
+  evalTactic (← `(tactic| try simp_mem))
+  let gs' ← getGoals
+  match gs' with
+  | [g'] =>
+     setGoals gs
+     return g'
+  | [] => do
+    throwError "expected `simp_mem` to produce exactly one goal, but it was too clever and closed all goals."
+  | _ => do
+    throwError "expected `simp_mem` to produce only a single goal, but it produced more than one goal: {gs'}"
+
+/-! ## sym_n preprocessing -/
+
 /-- `ensureLessThanRunSteps n` will
 - log a warning and return `m`, if `runSteps? = some m` and `m < n`, or
 - return `n` unchanged, otherwise  -/
@@ -341,6 +377,8 @@ def assertStepTheoremsGenerated : SymM Unit :=
   -- TODO: can we make this error ^^ into a `Try this:` suggestion that
   --       automatically adds the right command just before the theorem?
 
+/-! ## sym_n -/
+
 /- used in `sym_n` tactic to specify an initial state -/
 syntax sym_at := "at" ident
 
@@ -429,6 +467,7 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
       withMainContext' <|
       withTraceNode m!"aggregating (non-)effects" (tag := "aggregateEffects") <| do
         let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
+        let goal? ← goal?.mapM (simpMem c)
         replaceMainGoal goal?.toList
 
       traceHeartbeats "aggregation"

From 711cf60a37aa97114182a33c2ab9a6d9366f821d Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 7 Oct 2024 16:08:04 -0500
Subject: [PATCH 48/48] remove helpers that were moved to common file

---
 Tactics/Sym/AxEffects.lean | 32 --------------------------------
 1 file changed, 32 deletions(-)

diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index 4687a68f..24971f43 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -161,28 +161,6 @@ private def traceCurrentState (eff : AxEffects)
 
 /-! ## Helpers -/
 
-/--
-Rewrites `e` via some `eq`, producing a proof `e = e'` for some `e'`.
-
-Rewrites with a fresh metavariable as the ambient goal.
-Fails if the rewrite produces any subgoals.
--/
--- source: https://github.com/leanprover-community/mathlib4/blob/b35703fe5a80f1fa74b82a2adc22f3631316a5b3/Mathlib/Lean/Expr/Basic.lean#L476-L477
-private def rewrite (e eq : Expr) : MetaM Expr := do
-  let ⟨_, eq', []⟩ ← (← mkFreshExprMVar none).mvarId!.rewrite e eq
-    | throwError "Expr.rewrite may not produce subgoals."
-  return eq'
-
-/--
-Rewrites the type of `e` via some `eq`, then moves `e` into the new type via `Eq.mp`.
-
-Rewrites with a fresh metavariable as the ambient goal.
-Fails if the rewrite produces any subgoals.
--/
--- source: https://github.com/leanprover-community/mathlib4/blob/b35703fe5a80f1fa74b82a2adc22f3631316a5b3/Mathlib/Lean/Expr/Basic.lean#L476-L477
-private def rewriteType (e eq : Expr) : MetaM Expr := do
-  mkEqMP (← rewrite (← inferType e) eq) e
-
 /-- Given a `field` such that `fields ∉ eff.fields`, return a proof of
   `r field <currentState> = r field <initialState>`
 by constructing an application of `eff.nonEffectProof` -/
@@ -415,16 +393,6 @@ private def update_w (eff : AxEffects) (fld val : Expr) :
   eff.traceCurrentState "new state"
   return eff
 
-/-- Throw an error if `e` is not of type `expectedType` -/
-private def assertHasType (e expectedType : Expr) : MetaM Unit := do
-  let eType ← inferType e
-  if !(←isDefEq eType expectedType) then
-    throwError "{e} {← mkHasTypeButIsExpectedMsg eType expectedType}"
-
-private def assertIsDefEq (e expected : Expr) : MetaM Unit := do
-  if !(←isDefEq e expected) then
-    throwError "expected:\n  {expected}\nbut found:\n  {e}"
-
 /-- Given an expression `e : ArmState`,
 which is a sequence of `w`/`write_mem`s to `eff.currentState`,
 return an `AxEffects` where `e` is the new `currentState`. -/