From 73b588ef39427812dfe36a73c08f7a004b24950a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 20 Sep 2024 12:07:17 -0500
Subject: [PATCH 01/18] refactor `sym_n` implementation to use a state monad
 for the `SymContext`

---
 Tactics/Common.lean            |  9 ++++
 Tactics/Reflect/AxEffects.lean | 17 ++++++
 Tactics/Sym.lean               | 97 ++++++++++++++++++----------------
 Tactics/SymContext.lean        | 51 +++++++++++++++++-
 4 files changed, 127 insertions(+), 47 deletions(-)

diff --git a/Tactics/Common.lean b/Tactics/Common.lean
index 6bd98f26..5eaa5a93 100644
--- a/Tactics/Common.lean
+++ b/Tactics/Common.lean
@@ -269,3 +269,12 @@ def traceHeartbeats (cls : Name) (header : Option String := none) :
   let percent ← heartbeatsPercent
   trace cls fun _ =>
     m!"{header}used {heartbeats} heartbeats ({percent}% of maximum)"
+
+/-! ## `withMainContext'` -/
+
+variable {m} [Monad m] [MonadLiftT TacticM m] [MonadControlT MetaM m]
+/-- Execute `x` using the main goal local context and instances.
+
+Unlike the standard `withMainContext`, `x` may live in a generic monad `m`. -/
+def withMainContext' (x : m α) : m α := do
+  (← getMainGoal).withContext x
diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 442e0b74..8f6cc770 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -82,6 +82,23 @@ structure AxEffects where
 
 namespace AxEffects
 
+/-! ## Monad getters -/
+
+section Monad
+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 getProgramProof       : m Expr := do return (← read).programProof
+
+def getStackAlignmentProof? : m (Option Expr) := do
+  return (← read).stackAlignmentProof?
+
+end Monad
+
 /-! ## Initial Reflected State -/
 
 /-- An initial `AxEffects` state which has no writes.
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 8a7579bb..8ea70b42 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -14,6 +14,8 @@ open BitVec
 open Lean Meta
 open 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 :=
@@ -47,20 +49,20 @@ section stepiTac
 to obtain a new local hypothesis in terms of `w` and `write_mem`
   `h_step : ?s' = w _ _ (w _ _ (... ?s))`
 -/
-def stepiTac (stepi_eq h_step : Ident) (ctx : SymContext)
-  : TacticM Unit := withMainContext do
-  let pc := (Nat.toDigits 16 ctx.pc.toNat).asString
-  --  ^^ The PC in hex
-  let step_lemma := mkIdent <| Name.str ctx.program s!"stepi_eq_0x{pc}"
-
-  evalTacticAndTrace <|← `(tactic| (
-    have $h_step :=
-      _root_.Eq.trans (Eq.symm $stepi_eq)
-        ($step_lemma:ident
-          $ctx.h_program_ident:ident
-          $ctx.h_pc_ident:ident
-          $ctx.h_err_ident:ident)
-  ))
+def stepiTac (stepi_eq h_step : Ident) : SymReaderM Unit := fun ctx =>
+  withMainContext do
+    let pc := (Nat.toDigits 16 ctx.pc.toNat).asString
+    --  ^^ The PC in hex
+    let step_lemma := mkIdent <| Name.str ctx.program s!"stepi_eq_0x{pc}"
+
+    evalTacticAndTrace <|← `(tactic| (
+      have $h_step :=
+        _root_.Eq.trans (Eq.symm $stepi_eq)
+          ($step_lemma:ident
+            $ctx.h_program_ident:ident
+            $ctx.h_pc_ident:ident
+            $ctx.h_err_ident:ident)
+    ))
 
 end stepiTac
 
@@ -80,8 +82,7 @@ for some metavariable `?runSteps`, then create the proof obligation
 `?runSteps = _ + 1`, and attempt to close it using `whileTac`.
 Finally, we use this proof to change the type of `h_run` accordingly.
 -/
-def unfoldRun (c : SymContext) (whileTac : Unit → TacticM Unit) :
-    TacticM Unit :=
+def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
   let msg := m!"unfoldRun (runSteps? := {c.runSteps?})"
   withTraceNode `Tactic.sym (fun _ => pure msg) <|
   match c.runSteps? with
@@ -160,11 +161,12 @@ reads from `s{i+1}`.
 Return the context for the next step (see `SymContext.next`), where
 we attempt to determine the new PC by reflecting the obtained effects,
 falling back to incrementing the PC if reflection failed. -/
-def explodeStep (c : SymContext) (hStep : Expr) : TacticM SymContext :=
-  withMainContext do
+def explodeStep (hStep : Expr) : SymM Unit :=
+  withMainContext' do
+    let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep
 
-    let stateExpr ← c.stateExpr
+    let stateExpr ← getCurrentState
     /- Assert that the initial state of the obtained `AxEffects` is equal to
     the state tracked by `c`.
     This will catch and throw an error if the semantics of the current
@@ -198,12 +200,11 @@ def explodeStep (c : SymContext) (hStep : Expr) : TacticM SymContext :=
             let subGoal ← mkFreshMVarId
             -- subGoal.setTag <|
             let hAligned ← do
-              let name := Name.mkSimple s!"h_{c.next_state}_sp_aligned"
+              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
@@ -223,15 +224,15 @@ def explodeStep (c : SymContext) (hStep : Expr) : TacticM SymContext :=
             pure { eff with stackAlignmentProof? }
 
     -- Add new (non-)effect hyps to the context
-    let simpThms ← withMainContext <| do
+    let simpThms ← withMainContext' <| do
       if ←(getBoolOption `Tactic.sym.debug) then
         eff.validate
 
-      let eff ← eff.addHypothesesToLContext s!"h_{c.next_state}_"
+      let eff ← eff.addHypothesesToLContext s!"h_{← getNextStateName}_"
       withMainContext <| eff.toSimpTheorems
 
     -- Add the new (non-)effect hyps to the aggregation simp context
-    let c := c.addSimpTheorems simpThms
+    modifyThe SymContext (·.addSimpTheorems simpThms)
 
     -- Attempt to reflect the new PC
     let nextPc ← eff.getField .PC
@@ -245,7 +246,7 @@ def explodeStep (c : SymContext) (hStep : Expr) : TacticM SymContext :=
         {err.toMessageData}"
       pure none
 
-    return c.next nextPc?
+    modifyThe SymContext (·.next nextPc?)
 
 /-- A tactic wrapper around `explodeStep`.
 Note the use of `SymContext.fromLocalContext`,
@@ -257,39 +258,41 @@ elab "explode_step" h_step:term " at " state:term : tactic => withMainContext do
     | throwError "Expected fvar, found {state}"
   let stateDecl := (← getLCtx).get! stateFVar
   let c ← SymContext.fromLocalContext (some stateDecl.userName)
-
-  let _ ← explodeStep c hStep
-
+  let _ ← SymM.run c <| explodeStep hStep
 
 /--
 Symbolically simulate a single step, according the the symbolic simulation
 context `c`, returning the context for the next step in simulation. -/
