From 2f2093d92ebfb5413f8cf7ca63f908b1ad8fffb0 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 21:43:01 -0500
Subject: [PATCH] refactor: reimplement initNextStep without evalTactic

---
 Arm/Exec.lean              |  7 +++++
 Tactics/Sym.lean           | 58 +++++++++++++++++++++++---------------
 Tactics/Sym/AxEffects.lean |  2 +-
 Tactics/Sym/Context.lean   |  7 +++++
 4 files changed, 51 insertions(+), 23 deletions(-)

diff --git a/Arm/Exec.lean b/Arm/Exec.lean
index 18ff7562..d4de5a67 100644
--- a/Arm/Exec.lean
+++ b/Arm/Exec.lean
@@ -162,6 +162,13 @@ theorem run_onestep {s s': ArmState} {n : Nat} :
     (s' = run (n + 1) s) → ∃ s'', stepi s = s'' ∧ s' = run n s'' := by
   simp only [run, exists_eq_left', imp_self]
 
+/-- helper lemma for automation -/
+theorem run_of_run_succ_of_stepi_eq {s0 s1 sf : ArmState} {n : Nat}
+    (h_run : sf = run (n + 1) s0)
+    (h_stepi : stepi s0 = s1) :
+    sf = run n s1 := by
+  simp_all only [run]
+
 /-- helper lemma for automation -/
 theorem stepi_eq_of_fetch_inst_of_decode_raw_inst
     (s : ArmState) (addr : BitVec 64) (rawInst : BitVec 32) (inst : ArmInst)
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 0ee58f75..e2af7599 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -161,37 +161,47 @@ def unfoldRun (whileTac : Unit → TacticM Unit) : SymM Expr := do
         setGoals (rwRes.mvarIds ++ originalGoals ++ newGoal)
         instantiateMVars runStepsPred
 
-/--
+/-- TODO: better docstring
 `initNextStep` returns expressions
-- `sn : ArmState`, and
-- `stepiEq : stepi <currentState> = sn`
+- `nextState : ArmState`, and
+- `stepiEq : stepi <currentState> = nextState`
 In that order, it also modifies `hRun` to be of type:
   `<finalState> = hRun _ sn`
 -/
 def initNextStep (whileTac : TSyntax `tactic) : SymM (Expr × Expr) :=
   withMainContext' do
-    let stepi_eq := Lean.mkIdent (.mkSimple s!"stepi_{← getCurrentStateName}")
+    let goal ← getMainGoal
 
+    -- Add next state to local context
+    let currentState ← getCurrentState
+    let nextStateVal := -- `stepi <currentState>`
+        mkApp (mkConst ``stepi) currentState
+    let (nextStateId, goal) ← do
+      let name ← getNextStateName
+      goal.note name nextStateVal mkArmState
+    let nextState := Expr.fvar nextStateId
+
+    let stepiEq : Expr := -- `stepiEq : stepi <currentState> = nextState`
+      let ty := mkEqArmState nextStateVal nextState
+      mkApp2 (.const ``id [0]) ty <| mkEqReflArmState nextState
+    let (_, goal) ← do
+      let name := Name.mkSimple s!"stepi_{← getCurrentStateName}"
+      goal.note name stepiEq
+    replaceMainGoal [goal]
+
+    -- Ensure we have one more step to simulate
     let runStepsPred ← unfoldRun (fun _ => evalTacticAndTrace whileTac)
-    -- Add new state to local context
-    let hRunId      := mkIdent <|← getHRunName
-    let nextStateId := mkIdent <|← getNextStateName
-    withMainContext' <| evalTacticAndTrace <|← `(tactic|
-      init_next_step $hRunId:ident $stepi_eq:ident $nextStateId:ident
-    )
-    withMainContext' <| do
-      -- Update `hRun`
-      let hRun := hRunId.getId
-      let some hRun := (← getLCtx).findFromUserName? hRun
-        | throwError "internal error: couldn't find {hRun}"
-      modifyThe SymContext ({ · with
-        hRun := hRun.toExpr
-      })
 
-      let sn ← SymContext.findFromUserName nextStateId.getId
-      let stepiEq ← SymContext.findFromUserName stepi_eq.getId
+    -- Change `hRun`
+    let hRun ← getHRun
+    let hRun := -- `run_of_run_succ_of_stepi_eq <hRun> <stepiEq>`
+      mkAppN (mkConst ``run_of_run_succ_of_stepi_eq) #[
+        currentState, nextState, ← getFinalState, runStepsPred,
+        hRun, stepiEq
+      ]
+    modifyThe SymContext ({ · with hRun })
 
-      return (sn.toExpr, stepiEq.toExpr)
+    return (nextState, stepiEq)
 
 /-- 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}`,
@@ -422,7 +432,7 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
 
     traceHeartbeats "symbolic simulation total"
     withCurrHeartbeats <|
-    withTraceNode "Post processing" (tag := "postProccessing") <| do
+    Sym.withTraceNode "Post processing" (tag := "postProccessing") <| do
       let c ← getThe SymContext
       -- Check if we can substitute the final state
       if c.runSteps? = some 0 then
@@ -442,6 +452,10 @@ 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]
+      else -- Replace `h_run` in the local context
+        let goal ← getMainGoal
+        let res ← goal.replace c.hRunId c.hRun
+        replaceMainGoal [res.mvarId]
 
       -- Rudimentary aggregation: we feed all the axiomatic effect hypotheses
       -- added while symbolically evaluating to `simp`
diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index 275a0cb5..1caee504 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -107,7 +107,7 @@ def getCurrentStateName : m Name := do
   @id (MetaM _) <| do
     let state ← instantiateMVars state
     let Expr.fvar id := state.consumeMData
-      | throwError "error: expected a free variable, found:\n  {state} WHHOPS"
+      | throwError "error: expected a free variable, found:\n  {state}"
     let lctx ← getLCtx
     let some decl := lctx.find? id
       | throwError "error: unknown fvar: {state}"
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 2a68c5ca..5f43dac9 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -56,6 +56,9 @@ structure SymContext where
 
   See also `SymContext.runSteps?` -/
   hRun : Expr
+  /-- The id of the variable with which `hRun` was initialized -/
+  hRunId : FVarId
+
   /-- `programInfo` is the relevant cached `ProgramInfo` -/
   programInfo : ProgramInfo
 
@@ -169,6 +172,8 @@ variable {m} [Monad m] [MonadReaderOf SymContext m]
 
 def getCurrentStateNumber : m Nat := do return (← read).currentStateNumber
 
+def getFinalState : m Expr := do return (← read).finalState
+
 /-- Return an expression of type
   `<finalState> = run <runSteps> <initialState>` -/
 def getHRun : m Expr := do return (← read).hRun
@@ -265,6 +270,7 @@ private def initial (state : Expr) : MetaM SymContext := do
   return {
     finalState
     runSteps? := none
+    hRunId := ⟨.anonymous⟩
     hRun := ← mkFreshExprMVar none
     programInfo := {
       name := `dummyValue
@@ -334,6 +340,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
       (from {hRun.toExpr} : {hRun.type})"
 
   modifyThe SymContext ({ · with
+    hRunId := hRun.fvarId
     hRun := hRun.toExpr
     finalState := ←instantiateMVars c.finalState
     runSteps?