-def sym1 (c : SymContext) (whileTac : TSyntax `tactic) : TacticM SymContext :=
-  let msg := m!"(sym1): simulating step {c.curr_state_number}"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext do
+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
-      trace[Tactic.sym] "SymContext:\n{← c.toMessageData}"
+      traceSymContext
       trace[Tactic.sym] "Goal state:\n {← getMainGoal}"
 
-    let stepi_eq := Lean.mkIdent (.mkSimple s!"stepi_{c.state}")
-    let h_step   := Lean.mkIdent (.mkSimple s!"h_step_{c.curr_state_number + 1}")
+    let stepi_eq := Lean.mkIdent (.mkSimple s!"stepi_{← getCurrentStateName}")
+    let h_step   := Lean.mkIdent (.mkSimple s!"h_step_{stateNumber + 1}")
 
-    unfoldRun c (fun _ => evalTacticAndTrace whileTac)
+    unfoldRun (fun _ => evalTacticAndTrace whileTac)
     -- Add new state to local context
+    let hRunId      := mkIdent <|← getHRunName
+    let nextStateId := mkIdent <|← getNextStateName
     evalTacticAndTrace <|← `(tactic|
-      init_next_step $c.h_run_ident:ident $stepi_eq:ident $c.next_state_ident:ident
+      init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
     )
 
     -- Apply relevant pre-generated `stepi` lemma
-    stepiTac stepi_eq h_step c
+    stepiTac stepi_eq h_step
 
     -- WORKAROUND: eventually we'd like to eagerly simp away `if`s in the
     -- pre-generation of instruction semantics. For now, though, we keep a
     -- `simp` here
-    withMainContext <| do
+    withMainContext' <| do
       let hStep ← SymContext.findFromUserName h_step.getId
       let lctx ← getLCtx
-      let decls := (c.h_sp?.bind lctx.findFromUserName?).toArray
+      let decls := (← getThe SymContext).h_sp?.bind lctx.findFromUserName?
+      let decls := decls.toArray
       -- If we know SP is aligned, `simp` with that fact
 
       if !decls.isEmpty then
@@ -309,18 +312,17 @@ def sym1 (c : SymContext) (whileTac : TSyntax `tactic) : TacticM SymContext :=
           skipping simplification step"
 
     -- Prepare `h_program`,`h_err`,`h_pc`, etc. for next state
-    withMainContext <| do
+    withMainContext' <| do
       let hStep ← SymContext.findFromUserName h_step.getId
       -- ^^ we can't reuse `hStep` from before, since its fvarId might've been
       --    changed by `simp`
-      let c ← explodeStep c hStep.toExpr
+      explodeStep hStep.toExpr
 
       let goal ← getMainGoal
       let goal ← goal.clear hStep.fvarId
       replaceMainGoal [goal]
 
       traceHeartbeats
-      return c
 
 /- used in `sym_n` tactic to specify an initial state -/
 syntax sym_at := "at" ident
@@ -368,11 +370,13 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         omega;
         ))
 
-  Lean.Elab.Tactic.withMainContext <| do
-    let mut c ← SymContext.fromLocalContext s
-    c ← c.addGoalsForMissingHypotheses
+  let c ← withMainContext <| do
+    let c ← SymContext.fromLocalContext s
+    let c ← c.addGoalsForMissingHypotheses
     c.canonicalizeHypothesisTypes
+    pure c
 
+  let _ ← withMainContext <| SymM.run c <| do
     -- Check that we are not asked to simulate more steps than available
     let n ← do
       let n := n.getNat
@@ -400,9 +404,10 @@ Did you remember to generate step theorems with:
 
     -- The main loop
     for _ in List.range n do
-      c ← sym1 c whileTac
+      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
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 06e744f8..88d90d21 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -100,6 +100,24 @@ structure SymContext where
   to determine the name of the next state variable that is added by `sym` -/
   curr_state_number : Nat := 0
 
+/-! ## Monad -/
+
+/-- `SymM` is a wrapper around `TacticM` with a mutable `SymContext` state -/
+abbrev SymM := StateT SymContext TacticM
+
+/-- `SymM` is a wrapper around `TacticM` with a read-only `SymContext` state -/
+abbrev SymReaderM := ReaderT SymContext TacticM
+
+namespace SymM
+
+def run (ctx : SymContext) (k : SymM α) : TacticM (α × SymContext) :=
+  StateT.run k ctx
+
+instance : MonadLift SymReaderM SymM where
+  monadLift x c := do return (←x c, c)
+
+end SymM
+
 namespace SymContext
 
 /-! ## Simple projections -/
@@ -142,9 +160,35 @@ or throw an error if no local variable of that name exists -/
 def hRunDecl : MetaM LocalDecl := do
   findFromUserName c.h_run
 
+section Monad
+variable {m} [Monad m] [MonadReaderOf SymContext m]
+
+def getCurrentStateNumber : m Nat := do return (← read).curr_state_number
+
+/-- Retrieve the name of the current state -/
+def getCurrentStateName : m Name := do
+  return (← read).state
+
+/-- Retrieve an expression for the current state,
+or throw an error if no local variable of that name exists -/
+def getCurrentState [MonadLiftT MetaM m] : m Expr := do
+  (← read).stateExpr
+
+/-- Return the name of the hypothesis
+  `h_run : <finalState> = run <runSteps> <initialState>` -/
+def getHRunName : m Name := do return (← read).h_run
+
+/-- Retrieve the name for the next state
+
+NOTE: does not increment the state;
+consecutive calls to `getNextStateName` will give the same name -/
+def getNextStateName : m Name := do return (← read).next_state
+
+end Monad
+
 end
 
-/-! ## `ToMessageData` instance -/
+/-! ## `ToMessageData` instance and tracing -/
 
 /-- Convert a `SymContext` to `MessageData` for tracing.
 This is not a `ToMessageData` instance because we need access to `MetaM` -/
@@ -165,6 +209,11 @@ def toMessageData (c : SymContext) : MetaM MessageData := do
   state_prefix := {c.state_prefix},
   curr_state_number := {c.curr_state_number} }"
 
+def traceSymContext : SymM Unit :=
+  withTraceNode `Tactic.sym (fun _ => pure m!"SymContext: ") <| do
+    let m ← (← getThe SymContext).toMessageData
+    trace[Tactic.sym] m
+
 /-! ## Creating initial contexts -/
 
 /-- Infer `state_prefix` and `curr_state_number` from the `state` name

From 187d7ce683afab343ab6223fa65d4a5d17ccc13d Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 11:56:45 -0500
Subject: [PATCH 02/18] feat: add `AxEffects` field to SymContext

---
 Tactics/Common.lean     |  4 ++--
 Tactics/Sym.lean        |  3 ++-
 Tactics/SymContext.lean | 34 ++++++++++++++++++++++++++++++++--
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/Tactics/Common.lean b/Tactics/Common.lean
index 5eaa5a93..129ec6c6 100644
--- a/Tactics/Common.lean
+++ b/Tactics/Common.lean
@@ -223,9 +223,9 @@ def findLocalDeclOfType? (expectedType : Expr) : MetaM (Option LocalDecl) := do
     --    the local context, so we can safely pass it to `get!`
 
 def findLocalDeclOfTypeOrError (expectedType : Expr) : MetaM LocalDecl := do
-    let some name ← findLocalDeclOfType? expectedType
+    let some decl ← findLocalDeclOfType? expectedType
       | throwError "Failed to find a local hypothesis of type {expectedType}"
-    return name
+    return decl
 
 /-- `findProgramHyp` searches the local context for an hypothesis of type
   `state.program = ?concreteProgram`,
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 8ea70b42..8c68df3a 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -166,7 +166,7 @@ def explodeStep (hStep : Expr) : SymM Unit :=
     let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep
 
-    let stateExpr ← getCurrentState
+    let stateExpr ← SymContext.getCurrentState
     /- Assert that the initial state of the obtained `AxEffects` is equal to
     the state tracked by `c`.
     This will catch and throw an error if the semantics of the current
@@ -246,6 +246,7 @@ def explodeStep (hStep : Expr) : SymM Unit :=
         {err.toMessageData}"
       pure none
 
+    set eff
     modifyThe SymContext (·.next nextPc?)
 
 /-- A tactic wrapper around `explodeStep`.
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 88d90d21..9f64411d 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -61,6 +61,13 @@ structure SymContext where
   /-- `programInfo` is the relevant cached `ProgramInfo` -/
   programInfo : ProgramInfo
 
+  /-- the effects of the current state, such as:
+  - a proof that the PC is equal to `pc`
+  - a proof that the current state is valid (`read_err _ = .None`)
+  - a proof that the current state has the right program
+  - and more, see `AxEffects` for the full list -/
+  effects : AxEffects
+
   /-- `pc` is the *concrete* value of the program counter
 
   Note that for now we only support symbolic evaluation of programs
@@ -116,6 +123,16 @@ def run (ctx : SymContext) (k : SymM α) : TacticM (α × SymContext) :=
 instance : MonadLift SymReaderM SymM where
   monadLift x c := do return (←x c, c)
 
+instance : MonadStateOf AxEffects SymM where
+  get := do return (← getThe SymContext).effects
+  set effects := do modifyThe SymContext ({· with effects})
+  modifyGet f := do
+    let (a, effects) := f (← getThe SymContext).effects
+    modifyThe SymContext ({· with effects})
+    return a
+instance : MonadReaderOf AxEffects SymM where
+  read := getThe AxEffects
+
 end SymM
 
 namespace SymContext
@@ -332,6 +349,12 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
       (decls := axHyps)
       (noIndexAtArgs := false)
 
+  -- Build an initial AxEffects
+  let effects := {
+    AxEffects.initial stateExpr with
+      programProof := h_program.toExpr
+  }
+
   return inferStatePrefixAndNumber {
     state, finalState, runSteps?, program, pc,
     h_run := h_run.userName,
@@ -340,6 +363,7 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
     h_err? := (·.userName) <$> h_err?,
     h_sp? := (·.userName) <$> h_sp?,
     programInfo,
+    effects,
     aggregateSimpCtx, aggregateSimprocs
   }
 where
@@ -390,7 +414,10 @@ def addGoalsForMissingHypotheses (ctx : SymContext) (addHSp : Bool := false) :
             return goal'
 
           newGoals := newGoal :: newGoals
-          ctx := { ctx with h_err? }
+          ctx := { ctx with
+            h_err?
+            effects := ← ctx.effects.withField (.mvar newGoal)
+          }
       | some h_err =>
           let h_err ← userNameToMessageData h_err
           trace[Tactic.sym] "h_err? is {h_err}, no new goal needed"
@@ -411,7 +438,10 @@ def addGoalsForMissingHypotheses (ctx : SymContext) (addHSp : Bool := false) :
               return goal'
 
             newGoals := newGoal :: newGoals
-            ctx := { ctx with h_sp? }
+            ctx := { ctx with
+              h_sp?
+              effects.stackAlignmentProof? := some (Expr.mvar newGoal)
+            }
           else
             trace[Tactic.sym] "h_sp? is none, but addHSp is false, \
               so no new goal is added"

From b185bb074b5ebf0e81ba8b093f51db461e4efea6 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 18 Sep 2024 14:10:50 -0500
Subject: [PATCH 03/18] remove use of evalTactic in stepiTac

---
 Tactics/Reflect/AxEffects.lean |  5 ++++-
 Tactics/Sym.lean               | 37 ++++++++++++++++++++--------------
 Tactics/SymContext.lean        | 22 ++++++++++++++------
 3 files changed, 42 insertions(+), 22 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 8f6cc770..09dbba48 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -483,7 +483,10 @@ def withField (eff : AxEffects) (eq : Expr) : MetaM AxEffects := do
     trace[Tactic.sym] "current field effect: {fieldEff}"
 
     if field ∉ eff.fields then
-      let proof ← mkEqTrans fieldEff.proof eq
+      let proof ← if eff.currentState == eff.initialState then
+          pure eq
+        else
+          mkEqTrans fieldEff.proof eq
       let fields := eff.fields.insert field { value, proof }
       return { eff with fields }
     else
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 8c68df3a..b9bd9190 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -44,25 +44,30 @@ macro "init_next_step" h_run:ident stepi_eq:ident sn:ident : tactic =>
 
 section stepiTac
 
-/-- Apply the relevant pre-generated stepi lemma to a local hypothesis
+/-- Apply the relevant pre-generated stepi lemma to an expression
   `stepi_eq : stepi ?s = ?s'`
-to obtain a new local hypothesis in terms of `w` and `write_mem`
+to add a new local hypothesis in terms of `w` and `write_mem`
   `h_step : ?s' = w _ _ (w _ _ (... ?s))`
 -/
-def stepiTac (stepi_eq h_step : Ident) : SymReaderM Unit := fun ctx =>
-  withMainContext do
+def stepiTac (stepiEq : Expr) (hStep : Name) : SymReaderM Unit := fun ctx =>
+  withMainContext' do
     let pc := (Nat.toDigits 16 ctx.pc.toNat).asString
     --  ^^ The PC in hex
-    let step_lemma := mkIdent <| Name.str ctx.program s!"stepi_eq_0x{pc}"
-
-    evalTacticAndTrace <|← `(tactic| (
-      have $h_step :=
-        _root_.Eq.trans (Eq.symm $stepi_eq)
-          ($step_lemma:ident
-            $ctx.h_program_ident:ident
-            $ctx.h_pc_ident:ident
-            $ctx.h_err_ident:ident)
-    ))
+    let stepLemma := Name.str ctx.program s!"stepi_eq_0x{pc}"
+    -- let stepLemma := Expr.const stepLemma []
+
+    let eff := ctx.effects
+    let hStepExpr ← mkEqTrans
+      (← mkEqSymm stepiEq)
+      (← mkAppM stepLemma #[
+        eff.programProof,
+        (← eff.getField .PC).proof,
+        (← eff.getField .ERR).proof
+      ])
+
+    let goal ← getMainGoal
+    let ⟨_, goal⟩ ← goal.note hStep hStepExpr
+    replaceMainGoal [goal]
 
 end stepiTac
 
@@ -284,7 +289,9 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
     )
 
     -- Apply relevant pre-generated `stepi` lemma
-    stepiTac stepi_eq h_step
+    withMainContext' <| do
+      let stepiEq ← SymContext.findFromUserName stepi_eq.getId
+      stepiTac stepiEq.toExpr h_step.getId
 
     -- WORKAROUND: eventually we'd like to eagerly simp away `if`s in the
     -- pre-generation of instruction semantics. For now, though, we keep a
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 9f64411d..3ed5b266 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -320,12 +320,13 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
     findProgramHyp stateExpr
 
   -- Then, try to find `h_pc`
-  let pc ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
-  let h_pc ← findLocalDeclOfTypeOrError <| h_pc_type stateExpr pc
+  let pcE ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
+  let h_pc ← findLocalDeclOfTypeOrError <| h_pc_type stateExpr pcE
 
   -- Unwrap and reflect `pc`
-  let pc ← instantiateMVars pc
-  let pc ← withErrorContext h_pc.userName h_pc.type <| reflectBitVecLiteral 64 pc
+  let pcE ← instantiateMVars pcE
+  let pc ← withErrorContext h_pc.userName h_pc.type <|
+    reflectBitVecLiteral 64 pcE
 
   -- Attempt to find `h_err` and `h_sp`
   let h_err? ← findLocalDeclOfType? (h_err_type stateExpr)
@@ -350,9 +351,18 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
       (noIndexAtArgs := false)
 
   -- Build an initial AxEffects
-  let effects := {
-    AxEffects.initial stateExpr with
+  let effects := AxEffects.initial stateExpr
+  let mut fields :=
+    effects.fields.insert .PC { value := pcE, proof := h_pc.toExpr}
+  if let some hErr := h_err? then
+    fields := fields.insert .ERR {
+      value := mkConst ``StateError.None,
+      proof := hErr.toExpr
+    }
+  let effects := { effects with
       programProof := h_program.toExpr
+      stackAlignmentProof? := h_sp?.map (·.toExpr)
+      fields
   }
 
   return inferStatePrefixAndNumber {

From e3dd1f6f01aee94fae1d5c265d7c129c84bc5342 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 18 Sep 2024 14:16:50 -0500
Subject: [PATCH 04/18] get rid of dead code

---
 Tactics/Sym.lean        |  7 ++-----
 Tactics/SymContext.lean | 16 ----------------
 2 files changed, 2 insertions(+), 21 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index b9bd9190..857e99d6 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -180,11 +180,8 @@ def explodeStep (hStep : Expr) : SymM Unit :=
       throwError "[explodeStep] expected initial state {stateExpr}, but found:\n  \
         {eff.initialState}\nin\n\n{eff}"
 
-    let hProgram ← SymContext.findFromUserName c.h_program
-    eff ← eff.withProgramEq hProgram.toExpr
-
-    let hErr ← SymContext.findFromUserName c.h_err
-    eff ← eff.withField hErr.toExpr
+    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
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 3ed5b266..7030099d 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -56,8 +56,6 @@ structure SymContext where
   h_run : Name
   /-- `program` is a *constant* which represents the program being evaluated -/
   program : Name
-  /-- `h_program` is a local hypothesis of the form `state.program = program` -/
-  h_program : Name
   /-- `programInfo` is the relevant cached `ProgramInfo` -/
   programInfo : ProgramInfo
 
@@ -146,19 +144,9 @@ variable (c : SymContext)
 def next_state (c : SymContext) : Name :=
   .mkSimple s!"{c.state_prefix}{c.curr_state_number + 1}"
 
-/-- return `h_err?` if given, or a default hardcoded name -/
-def h_err : Name := c.h_err?.getD (.mkSimple s!"h_{c.state}_err")
-
-/-- return `h_sp?` if given, or a default hardcoded name -/
-def h_sp  : Name := c.h_err?.getD (.mkSimple s!"h_{c.state}_sp")
-
 def state_ident       : Ident := mkIdent c.state
 def next_state_ident  : Ident := mkIdent c.next_state
 def h_run_ident       : Ident := mkIdent c.h_run
-def h_program_ident   : Ident := mkIdent c.h_program
-def h_pc_ident        : Ident := mkIdent c.h_pc
-def h_err_ident       : Ident := mkIdent c.h_err
-def h_sp_ident        : Ident := mkIdent c.h_sp
 
 /-- Find the local declaration that corresponds to a given name,
 or throw an error if no local variable of that name exists -/
@@ -368,7 +356,6 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
   return inferStatePrefixAndNumber {
     state, finalState, runSteps?, program, pc,
     h_run := h_run.userName,
-    h_program := h_program.userName,
     h_pc := h_pc.userName
     h_err? := (·.userName) <$> h_err?,
     h_sp? := (·.userName) <$> h_sp?,
@@ -472,10 +459,8 @@ def canonicalizeHypothesisTypes (c : SymContext) : TacticM Unit := withMainConte
   let lctx ← getLCtx
   let mut goal ← getMainGoal
   let state ← c.stateExpr
-  let program := mkConst c.program
 
   let mut hyps := #[
-    (c.h_program, h_program_type state program),
     (c.h_pc, h_pc_type state (toExpr c.pc))
   ]
   if let some runSteps := c.runSteps? then
@@ -503,7 +488,6 @@ def next (c : SymContext) (nextPc? : Option (BitVec 64) := none) :
   let s := c.next_state
   { c with
     state       := s
-    h_program   := .mkSimple s!"h_{s}_program"
     h_pc        := .mkSimple s!"h_{s}_pc"
     h_err?      := c.h_err?.map (fun _ => .mkSimple s!"h_{s}_err")
     h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")

From bc6a373cd31258718d434fca59c5b7d38f8d7f7f Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 18 Sep 2024 14:20:05 -0500
Subject: [PATCH 05/18] fix: update should expect currentState, not
 initialState

---
 Tactics/Reflect/AxEffects.lean | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 09dbba48..c6cb5ab1 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -376,6 +376,28 @@ 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 `currentState`.
+
+Note that as soon as an unsupported expression (e.g., an `if`) is encountered,
+the whole expression is taken to be the initial state,
+even if there might be more `w`/`write_mem`s in sub-expressions. -/
+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
+
+    | w field value e =>
+        let eff ← eff.updateWithExpr e
+        eff.update_w field value
+
+    | _ =>
+        assertIsDefEq e eff.currentState
+        return eff
+
 /-- Given an expression `e : ArmState`,
 which is a sequence of `w`/`write_mem`s to the some state `s`,
 return an `AxEffects` where `s` is the intial state, and `e` is `currentState`.

From d8d780c0d61503d9820c7d8b5b8ea203cad2cc15 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 20 Sep 2024 21:01:04 -0500
Subject: [PATCH 06/18] refactor: remove `state` field from `SymContext`, and
 replace uses with `getCurrentState`

The latter refers to the `currentState` field from `AxEffects`, which stores the same state, but as an Expr
---
 Tactics/Reflect/AxEffects.lean |  14 ++++
 Tactics/Sym.lean               | 148 ++++++++++++++++++---------------
 Tactics/SymContext.lean        | 137 +++++++++++++-----------------
 3 files changed, 154 insertions(+), 145 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index c6cb5ab1..8f3ea555 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -97,6 +97,20 @@ def getProgramProof       : m Expr := do return (← read).programProof
 def getStackAlignmentProof? : m (Option Expr) := do
   return (← read).stackAlignmentProof?
 
+variable [MonadLiftT MetaM m] in
+/-- Retrieve the user-facing name of the current state, assuming that
+the current state is a free variable in the ambient local context -/
+def getCurrentStateName : m Name := do
+  let state ← getCurrentState
+  @id (MetaM _) <| do
+    let state ← instantiateMVars state
+    let Expr.fvar id := state.consumeMData
+      | throwError "error: expected a free variable, found:\n  {state} WHHOPS"
+    let lctx ← getLCtx
+    let some decl := lctx.find? id
+      | throwError "error: unknown fvar: {state}"
+    return decl.userName
+
 end Monad
 
 /-! ## Initial Reflected State -/
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 857e99d6..f0ceeaaa 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -87,7 +87,8 @@ for some metavariable `?runSteps`, then create the proof obligation
 `?runSteps = _ + 1`, and attempt to close it using `whileTac`.
 Finally, we use this proof to change the type of `h_run` accordingly.
 -/
-def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
+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) <|
   match c.runSteps? with
@@ -100,7 +101,7 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
         -- NOTE: this error shouldn't occur, as we should have checked in
         -- `sym_n` that, if the number of runSteps is statically known,
         -- that we never simulate more than that many steps
-    | none => withMainContext do
+    | none => withMainContext' do
         let mut goal :: originalGoals ← getGoals
           | throwNoGoalsToBeSolved
         let hRunDecl ← c.hRunDecl
@@ -110,7 +111,7 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
         guard <|← isDefEq hRunDecl.type (
           mkApp3 (.const ``Eq [1]) (mkConst ``ArmState)
             c.finalState
-            (mkApp2 (mkConst ``_root_.run) runSteps (←c.stateExpr)))
+            (mkApp2 (mkConst ``_root_.run) runSteps (← getCurrentState)))
         -- NOTE: ^^ Since we check for def-eq on SymContext construction,
         --          this check should never fail
 
@@ -138,8 +139,9 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
           runStepsPredId.assign default
 
         -- Change the type of `h_run`
+        let state ← getCurrentState
         let typeNew ← do
-          let rhs := mkApp2 (mkConst ``_root_.run) subGoalTyRhs (←c.stateExpr)
+          let rhs := mkApp2 (mkConst ``_root_.run) subGoalTyRhs state
           mkEq c.finalState rhs
         let eqProof ← do
           let f := -- `fun s => <finalState> = s`
@@ -147,9 +149,8 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := fun c =>
                         c.finalState (.bvar 0)
             .lam `s (mkConst ``ArmState) eq .default
           let g := mkConst ``_root_.run
-          let s ← c.stateExpr
           let h ← instantiateMVars (.mvar subGoal)
-          mkCongrArg f (←mkCongrFun (←mkCongrArg g h) s)
+          mkCongrArg f (←mkCongrFun (←mkCongrArg g h) state)
         let res ← goal.replaceLocalDecl hRunDecl.fvarId typeNew eqProof
 
         -- Restore goal state
@@ -171,7 +172,7 @@ def explodeStep (hStep : Expr) : SymM Unit :=
     let c ← getThe SymContext
     let mut eff ← AxEffects.fromEq hStep
 
-    let stateExpr ← SymContext.getCurrentState
+    let stateExpr ← getCurrentState
     /- Assert that the initial state of the obtained `AxEffects` is equal to
     the state tracked by `c`.
     This will catch and throw an error if the semantics of the current
@@ -225,31 +226,32 @@ def explodeStep (hStep : Expr) : SymM Unit :=
                 #[eff.currentState, spEff.value, spEff.proof, hAligned]
             pure { eff with stackAlignmentProof? }
 
-    -- Add new (non-)effect hyps to the context
-    let simpThms ← withMainContext' <| do
+    -- Add new (non-)effect hyps to the context, and to the aggregation simpset
+    withMainContext' <| do
       if ←(getBoolOption `Tactic.sym.debug) then
         eff.validate
 
       let eff ← eff.addHypothesesToLContext s!"h_{← getNextStateName}_"
-      withMainContext <| eff.toSimpTheorems
-
-    -- Add the new (non-)effect hyps to the aggregation simp context
-    modifyThe SymContext (·.addSimpTheorems simpThms)
-
-    -- Attempt to reflect the new PC
-    let nextPc ← eff.getField .PC
-    let nextPc? ← try
-      let nextPc ← reflectBitVecLiteral 64 nextPc.value
-      -- NOTE: `reflectBitVecLiteral` is fast when the value is a literal,
-      -- but might involve an expensive reduction when it is not
-      pure <| some nextPc
-    catch err =>
-      trace[Tactic.sym] "failed to reflect {nextPc.value}\n\n\
-        {err.toMessageData}"
-      pure none
-
-    set eff
-    modifyThe SymContext (·.next nextPc?)
+      withMainContext' <| do
+        let simpThms ← eff.toSimpTheorems
+        modifyThe SymContext (·.addSimpTheorems simpThms)
+      set eff
+
+    -- Prepare sym context for the next step
+    withMainContext' <| do
+      -- Attempt to reflect the new PC
+      let nextPc ← eff.getField .PC
+      let nextPc? ← try
+        let nextPc ← reflectBitVecLiteral 64 nextPc.value
+        -- NOTE: `reflectBitVecLiteral` is fast when the value is a literal,
+        -- but might involve an expensive reduction when it is not
+        pure <| some nextPc
+      catch err =>
+        trace[Tactic.sym] "failed to reflect {nextPc.value}\n\n\
+          {err.toMessageData}"
+        pure none
+
+      modifyThe SymContext (·.next nextPc?)
 
 /-- A tactic wrapper around `explodeStep`.
 Note the use of `SymContext.fromLocalContext`,
@@ -329,6 +331,42 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
 
       traceHeartbeats
 
+/-- `ensureLessThanRunSteps n` will
+- 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
+
+/-- 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
+  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}
+
+Alternatively, are you sure that {pc} is a valid value for the program counter?"
+-- 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
 
@@ -375,41 +413,20 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         omega;
         ))
 
-  let c ← withMainContext <| do
-    let c ← SymContext.fromLocalContext s
-    let c ← c.addGoalsForMissingHypotheses
-    c.canonicalizeHypothesisTypes
-    pure c
-
-  let _ ← withMainContext <| SymM.run c <| do
-    -- Check that we are not asked to simulate more steps than available
-    let n ← do
-      let n := n.getNat
-      match c.runSteps? with
-        | none => pure n
-        | some runSteps =>
-            if n ≤ runSteps then
-              pure n
-            else
-              let h_run ← userNameToMessageData c.h_run
-              logWarning m!"Symbolic simulation using {h_run} is limited to at most {runSteps} steps"
-              pure runSteps
-
-    -- Check that step theorems have been pre-generated
-    try
-      let pc := c.pc.toHexWithoutLeadingZeroes
-      let step_thm := Name.str c.program ("stepi_eq_0x" ++ pc)
-      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?
+  let c ← withMainContext <| SymContext.fromLocalContext s
+  SymM.run' c <| do
+    -- Context preparation
+    addGoalsForMissingHypotheses
+    canonicalizeHypothesisTypes
+
+    -- Check pre-conditions
+    assertStepTheoremsGenerated
+    let n ← ensureAtMostRunSteps n.getNat
 
-    -- The main loop
-    for _ in List.range n do
-      sym1 whileTac
+    withMainContext' <| do
+      -- The main loop
+      for _ in List.range n do
+        sym1 whileTac
 
     traceHeartbeats "symbolic simulation total"
     let c ← getThe SymContext
@@ -418,9 +435,8 @@ Did you remember to generate step theorems with:
       let msg := do
         let hRun ← userNameToMessageData c.h_run
         pure m!"runSteps := 0, substituting along {hRun}"
-      withTraceNode `Tactic.sym (fun _ => msg) <| withMainContext do
-        let s ← SymContext.findFromUserName c.state
-        let sfEq ← mkEq s.toExpr c.finalState
+      withTraceNode `Tactic.sym (fun _ => msg) <| withMainContext' do
+        let sfEq ← mkEq (← getCurrentState) c.finalState
 
         let goal ← getMainGoal
         trace[Tactic.sym] "original goal:\n{goal}"
@@ -440,7 +456,7 @@ Did you remember to generate step theorems with:
     -- 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 `Tactic.sym (fun _ => pure msg) <| withMainContext' do
       traceHeartbeats "pre"
       let goal? ← LNSymSimp (← getMainGoal) c.aggregateSimpCtx c.aggregateSimprocs
       replaceMainGoal goal?.toList
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 7030099d..c6daf9d4 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -35,8 +35,6 @@ open BitVec
 /-- A `SymContext` collects the names of various variables/hypotheses in
 the local context required for symbolic evaluation -/
 structure SymContext where
-  /-- `state` is a local variable of type `ArmState` -/
-  state : Name
   /-- `finalState` is an expression of type `ArmState` -/
   finalState : Expr
   /-- `runSteps?` stores the number of steps that we can *maximally* simulate,
@@ -54,8 +52,6 @@ structure SymContext where
 
   See also `SymContext.runSteps?` -/
   h_run : Name
-  /-- `program` is a *constant* which represents the program being evaluated -/
-  program : Name
   /-- `programInfo` is the relevant cached `ProgramInfo` -/
   programInfo : ProgramInfo
 
@@ -118,18 +114,22 @@ namespace SymM
 def run (ctx : SymContext) (k : SymM α) : TacticM (α × SymContext) :=
   StateT.run k ctx
 
+def run' (ctx : SymContext) (k : SymM α) : TacticM α :=
+  StateT.run' k ctx
+
 instance : MonadLift SymReaderM SymM where
   monadLift x c := do return (←x c, c)
 
+instance : MonadReaderOf AxEffects SymReaderM where
+  read := do return (← readThe SymContext).effects
+
 instance : MonadStateOf AxEffects SymM where
-  get := do return (← getThe SymContext).effects
+  get := readThe AxEffects
   set effects := do modifyThe SymContext ({· with effects})
   modifyGet f := do
     let (a, effects) := f (← getThe SymContext).effects
     modifyThe SymContext ({· with effects})
     return a
-instance : MonadReaderOf AxEffects SymM where
-  read := getThe AxEffects
 
 end SymM
 
@@ -140,14 +140,13 @@ section
 open Lean (Ident mkIdent)
 variable (c : SymContext)
 
+/-- `program` is a *constant* which represents the program being evaluated -/
+def program : Name := c.programInfo.name
+
 /-- `next_state` generates the name for the next intermediate state -/
 def next_state (c : SymContext) : Name :=
   .mkSimple s!"{c.state_prefix}{c.curr_state_number + 1}"
 
-def state_ident       : Ident := mkIdent c.state
-def next_state_ident  : Ident := mkIdent c.next_state
-def h_run_ident       : Ident := mkIdent c.h_run
-
 /-- 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
@@ -155,11 +154,6 @@ def findFromUserName (name : Name) : MetaM LocalDecl := do
     | throwError "Unknown local variable `{name}`"
   return decl
 
-/-- Return an expression for `c.state`,
-or throw an error if no local variable of that name exists -/
-def stateExpr : MetaM Expr :=
-  (·.toExpr) <$> findFromUserName c.state
-
 /-- Find the local declaration that corresponds to `c.h_run`,
 or throw an error if no local variable of that name exists -/
 def hRunDecl : MetaM LocalDecl := do
@@ -170,15 +164,6 @@ variable {m} [Monad m] [MonadReaderOf SymContext m]
 
 def getCurrentStateNumber : m Nat := do return (← read).curr_state_number
 
-/-- Retrieve the name of the current state -/
-def getCurrentStateName : m Name := do
-  return (← read).state
-
-/-- Retrieve an expression for the current state,
-or throw an error if no local variable of that name exists -/
-def getCurrentState [MonadLiftT MetaM m] : m Expr := do
-  (← read).stateExpr
-
 /-- Return the name of the hypothesis
   `h_run : <finalState> = run <runSteps> <initialState>` -/
 def getHRunName : m Name := do return (← read).h_run
@@ -198,21 +183,24 @@ end
 /-- Convert a `SymContext` to `MessageData` for tracing.
 This is not a `ToMessageData` instance because we need access to `MetaM` -/
 def toMessageData (c : SymContext) : MetaM MessageData := do
-  let state ← c.stateExpr
   let h_run ← userNameToMessageData c.h_run
-  let h_err? ← c.h_err?.mapM userNameToMessageData
   let h_sp?  ← c.h_sp?.mapM userNameToMessageData
 
-  return m!"\{ state := {state},
-  finalState := {c.finalState},
+  return m!"\{ finalState := {c.finalState},
   runSteps? := {c.runSteps?},
   h_run := {h_run},
   program := {c.program},
   pc := {c.pc},
-  h_err? := {h_err?},
   h_sp? := {h_sp?},
   state_prefix := {c.state_prefix},
-  curr_state_number := {c.curr_state_number} }"
+  curr_state_number := {c.curr_state_number},
+  effects := {c.effects} }"
+
+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
 
 def traceSymContext : SymM Unit :=
   withTraceNode `Tactic.sym (fun _ => pure m!"SymContext: ") <| do
@@ -221,22 +209,29 @@ def traceSymContext : SymM Unit :=
 
 /-! ## Creating initial contexts -/
 
+/-- Modify a `SymContext` with a monadic action `k : SymM Unit` -/
+def modify (ctxt : SymContext) (k : SymM Unit) : TacticM SymContext := do
+  let ((), ctxt) ← SymM.run ctxt k
+  return ctxt
+
 /-- Infer `state_prefix` and `curr_state_number` from the `state` name
 as follows: if `state` is `s{i}` for some number `i` and a single character `s`,
 then `s` is the prefix and `i` the number,
 otherwise ignore `state`, and start counting from `s1` -/
-def inferStatePrefixAndNumber (ctxt : SymContext) : SymContext :=
-  let state := ctxt.state.toString
+def inferStatePrefixAndNumber : SymM Unit := do
+  let state ← AxEffects.getCurrentStateName
+  let state := state.toString
   let tail := state.toSubstring.drop 1
 
-  if let some curr_state_number := tail.toNat? then
-    { ctxt with
-      state_prefix := (state.get? 0).getD 's' |>.toString,
-      curr_state_number }
-  else
-    { ctxt with
-      state_prefix := "s",
-      curr_state_number := 1 }
+  modifyThe SymContext fun ctxt =>
+    if let some curr_state_number := tail.toNat? then
+      { ctxt with
+        state_prefix := (state.get? 0).getD 's' |>.toString,
+        curr_state_number }
+    else
+      { ctxt with
+        state_prefix := "s",
+        curr_state_number := 1 }
 
 /-- Annotate any errors thrown by `k` with a local variable (and its type) -/
 private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α) :
@@ -249,8 +244,13 @@ private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α)
     throwErrorAt e.getRef "{e.toMessageData}\n\nIn {h}{type}"
 
 /-- Build a `SymContext` by searching the local context for hypotheses of the
-required types (up-to defeq) -/
-def fromLocalContext (state? : Option Name) : MetaM SymContext := do
+required types (up-to defeq) . The local context is modified to unfold the types
+to be syntactically equal to the expected type.
+
+If an hypothesis `h_err : r <state> .ERR = None` is not found,
+we create a new subgoal of this type
+-/
+def fromLocalContext (state? : Option Name) : TacticM SymContext := do
   let msg := m!"Building a `SymContext` from the local context"
   withTraceNode `Tactic.sym (fun _ => pure msg) do
   trace[Tactic.Sym] "state? := {state?}"
@@ -283,25 +283,6 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
   -- At this point, `stateExpr` should have been assigned (if it was an mvar),
   -- so we can unwrap it to get the underlying name
   let stateExpr ← instantiateMVars stateExpr
-  let state ← state?.getDM <| do
-    let .fvar state := stateExpr
-      | let h_run_type ← instantiateMVars h_run.type
-        let h_run := h_run.toExpr
-        throwError
-  "Expected a free variable, found:
-    {stateExpr}
-  We inferred this as the initial state because we found:
-    {h_run} : {h_run_type}
-  in the local context.
-
-  If this is wrong, please explicitly provide the right initial state,
-  as in `sym {runSteps} at ?s0`
-  "
-    let some state := lctx.find? state
-      /- I don't expect this error to be possible in a well-formed state,
-      but you never know -/
-      | throwError "Failed to find local variable for state {stateExpr}"
-    pure state.userName
 
   -- Try to find `h_program`, and infer `program` from it
   let ⟨h_program, program⟩ ← withErrorContext h_run.userName h_run.type <|
@@ -352,9 +333,8 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
       stackAlignmentProof? := h_sp?.map (·.toExpr)
       fields
   }
-
-  return inferStatePrefixAndNumber {
-    state, finalState, runSteps?, program, pc,
+  let c : SymContext := {
+    finalState, runSteps?, pc,
     h_run := h_run.userName,
     h_pc := h_pc.userName
     h_err? := (·.userName) <$> h_err?,
@@ -363,6 +343,8 @@ def fromLocalContext (state? : Option Name) : MetaM SymContext := do
     effects,
     aggregateSimpCtx, aggregateSimprocs
   }
+  c.modify <|
+    inferStatePrefixAndNumber
 where
   findLocalDeclOfType? (expectedType : Expr) : MetaM (Option LocalDecl) := do
     let msg := m!"Searching for hypothesis of type: {expectedType}"
@@ -384,23 +366,21 @@ where
 /-- If `h_sp` or `h_err` are missing from the `SymContext`,
 add new goals of the expected types,
 and use these to add `h_sp` and `h_err` to the main goal context -/
-def addGoalsForMissingHypotheses (ctx : SymContext) (addHSp : Bool := false) :
-    TacticM SymContext :=
+def addGoalsForMissingHypotheses (addHSp : Bool := false) : SymM Unit :=
   let msg := "Adding goals for missing hypotheses"
-  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext do
-    let mut ctx := ctx
+  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
+    let mut ctx ← getThe SymContext
     let mut goal ← getMainGoal
     let mut newGoals := []
     let lCtx ← getLCtx
-    let some stateExpr :=
-      (Expr.fvar ·.fvarId) <$> lCtx.findFromUserName? ctx.state
-      | throwError "Could not find '{ctx.state}' in the local context"
+    let stateExpr ← AxEffects.getCurrentState
+    let stateName ← AxEffects.getCurrentStateName
 
     match ctx.h_err? with
       | none =>
           trace[Tactic.sym] "h_err? is none, adding a new goal"
 
-          let h_err? := Name.mkSimple s!"h_{ctx.state}_run"
+          let h_err? := Name.mkSimple s!"h_{stateName}_err"
           let newGoal ← mkFreshMVarId
 
           goal := ← do
@@ -424,7 +404,7 @@ def addGoalsForMissingHypotheses (ctx : SymContext) (addHSp : Bool := false) :
           if addHSp then
             trace[Tactic.sym] "h_sp? is none, adding a new goal"
 
-            let h_sp? := Name.mkSimple s!"h_{ctx.state}_sp"
+            let h_sp? := Name.mkSimple s!"h_{stateName}_sp"
             let newGoal ← mkFreshMVarId
 
             goal := ← do
@@ -447,7 +427,7 @@ def addGoalsForMissingHypotheses (ctx : SymContext) (addHSp : Bool := false) :
           trace[Tactic.sym] "h_sp? is {h_sp}, no new goal needed"
 
     replaceMainGoal (goal :: newGoals)
-    return ctx
+    set ctx
 
 /-- change the type (in the local context of the main goal)
 of the hypotheses tracked by the given `SymContext` to be *exactly* of the shape
@@ -455,10 +435,10 @@ described in the relevant docstrings.
 
 That is, (un)fold types which were definitionally, but not syntactically,
 equal to the expected shape. -/
-def canonicalizeHypothesisTypes (c : SymContext) : TacticM Unit := withMainContext do
+def canonicalizeHypothesisTypes : SymReaderM Unit := fun c => withMainContext do
   let lctx ← getLCtx
   let mut goal ← getMainGoal
-  let state ← c.stateExpr
+  let state := c.effects.currentState
 
   let mut hyps := #[
     (c.h_pc, h_pc_type state (toExpr c.pc))
@@ -472,7 +452,7 @@ def canonicalizeHypothesisTypes (c : SymContext) : TacticM Unit := withMainConte
 
   for ⟨name, type⟩ in hyps do
     let some decl := lctx.findFromUserName? name
-      | throwError "Unknown local hypothesis `{c.state}`"
+      | throwError "Unknown local hypothesis `{name}`"
     goal ← goal.replaceLocalDeclDefEq decl.fvarId type
   replaceMainGoal [goal]
 
@@ -487,7 +467,6 @@ def next (c : SymContext) (nextPc? : Option (BitVec 64) := none) :
   let curr_state_number := c.curr_state_number + 1
   let s := c.next_state
   { c with
-    state       := s
     h_pc        := .mkSimple s!"h_{s}_pc"
     h_err?      := c.h_err?.map (fun _ => .mkSimple s!"h_{s}_err")
     h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")

From 34f32260d76f83997c205448bd72dcedbb8be680 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 20 Sep 2024 21:06:31 -0500
Subject: [PATCH 07/18] refactor: remove h_pc field from SymContext

It was unused
---
 Tactics/SymContext.lean | 8 +-------
 1 file changed, 1 insertion(+), 7 deletions(-)

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index c6daf9d4..71043878 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -75,8 +75,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_pc` is a local hypothesis of the form `r StateField.PC state = pc` -/
-  h_pc  : Name
   /-- `h_err?`, if present, is a local hypothesis of the form
   `r StateField.ERR state = .None` -/
   h_err? : Option Name
@@ -336,7 +334,6 @@ def fromLocalContext (state? : Option Name) : TacticM SymContext := do
   let c : SymContext := {
     finalState, runSteps?, pc,
     h_run := h_run.userName,
-    h_pc := h_pc.userName
     h_err? := (·.userName) <$> h_err?,
     h_sp? := (·.userName) <$> h_sp?,
     programInfo,
@@ -440,9 +437,7 @@ def canonicalizeHypothesisTypes : SymReaderM Unit := fun c => withMainContext do
   let mut goal ← getMainGoal
   let state := c.effects.currentState
 
-  let mut hyps := #[
-    (c.h_pc, h_pc_type state (toExpr c.pc))
-  ]
+  let mut hyps := #[]
   if let some runSteps := c.runSteps? then
     hyps := hyps.push (c.h_run, h_run_type c.finalState (toExpr runSteps) state)
   if let some h_err := c.h_err? then
@@ -467,7 +462,6 @@ def next (c : SymContext) (nextPc? : Option (BitVec 64) := none) :
   let curr_state_number := c.curr_state_number + 1
   let s := c.next_state
   { c with
-    h_pc        := .mkSimple s!"h_{s}_pc"
     h_err?      := c.h_err?.map (fun _ => .mkSimple s!"h_{s}_err")
     h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")
     runSteps?   := (· - 1) <$> c.runSteps?

From e5becf10f845dc1a3bb4b97ab0d162e05f433704 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 20 Sep 2024 21:26:29 -0500
Subject: [PATCH 08/18] refactor: rename `curr_state_number` to
 `currentStateNumber`, make `SymContext.next` monadic, and fold-in the
 reflection of the pc (removing the nextPc argument)

---
 Tactics/Reflect/AxEffects.lean |  7 ++++-
 Tactics/Sym.lean               | 29 ++++--------------
 Tactics/SymContext.lean        | 55 +++++++++++++++++++---------------
 3 files changed, 43 insertions(+), 48 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 8f3ea555..9ba29735 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -216,7 +216,7 @@ partial def mkAppNonEffect (eff : AxEffects) (field : Expr) : MetaM Expr := do
       return nonEffectProof
 
 /-- Get the value for a field, if one is stored in `eff.fields`,
-or assemble an instantiation of the non-effects proof -/
+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
@@ -231,6 +231,11 @@ def getField (eff : AxEffects) (fld : StateField) : MetaM FieldEffect :=
       let proof  ← eff.mkAppNonEffect (toExpr fld)
       pure { value, proof }
 
+variable {m} [Monad m] [MonadStateOf AxEffects m] [MonadLiftT MetaM m] in
+@[inherit_doc getField]
+def getFieldM (field : StateField) : m FieldEffect := do
+  (← get).getField field
+
 /-! ## Update a Reflected State -/
 
 /-- Execute `write_mem <n> <addr> <val>` against the state stored in `eff`
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index f0ceeaaa..ae7c2f5f 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -160,13 +160,11 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymReaderM Unit := do
           originalGoals := originalGoals.concat subGoal
         setGoals (res.mvarId :: originalGoals)
 
-/-- Given an equality `h_step : s{i+1} = w ... (... (w ... s{i})...)`,
-add hypotheses that axiomatically describe the effects in terms of
-reads from `s{i+1}`.
-
-Return the context for the next step (see `SymContext.next`), where
-we attempt to determine the new PC by reflecting the obtained effects,
-falling back to incrementing the PC if reflection failed. -/
+/-- Break an equality `h_step : s{i+1} = w ... (... (w ... s{i})...)` into an
+`AxEffects` that characterizes the effects in terms of reads from `s{i+1}`,
+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
     let c ← getThe SymContext
@@ -237,22 +235,6 @@ def explodeStep (hStep : Expr) : SymM Unit :=
         modifyThe SymContext (·.addSimpTheorems simpThms)
       set eff
 
-    -- Prepare sym context for the next step
-    withMainContext' <| do
-      -- Attempt to reflect the new PC
-      let nextPc ← eff.getField .PC
-      let nextPc? ← try
-        let nextPc ← reflectBitVecLiteral 64 nextPc.value
-        -- NOTE: `reflectBitVecLiteral` is fast when the value is a literal,
-        -- but might involve an expensive reduction when it is not
-        pure <| some nextPc
-      catch err =>
-        trace[Tactic.sym] "failed to reflect {nextPc.value}\n\n\
-          {err.toMessageData}"
-        pure none
-
-      modifyThe SymContext (·.next nextPc?)
-
 /-- A tactic wrapper around `explodeStep`.
 Note the use of `SymContext.fromLocalContext`,
 so the local context is assumed to be of the same shape as for `sym_n` -/
@@ -324,6 +306,7 @@ def sym1 (whileTac : TSyntax `tactic) : SymM Unit := do
       -- ^^ we can't reuse `hStep` from before, since its fvarId might've been
       --    changed by `simp`
       explodeStep hStep.toExpr
+      prepareForNextStep
 
       let goal ← getMainGoal
       let goal ← goal.clear hStep.fvarId
diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 71043878..57ad71ef 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -97,7 +97,7 @@ structure SymContext where
   /-- `curr_state_number` is incremented each simulation step,
   and used together with `curr_state_number`
   to determine the name of the next state variable that is added by `sym` -/
-  curr_state_number : Nat := 0
+  currentStateNumber : Nat := 0
 
 /-! ## Monad -/
 
@@ -141,9 +141,9 @@ variable (c : SymContext)
 /-- `program` is a *constant* which represents the program being evaluated -/
 def program : Name := c.programInfo.name
 
-/-- `next_state` generates the name for the next intermediate state -/
+/-- `next_state` generates the name for the next intermediate state. -/
 def next_state (c : SymContext) : Name :=
-  .mkSimple s!"{c.state_prefix}{c.curr_state_number + 1}"
+  .mkSimple s!"{c.state_prefix}{c.currentStateNumber + 1}"
 
 /-- Find the local declaration that corresponds to a given name,
 or throw an error if no local variable of that name exists -/
@@ -160,7 +160,7 @@ def hRunDecl : MetaM LocalDecl := do
 section Monad
 variable {m} [Monad m] [MonadReaderOf SymContext m]
 
-def getCurrentStateNumber : m Nat := do return (← read).curr_state_number
+def getCurrentStateNumber : m Nat := do return (← read).currentStateNumber
 
 /-- Return the name of the hypothesis
   `h_run : <finalState> = run <runSteps> <initialState>` -/
@@ -168,8 +168,9 @@ def getHRunName : m Name := do return (← read).h_run
 
 /-- Retrieve the name for the next state
 
-NOTE: does not increment the state;
-consecutive calls to `getNextStateName` will give the same name -/
+NOTE: `getNextStateName` does not increment the state, so consecutive calls
+will give the same name. Calling `prepareForNextStep` will increment the state.
+-/
 def getNextStateName : m Name := do return (← read).next_state
 
 end Monad
@@ -191,7 +192,7 @@ def toMessageData (c : SymContext) : MetaM MessageData := do
   pc := {c.pc},
   h_sp? := {h_sp?},
   state_prefix := {c.state_prefix},
-  curr_state_number := {c.curr_state_number},
+  curr_state_number := {c.currentStateNumber},
   effects := {c.effects} }"
 
 variable {α : Type} {m : Type → Type} [Monad m] [MonadTrace m] [MonadLiftT IO m]
@@ -222,14 +223,14 @@ def inferStatePrefixAndNumber : SymM Unit := do
   let tail := state.toSubstring.drop 1
 
   modifyThe SymContext fun ctxt =>
-    if let some curr_state_number := tail.toNat? then
+    if let some currentStateNumber := tail.toNat? then
       { ctxt with
         state_prefix := (state.get? 0).getD 's' |>.toString,
-        curr_state_number }
+        currentStateNumber }
     else
       { ctxt with
         state_prefix := "s",
-        curr_state_number := 1 }
+        currentStateNumber := 1 }
 
 /-- Annotate any errors thrown by `k` with a local variable (and its type) -/
 private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α) :
@@ -369,7 +370,6 @@ def addGoalsForMissingHypotheses (addHSp : Bool := false) : SymM Unit :=
     let mut ctx ← getThe SymContext
     let mut goal ← getMainGoal
     let mut newGoals := []
-    let lCtx ← getLCtx
     let stateExpr ← AxEffects.getCurrentState
     let stateName ← AxEffects.getCurrentStateName
 
@@ -453,21 +453,28 @@ def canonicalizeHypothesisTypes : SymReaderM Unit := fun c => withMainContext do
 
 /-! ## Incrementing the context to the next state -/
 
-/-- `c.next` generates names for the next intermediate state and its hypotheses
-
-`nextPc?`, if given, will be the pc of the next context.
-If `nextPC?` is `none`, then the previous pc is incremented by 4 -/
-def next (c : SymContext) (nextPc? : Option (BitVec 64) := none) :
-    SymContext :=
-  let curr_state_number := c.curr_state_number + 1
-  let s := c.next_state
-  { c with
-    h_err?      := c.h_err?.map (fun _ => .mkSimple s!"h_{s}_err")
+/-- `prepareForNextStep` prepares the state for the next step of symbolic
+evaluation:
+  * `pc` is reflected from the stored `effects`
+  * `runSteps?`, if specified, is decremented,
+  * the `currentStateNumber` is incremented
+-/
+def prepareForNextStep : SymM Unit := do
+  let s ← getNextStateName
+  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
     h_sp?       := c.h_sp?.map (fun _ => .mkSimple s!"h_{s}_sp_aligned")
     runSteps?   := (· - 1) <$> c.runSteps?
-    pc          := nextPc?.getD (c.pc + 4#64)
-    curr_state_number
-  }
+    currentStateNumber := c.currentStateNumber + 1
+  })
 
 /-- Add a set of new simp-theorems to the simp-theorems used
 for effect aggregation -/

From c20824811f2d1ae17239fa4ae5f571530480982c Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Fri, 20 Sep 2024 21:36:51 -0500
Subject: [PATCH 09/18] remove `next_state`

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

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 57ad71ef..b959afab 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -141,10 +141,6 @@ variable (c : SymContext)
 /-- `program` is a *constant* which represents the program being evaluated -/
 def program : Name := c.programInfo.name
 
-/-- `next_state` generates the name for the next intermediate state. -/
-def next_state (c : SymContext) : Name :=
-  .mkSimple s!"{c.state_prefix}{c.currentStateNumber + 1}"
-
 /-- 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
@@ -171,7 +167,9 @@ def getHRunName : m Name := do return (← read).h_run
 NOTE: `getNextStateName` does not increment the state, so consecutive calls
 will give the same name. Calling `prepareForNextStep` will increment the state.
 -/
-def getNextStateName : m Name := do return (← read).next_state
+def getNextStateName : m Name := do
+  let c ← read
+  return Name.mkSimple s!"{c.state_prefix}{c.currentStateNumber + 1}"
 
 end Monad
 
@@ -357,8 +355,6 @@ where
       trace[Tactic.sym] "Found: {decl.toExpr}"
       return decl
 
-
-
 /-! ## Massaging the local context -/
 
 /-- If `h_sp` or `h_err` are missing from the `SymContext`,

From ded5d9aaa2f8c2de86729873bf8f18dcf3465aac Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:01:06 -0500
Subject: [PATCH 10/18] refactor: downgrade getFieldM to a ReaderM

---
 Tactics/Reflect/AxEffects.lean | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 9ba29735..cce9552d 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -231,10 +231,10 @@ def getField (eff : AxEffects) (fld : StateField) : MetaM FieldEffect :=
       let proof  ← eff.mkAppNonEffect (toExpr fld)
       pure { value, proof }
 
-variable {m} [Monad m] [MonadStateOf AxEffects m] [MonadLiftT MetaM m] in
+variable {m} [Monad m] [MonadReaderOf AxEffects m] [MonadLiftT MetaM m] in
 @[inherit_doc getField]
 def getFieldM (field : StateField) : m FieldEffect := do
-  (← get).getField field
+  (← read).getField field
 
 /-! ## Update a Reflected State -/
 

From 21edc25424abd602c790a9d4bf6ef673d3aa8d53 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:35:08 -0500
Subject: [PATCH 11/18] improve assertStepTheoremsGenerated error message

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

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index ae7c2f5f..1b74546a 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -338,15 +338,22 @@ 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}
-
-Alternatively, are you sure that {pc} is a valid value for the program counter?"
+  #genStepEqTheorems {c.program}"
 -- TODO: can we make this error ^^ into a `Try this:` suggestion that
 --       automatically adds the right command just before the theorem?
 

From 71b7ae9007b07a5a83d6214a42f58f87fbf3f027 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:38:50 -0500
Subject: [PATCH 12/18] Update Tactics/SymContext.lean

Co-authored-by: Shilpi Goel <shigoel@gmail.com>
---
 Tactics/SymContext.lean | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index b959afab..317883f8 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -104,7 +104,7 @@ structure SymContext where
 /-- `SymM` is a wrapper around `TacticM` with a mutable `SymContext` state -/
 abbrev SymM := StateT SymContext TacticM
 
-/-- `SymM` is a wrapper around `TacticM` with a read-only `SymContext` state -/
+/-- `SymReaderM` is a wrapper around `TacticM` with a read-only `SymContext` state -/
 abbrev SymReaderM := ReaderT SymContext TacticM
 
 namespace SymM

From c243bd39b6b01db6caa7cca6370e9bfe237f7525 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:39:32 -0500
Subject: [PATCH 13/18] limit scope of withMainContext' variable

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

diff --git a/Tactics/Common.lean b/Tactics/Common.lean
index 129ec6c6..ac1432ee 100644
--- a/Tactics/Common.lean
+++ b/Tactics/Common.lean
@@ -272,7 +272,7 @@ def traceHeartbeats (cls : Name) (header : Option String := none) :
 
 /-! ## `withMainContext'` -/
 
-variable {m} [Monad m] [MonadLiftT TacticM m] [MonadControlT MetaM m]
+variable {m} [Monad m] [MonadLiftT TacticM m] [MonadControlT MetaM m] in
 /-- Execute `x` using the main goal local context and instances.
 
 Unlike the standard `withMainContext`, `x` may live in a generic monad `m`. -/

From 89df82a476ec001716c8222964d9adcdb713453c Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:41:25 -0500
Subject: [PATCH 14/18] make updateWithExpr docstring more accurate

---
 Tactics/Reflect/AxEffects.lean | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index cce9552d..69cb5238 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -397,11 +397,7 @@ 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 `currentState`.
-
-Note that as soon as an unsupported expression (e.g., an `if`) is encountered,
-the whole expression is taken to be the initial state,
-even if there might be more `w`/`write_mem`s in sub-expressions. -/
+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

From 523dc74bf31c5d88f8c16602afa9e37224088820 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:51:33 -0500
Subject: [PATCH 15/18] log warning when `inferStatePrefixAndNumber` falls back
 to the default numbering scheme

---
 Tactics/SymContext.lean | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 317883f8..1be165a0 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -220,15 +220,23 @@ def inferStatePrefixAndNumber : SymM Unit := do
   let state := state.toString
   let tail := state.toSubstring.drop 1
 
-  modifyThe SymContext fun ctxt =>
+  let ctxt ← getThe SymContext
+  let ctxt : SymContext ← do
     if let some currentStateNumber := tail.toNat? then
-      { ctxt with
+      pure { ctxt with
         state_prefix := (state.get? 0).getD 's' |>.toString,
         currentStateNumber }
     else
-      { ctxt with
+      logWarning "\
+        Expected state to be a single letter followed by a number, but found:
+          {state}
+
+        Falling back to the default numbering schema,
+        with `s1` as the first new intermediate state"
+      pure { ctxt with
         state_prefix := "s",
         currentStateNumber := 1 }
+  set ctxt
 
 /-- Annotate any errors thrown by `k` with a local variable (and its type) -/
 private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α) :

From 3a732b11b3f950dacef170ceb5707a507a92aaa7 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:54:34 -0500
Subject: [PATCH 16/18] minor cleanup

---
 Tactics/SymContext.lean | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 1be165a0..fbd64b23 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -220,23 +220,20 @@ def inferStatePrefixAndNumber : SymM Unit := do
   let state := state.toString
   let tail := state.toSubstring.drop 1
 
-  let ctxt ← getThe SymContext
-  let ctxt : SymContext ← do
-    if let some currentStateNumber := tail.toNat? then
-      pure { ctxt with
-        state_prefix := (state.get? 0).getD 's' |>.toString,
-        currentStateNumber }
-    else
-      logWarning "\
-        Expected state to be a single letter followed by a number, but found:
-          {state}
-
-        Falling back to the default numbering schema,
-        with `s1` as the first new intermediate state"
-      pure { ctxt with
-        state_prefix := "s",
-        currentStateNumber := 1 }
-  set ctxt
+  if let some currentStateNumber := tail.toNat? then
+    modifyThe SymContext ({ · with
+      state_prefix := (state.get? 0).getD 's' |>.toString,
+      currentStateNumber })
+  else
+    logWarning "\
+      Expected state to be a single letter followed by a number, but found:
+        {state}
+
+      Falling back to the default numbering schema,
+      with `s1` as the first new intermediate state"
+    modifyThe SymContext ({ · with
+      state_prefix := "s",
+      currentStateNumber := 1 })
 
 /-- Annotate any errors thrown by `k` with a local variable (and its type) -/
 private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α) :

From bc1375b97cc928adc2c14c6716a80129cb3f2051 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Mon, 23 Sep 2024 15:56:45 -0500
Subject: [PATCH 17/18] document that `inferStatePrefixAndNumber` might log a
 warning

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

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index fbd64b23..41075ce8 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -214,7 +214,7 @@ def modify (ctxt : SymContext) (k : SymM Unit) : TacticM SymContext := do
 /-- Infer `state_prefix` and `curr_state_number` from the `state` name
 as follows: if `state` is `s{i}` for some number `i` and a single character `s`,
 then `s` is the prefix and `i` the number,
-otherwise ignore `state`, and start counting from `s1` -/
+otherwise ignore `state`, log a warning, and start counting from `s1` -/
 def inferStatePrefixAndNumber : SymM Unit := do
   let state ← AxEffects.getCurrentStateName
   let state := state.toString

From 91d914d47ce8d798a7f6d3c16598b7f489804e6a Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 14:49:57 -0500
Subject: [PATCH 18/18] workaround for logWarning being slow

---
 Tactics/SymContext.lean | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/Tactics/SymContext.lean b/Tactics/SymContext.lean
index 41075ce8..e672e47f 100644
--- a/Tactics/SymContext.lean
+++ b/Tactics/SymContext.lean
@@ -129,6 +129,23 @@ instance : MonadStateOf AxEffects SymM where
     modifyThe SymContext ({· with effects})
     return a
 
+/-!
+## WORKAROUND for https://github.com/leanprover/lean4/issues/5457
+For some reason, `logWarning` is very slow to elaborate,
+so we add a specialized `SymM.logWarning` with a specific instance of `MonadLog`
+hidden behind a def. For some reason this is fast to elaborate.
+-/
+
+/-- This def may seem pointless, but it is in-fact load-bearing.
+
+Furthermore, making it an `instance` will cause `logWarning` below to be
+very slow to elaborate. Why? No clue. -/
+protected def instMonadLog : MonadLog SymM := inferInstance
+
+@[inherit_doc Lean.logWarning]
+def logWarning (msg : MessageData) : SymM Unit :=
+  @Lean.logWarning SymM _ SymM.instMonadLog _ _ msg
+
 end SymM
 
 namespace SymContext
@@ -225,7 +242,7 @@ def inferStatePrefixAndNumber : SymM Unit := do
       state_prefix := (state.get? 0).getD 's' |>.toString,
       currentStateNumber })
   else
-    logWarning "\
+    SymM.logWarning "\
       Expected state to be a single letter followed by a number, but found:
         {state}