From 70f2410d89277722826db8a4e549cd1623766434 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 17:30:09 -0500
Subject: [PATCH 01/14] feat: `LCtxSearch` abstraction to search the local
 context in a single pass

---
 Tactics/Sym/LCtxSearch.lean | 169 ++++++++++++++++++++++++++++++++++++
 1 file changed, 169 insertions(+)
 create mode 100644 Tactics/Sym/LCtxSearch.lean

diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
new file mode 100644
index 00000000..13541b16
--- /dev/null
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -0,0 +1,169 @@
+/-
+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
+
+/-!
+## Local Context Search
+
+In this module we build an abstraction around searching the local context for
+multiple local variables or hypotheses at the same time.
+
+The main entry point to search is `searchLocalContext`.
+`searchFor` is the main way to register which patterns to search for,
+and what actions to perform if the variable is found (or not found).
+-/
+
+variable (m) [Monad m]
+
+inductive LCtxSearchResult
+  /-- This occurence of the pattern should be ignored -/
+  | skip
+  /-- This should be counted as a successful occurence -/
+  | done
+  deriving DecidableEq
+
+structure LCtxSearchState.Pattern where
+  /-- The type to search for (up to def-eq!).
+  Notice that `expectedType` is stored as a monadic value,
+  so that we can create fresh metavariables for each search -/
+  expectedType : m Expr
+  /-- A cached result of `expectedType`;
+  this should be regenerated after every match! -/
+  cachedExpectedType : Expr
+  /-- `whenFound` will be called whenever a match for `pattern` is found,
+  with the instantiated pattern as an argument.
+  The returned `LCtxSearchResult` determines if we count this as a successful
+  occurence of the pattern, which is relevant for if `whenNotFound` is called.
+
+  NOTE: We give the (instantiated) pattern as an arg, *not* the expression that
+  we matched against. This way, implementors can recover information through
+  syntactic destructuring.
+
+  An alternative design would have `pattern : MetaM (List Expr × Expr)`,
+  where the list is intended to be a list of meta-variables, and
+  `whenFound : List Expr → Expr → m Unit`, where we would call
+  `whenFound` with the list returned by `pattern` (which has the metavariables
+  that should now have been instantiated with subexpressions of interest).
+  -/
+  whenFound : Expr → m LCtxSearchResult
+  /-- `whenNotFound` will be called if no successful occurence of the pattern
+  (as determined by the return value of `whenFound`)
+  could be found in the local context -/
+  whenNotFound : Unit → m Unit
+  /-- How many times have we (successfully) found the pattern -/
+  occurences : Nat := 0
+
+structure LCtxSearchState where
+  patterns : Array (LCtxSearchState.Pattern m)
+
+abbrev SearchLCtxForM := StateM (LCtxSearchState m)
+
+variable {m}
+
+/-- register a new expression pattern to search for:
+- `expectedType` should give an expression, with meta-variables, which is the
+  type to search for.
+
+  Note that, once a match is found, any meta-variables in `expectedType` will be
+  assigned, and thus further matches will now need to match those same concrete
+  values. That is why `expectedType` is a monadic value, which is re-evaluated
+  after each successful match.
+  If multiple matches, with distinct instantiations of a meta-variable, are
+  desired, it's important that meta-variable is created *inside* the
+  `expectedType` action.
+  If, on the other hand, a single instantiation accross all variables is
+  desired, the meta-variable should be created *outside*.
+- `whenFound` will be called whenever a local variable whose type is def-eq
+  to `expectedType` is found, with as argument the instantiated `expectedType`.
+  The return value of `whenFound` is used to determine if a match is considered
+  succesful.
+- `whenNotFound` will be called if no local variable could be found with a type
+  def-eq to the pattern, or if `whenFound` returned `skip` for all variables
+  that were found.
+
+WARNING: Once a pattern is found for which `whenFound` returns `done`, that
+particular variable will not be matched for any other patterns.
+In case of overlapping patterns, the pattern which was added first will be
+tried first
+-/
+def searchLCtxFor
+    (expectedType : m Expr)
+    (whenFound : Expr → m LCtxSearchResult)
+    (whenNotFound : Unit → m Unit := pure)
+    : SearchLCtxForM m Unit := do
+  let pattern := {
+    -- Placeholder value, since we can't evaluate `m` inside of `LCtxSearchM`
+    cachedExpectedType := .bvar 0
+    expectedType, whenFound, whenNotFound
+  }
+  modify fun state => { state with
+    patterns := state.patterns.push pattern
+  }
+
+section Run
+open Meta (isDefEq)
+variable [MonadLCtx m] [MonadLift MetaM m]
+
+/--
+Attempt to match `e` against the given pattern:
+- if `e` is def-eq to `pat.cachedExpectedType`, then return
+    the updated pattern state (with a fresh `cachedExpectedType`), and
+    the result of `whenFound`
+- Otherwise, if `e` is not def-eq, return `none`
+-/
+def LCtxSearchState.Pattern.match? (pat : Pattern m) (e : Expr) :
+    m (Option (Pattern m × LCtxSearchResult)) := do
+  if !(← isDefEq e pat.cachedExpectedType) then
+    return none
+  else
+    let cachedExpectedType ← pat.expectedType
+    let res ← pat.whenFound pat.cachedExpectedType
+    let occurences := match res with
+      | .skip => pat.occurences
+      | .done => pat.occurences + 1
+    return some ({pat with cachedExpectedType, occurences}, res)
+
+/-- Search the local context for variables of certain types, in a single pass.
+`k` is a monadic continuation that determines the patterns to search for,
+see `searchLCtxFor` to see how to register those patterns
+-/
+def searchLCtx (k : SearchLCtxForM m Unit) : m Unit := do
+  let ((), { patterns }) := StateT.run k ⟨#[]⟩
+  -- Properly initialize the `cachedExpectedType`s
+  let patterns ← patterns.mapM fun pat => do
+    pure { pat with cachedExpectedType := ← pat.expectedType }
+
+  -- We have to keep `patterns` in a Subtype to be able to prove our indexes
+  -- are valid even after mutation
+  -- TODO(@alexkeizer): consider using `Batteries.Data.Vector`, if we can
+  --                    justify a batteries dependency
+  let n := patterns.size
+  let mut patterns : { as : Array _ // as.size = n} :=
+    ⟨patterns, rfl⟩
+
+  -- The main search
+  for decl in ← getLCtx do
+    for hi : i in [0:n] do
+      have hi : i < patterns.val.size := by
+        rw [patterns.property]; get_elem_tactic
+      let pat := patterns.val[i]
+      if let some (pat, res) ← pat.match? decl.type then
+        patterns := ⟨
+          patterns.val.set ⟨i, hi⟩ pat,
+          by simp[patterns.property]
+        ⟩
+        if res = .done then
+          break -- break out of the inner loop
+
+  -- Finally, check each pattern and call `whenNotFound` if appropriate
+  for pat in patterns.val do
+    if pat.occurences = 0 then
+      pat.whenNotFound ()
+  return ()
+
+end Run

From 3e51461ee399b4a280cbb6e67146f3f39f0e561e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 17:53:08 -0500
Subject: [PATCH 02/14] feat: implement at-most-once matching, allow running
 `m` inside `SearchLCtxForM m`

---
 Tactics/Sym/LCtxSearch.lean | 38 ++++++++++++++++++++++++++-----------
 1 file changed, 27 insertions(+), 11 deletions(-)

diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
index 13541b16..48cb0190 100644
--- a/Tactics/Sym/LCtxSearch.lean
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -23,7 +23,11 @@ variable (m) [Monad m]
 inductive LCtxSearchResult
   /-- This occurence of the pattern should be ignored -/
   | skip
-  /-- This should be counted as a successful occurence -/
+  /-- This should be counted as a successful occurence,
+  and we should *continue* matching for more variables -/
+  | continu
+  /-- This should be counted as a successful occurence,
+  and we can *stop* matching against this particular pattern -/
   | done
   deriving DecidableEq
 
@@ -57,11 +61,14 @@ structure LCtxSearchState.Pattern where
   whenNotFound : Unit → m Unit
   /-- How many times have we (successfully) found the pattern -/
   occurences : Nat := 0
+  /-- Whether the pattern is active; is `isActive = false`,
+  then no further matches are attempted -/
+  isActive : Bool := true
 
 structure LCtxSearchState where
   patterns : Array (LCtxSearchState.Pattern m)
 
-abbrev SearchLCtxForM := StateM (LCtxSearchState m)
+abbrev SearchLCtxForM := StateT (LCtxSearchState m) m
 
 variable {m}
 
@@ -98,13 +105,24 @@ def searchLCtxFor
     : SearchLCtxForM m Unit := do
   let pattern := {
     -- Placeholder value, since we can't evaluate `m` inside of `LCtxSearchM`
-    cachedExpectedType := .bvar 0
+    cachedExpectedType :=← expectedType
     expectedType, whenFound, whenNotFound
   }
   modify fun state => { state with
     patterns := state.patterns.push pattern
   }
 
+/-- A wrapper around `searchLCtxFor`, which is simplified for matching at most
+one occurence of `expectedType` -/
+def searchLCtxForOnce
+    (expectedType : Expr)
+    (whenFound : Expr → m Unit)
+    (whenNotFound : Unit → m Unit := pure)
+    : SearchLCtxForM m Unit := do
+  searchLCtxFor (pure expectedType)
+    (fun e => do whenFound e; return .done)
+    whenNotFound
+
 section Run
 open Meta (isDefEq)
 variable [MonadLCtx m] [MonadLift MetaM m]
@@ -118,14 +136,16 @@ Attempt to match `e` against the given pattern:
 -/
 def LCtxSearchState.Pattern.match? (pat : Pattern m) (e : Expr) :
     m (Option (Pattern m × LCtxSearchResult)) := do
-  if !(← isDefEq e pat.cachedExpectedType) then
+  if !pat.isActive then
+    return none
+  else if !(← isDefEq e pat.cachedExpectedType) then
     return none
   else
     let cachedExpectedType ← pat.expectedType
     let res ← pat.whenFound pat.cachedExpectedType
     let occurences := match res with
       | .skip => pat.occurences
-      | .done => pat.occurences + 1
+      | .done | .continu => pat.occurences + 1
     return some ({pat with cachedExpectedType, occurences}, res)
 
 /-- Search the local context for variables of certain types, in a single pass.
@@ -133,11 +153,7 @@ def LCtxSearchState.Pattern.match? (pat : Pattern m) (e : Expr) :
 see `searchLCtxFor` to see how to register those patterns
 -/
 def searchLCtx (k : SearchLCtxForM m Unit) : m Unit := do
-  let ((), { patterns }) := StateT.run k ⟨#[]⟩
-  -- Properly initialize the `cachedExpectedType`s
-  let patterns ← patterns.mapM fun pat => do
-    pure { pat with cachedExpectedType := ← pat.expectedType }
-
+  let ((), { patterns }) ← StateT.run k ⟨#[]⟩
   -- We have to keep `patterns` in a Subtype to be able to prove our indexes
   -- are valid even after mutation
   -- TODO(@alexkeizer): consider using `Batteries.Data.Vector`, if we can
@@ -157,7 +173,7 @@ def searchLCtx (k : SearchLCtxForM m Unit) : m Unit := do
           patterns.val.set ⟨i, hi⟩ pat,
           by simp[patterns.property]
         ⟩
-        if res = .done then
+        if res = .done || res = .continu then
           break -- break out of the inner loop
 
   -- Finally, check each pattern and call `whenNotFound` if appropriate

From 8c2690a5799701f52a7bab90a65b80820e52f11e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 19:02:12 -0500
Subject: [PATCH 03/14] feat: throwNotFound helper

---
 Tactics/Sym/LCtxSearch.lean | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
index 48cb0190..ac0fabfe 100644
--- a/Tactics/Sym/LCtxSearch.lean
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -58,7 +58,7 @@ structure LCtxSearchState.Pattern where
   /-- `whenNotFound` will be called if no successful occurence of the pattern
   (as determined by the return value of `whenFound`)
   could be found in the local context -/
-  whenNotFound : Unit → m Unit
+  whenNotFound : Expr → m Unit
   /-- How many times have we (successfully) found the pattern -/
   occurences : Nat := 0
   /-- Whether the pattern is active; is `isActive = false`,
@@ -91,7 +91,8 @@ variable {m}
   succesful.
 - `whenNotFound` will be called if no local variable could be found with a type
   def-eq to the pattern, or if `whenFound` returned `skip` for all variables
-  that were found.
+  that were found. For convenience, we pass in the `expectedType` here as well.
+  See `throwNotFound` for a convenient way to throw an error here.
 
 WARNING: Once a pattern is found for which `whenFound` returns `done`, that
 particular variable will not be matched for any other patterns.
@@ -101,7 +102,7 @@ tried first
 def searchLCtxFor
     (expectedType : m Expr)
     (whenFound : Expr → m LCtxSearchResult)
-    (whenNotFound : Unit → m Unit := pure)
+    (whenNotFound : Expr → m Unit := fun _ => pure ())
     : SearchLCtxForM m Unit := do
   let pattern := {
     -- Placeholder value, since we can't evaluate `m` inside of `LCtxSearchM`
@@ -117,7 +118,7 @@ one occurence of `expectedType` -/
 def searchLCtxForOnce
     (expectedType : Expr)
     (whenFound : Expr → m Unit)
-    (whenNotFound : Unit → m Unit := pure)
+    (whenNotFound : Expr → m Unit := fun _ => pure ())
     : SearchLCtxForM m Unit := do
   searchLCtxFor (pure expectedType)
     (fun e => do whenFound e; return .done)
@@ -179,7 +180,14 @@ def searchLCtx (k : SearchLCtxForM m Unit) : m Unit := do
   -- Finally, check each pattern and call `whenNotFound` if appropriate
   for pat in patterns.val do
     if pat.occurences = 0 then
-      pat.whenNotFound ()
+      pat.whenNotFound pat.cachedExpectedType
   return ()
 
+variable [MonadError m] in
+/-- Throw an error complaining that no variable of `expectedType` could
+be found -/
+def throwNotFound (expectedType : Expr) : m Unit :=
+  throwError "Expected a local variable of type:\n  {expectedType}\n\
+    but no such variable was found in the local context"
+
 end Run

From 5c8a0696df9de2a60de589c02b18a94fcc85171e Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 19:17:21 -0500
Subject: [PATCH 04/14] fix: actually pass the LocalDecl when a match is found

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

diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
index ac0fabfe..5a0c9273 100644
--- a/Tactics/Sym/LCtxSearch.lean
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -54,7 +54,7 @@ structure LCtxSearchState.Pattern where
   `whenFound` with the list returned by `pattern` (which has the metavariables
   that should now have been instantiated with subexpressions of interest).
   -/
-  whenFound : Expr → m LCtxSearchResult
+  whenFound : LocalDecl → Expr → m LCtxSearchResult
   /-- `whenNotFound` will be called if no successful occurence of the pattern
   (as determined by the return value of `whenFound`)
   could be found in the local context -/
@@ -101,7 +101,7 @@ tried first
 -/
 def searchLCtxFor
     (expectedType : m Expr)
-    (whenFound : Expr → m LCtxSearchResult)
+    (whenFound : LocalDecl → Expr → m LCtxSearchResult)
     (whenNotFound : Expr → m Unit := fun _ => pure ())
     : SearchLCtxForM m Unit := do
   let pattern := {
@@ -117,11 +117,11 @@ def searchLCtxFor
 one occurence of `expectedType` -/
 def searchLCtxForOnce
     (expectedType : Expr)
-    (whenFound : Expr → m Unit)
+    (whenFound : LocalDecl → Expr → m Unit)
     (whenNotFound : Expr → m Unit := fun _ => pure ())
     : SearchLCtxForM m Unit := do
   searchLCtxFor (pure expectedType)
-    (fun e => do whenFound e; return .done)
+    (fun d e => do whenFound d e; return .done)
     whenNotFound
 
 section Run
@@ -135,15 +135,15 @@ Attempt to match `e` against the given pattern:
     the result of `whenFound`
 - Otherwise, if `e` is not def-eq, return `none`
 -/
-def LCtxSearchState.Pattern.match? (pat : Pattern m) (e : Expr) :
+def LCtxSearchState.Pattern.match? (pat : Pattern m) (decl : LocalDecl) :
     m (Option (Pattern m × LCtxSearchResult)) := do
   if !pat.isActive then
     return none
-  else if !(← isDefEq e pat.cachedExpectedType) then
+  else if !(← isDefEq decl.type pat.cachedExpectedType) then
     return none
   else
     let cachedExpectedType ← pat.expectedType
-    let res ← pat.whenFound pat.cachedExpectedType
+    let res ← pat.whenFound decl pat.cachedExpectedType
     let occurences := match res with
       | .skip => pat.occurences
       | .done | .continu => pat.occurences + 1
@@ -169,7 +169,7 @@ def searchLCtx (k : SearchLCtxForM m Unit) : m Unit := do
       have hi : i < patterns.val.size := by
         rw [patterns.property]; get_elem_tactic
       let pat := patterns.val[i]
-      if let some (pat, res) ← pat.match? decl.type then
+      if let some (pat, res) ← pat.match? decl then
         patterns := ⟨
           patterns.val.set ⟨i, hi⟩ pat,
           by simp[patterns.property]

From b703e6d58946b569e10d76e51a807fa549cd45a3 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Tue, 24 Sep 2024 19:17:42 -0500
Subject: [PATCH 05/14] drop dead code in `findProgramHyp`

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

diff --git a/Tactics/Common.lean b/Tactics/Common.lean
index ac1432ee..1af7140d 100644
--- a/Tactics/Common.lean
+++ b/Tactics/Common.lean
@@ -242,8 +242,7 @@ def findProgramHyp (state : Expr) : MetaM (LocalDecl × Name) := do
   -- Assert that `program` is a(n application of a) constant, and find its name
   let program := (← instantiateMVars program).getAppFn
   let .const program _ := program
-    |  -- withErrorContext h_run h_run_type <|
-        throwError "Expected a constant, found:\n\t{program}"
+    | throwError "Expected a constant, found:\n\t{program}"
 
   return ⟨h_program, program⟩
 

From 9bdcb1680662b38d162ebbd6b4fabca7958d04e8 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 14:35:32 -0500
Subject: [PATCH 06/14] refactor: move `addSimpTheorems` upwards

---
 Tactics/Sym/Context.lean | 30 ++++++++++++++++--------------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 39f6f4ed..42ce0e34 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -218,6 +218,22 @@ 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
+for effect aggregation -/
+def addSimpTheorems (c : SymContext) (simpThms : Array SimpTheorem) : SymContext :=
+  let addSimpThms := simpThms.foldl addSimpTheoremEntry
+
+  let oldSimpTheorems := c.aggregateSimpCtx.simpTheorems
+  let simpTheorems :=
+    if oldSimpTheorems.isEmpty then
+      oldSimpTheorems.push <| addSimpThms {}
+    else
+      oldSimpTheorems.modify (oldSimpTheorems.size - 1) addSimpThms
+
+  { c with aggregateSimpCtx.simpTheorems := simpTheorems }
+
 /-! ## Creating initial contexts -/
 
 /-- Modify a `SymContext` with a monadic action `k : SymM Unit` -/
@@ -431,17 +447,3 @@ def prepareForNextStep : SymM Unit := do
     runSteps?   := (· - 1) <$> c.runSteps?
     currentStateNumber := c.currentStateNumber + 1
   })
-
-/-- Add a set of new simp-theorems to the simp-theorems used
-for effect aggregation -/
-def addSimpTheorems (c : SymContext) (simpThms : Array SimpTheorem) : SymContext :=
-  let addSimpThms := simpThms.foldl addSimpTheoremEntry
-
-  let oldSimpTheorems := c.aggregateSimpCtx.simpTheorems
-  let simpTheorems :=
-    if oldSimpTheorems.isEmpty then
-      oldSimpTheorems.push <| addSimpThms {}
-    else
-      oldSimpTheorems.modify (oldSimpTheorems.size - 1) addSimpThms
-
-  { c with aggregateSimpCtx.simpTheorems := simpTheorems }

From 6192c820854bc9e28e1f87addf732713ad977761 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 14:54:15 -0500
Subject: [PATCH 07/14] add `traceNotFound` helper, fix `MonadLift` ->
 `MonadLiftT` bound

---
 Tactics/Sym/LCtxSearch.lean | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
index 5a0c9273..34e33d72 100644
--- a/Tactics/Sym/LCtxSearch.lean
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -126,7 +126,7 @@ def searchLCtxForOnce
 
 section Run
 open Meta (isDefEq)
-variable [MonadLCtx m] [MonadLift MetaM m]
+variable [MonadLCtx m] [MonadLiftT MetaM m]
 
 /--
 Attempt to match `e` against the given pattern:
@@ -190,4 +190,11 @@ def throwNotFound (expectedType : Expr) : m Unit :=
   throwError "Expected a local variable of type:\n  {expectedType}\n\
     but no such variable was found in the local context"
 
+/-- Add a message to the trace that we searched for, but couldn't find,
+a variable of `expectedType`, and continue execution. -/
+def traceNotFound (cls : Name) (expectedType : Expr) : m Unit :=
+  trace (m:=MetaM) cls fun _ =>
+    m!"Unable to find a variable of type {expectedType} in the local context"
+
+
 end Run

From 4e2a480f2b8c398864a28c6183f45f5ecb65d841 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 15:03:05 -0500
Subject: [PATCH 08/14] feat: `AxEffects.setFieldEffect` helper

---
 Tactics/Sym/AxEffects.lean | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/Tactics/Sym/AxEffects.lean b/Tactics/Sym/AxEffects.lean
index a0653c11..8c8a6aa5 100644
--- a/Tactics/Sym/AxEffects.lean
+++ b/Tactics/Sym/AxEffects.lean
@@ -231,11 +231,36 @@ def getField (eff : AxEffects) (fld : StateField) : MetaM FieldEffect :=
       let proof  ← eff.mkAppNonEffect (toExpr fld)
       pure { value, proof }
 
-variable {m} [Monad m] [MonadReaderOf AxEffects m] [MonadLiftT MetaM m] in
+section Monad
+variable {m} [Monad m] [MonadLiftT MetaM m]
+
+variable [MonadReaderOf AxEffects m] in
 @[inherit_doc getField]
 def getFieldM (field : StateField) : m FieldEffect := do
   (← read).getField field
 
+variable [MonadStateOf AxEffects m]
+
+/-- Set the effect of a specific field in the monad state, overwriting any
+previous value for that field.
+
+NOTE: the proof in `effect` is assumed to be valid for the current state,
+this is not eagerly checked (but the kernel will of course eventually reject
+a proof if it used a malformed field-effect; a mallformed proof does not
+compromise soundness, but it will cause obscure errors) -/
+def setFieldEffect (field : StateField) (effect : FieldEffect) : m Unit :=
+  modify fun eff => { eff with
+    fields := eff.fields.insert field effect }
+
+/-- Given a proof that `r .ERR <currentState> = None`, set the effect of the
+`ERR` field accordingly.
+
+This is a specialization of `setFieldEffect`. -/
+def setErrorProof (proof : Expr) : m Unit :=
+  setFieldEffect .ERR { value := mkConst ``StateError.None, proof }
+
+end Monad
+
 /-! ## Update a Reflected State -/
 
 /-- Execute `write_mem <n> <addr> <val>` against the state stored in `eff`

From 77ff5932ba108e090332155325a80e0dcdfde8b2 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 15:22:28 -0500
Subject: [PATCH 09/14] feat: use `searchLCtx` to build a SymContext

---
 Tactics/Sym/Context.lean | 219 +++++++++++++++++++++++++--------------
 1 file changed, 141 insertions(+), 78 deletions(-)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 42ce0e34..ddfa10d2 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -11,8 +11,10 @@ import Tactics.Common
 import Tactics.Attr
 import Tactics.Sym.ProgramInfo
 import Tactics.Sym.AxEffects
+import Tactics.Sym.LCtxSearch
 import Tactics.Simp
 
+
 /-!
 This files defines the `SymContext` structure,
 which collects the names of various
@@ -241,6 +243,35 @@ def modify (ctxt : SymContext) (k : SymM Unit) : TacticM SymContext := do
   let ((), ctxt) ← SymM.run ctxt k
   return ctxt
 
+private def initial (state : Expr) : MetaM SymContext := do
+  /- Create an mvar for the final state -/
+  let finalState ← mkFreshExprMVar mkArmState
+  /- Get the default simp lemmas & simprocs for aggregation -/
+  let (aggregateSimpCtx, aggregateSimprocs) ←
+    LNSymSimpContext (config := {decide := true, failIfUnchanged := false})
+  let aggregateSimpCtx := { aggregateSimpCtx with
+    -- Create a new discrtree for effect hypotheses to be added to.
+    -- TODO(@alexkeizer): I put this here, since the previous version kept
+    -- a seperate discrtree for lemmas. I should run benchmarks to see what
+    -- happens if we keep everything in one simpset.
+    simpTheorems := aggregateSimpCtx.simpTheorems.push {}
+  }
+  return {
+    finalState
+    runSteps? := none
+    h_run := `dummyValue
+    programInfo := {
+      name := `dummyValue
+      instructions := ∅
+    }
+    pc := 0
+    h_sp? := none
+    aggregateSimpCtx,
+    aggregateSimprocs,
+    effects := AxEffects.initial state
+  }
+
+
 /-- 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,
@@ -275,6 +306,109 @@ private def withErrorContext (name : Name) (type? : Option Expr) (k : MetaM α)
       | none      => m!""
     throwErrorAt e.getRef "{e.toMessageData}\n\nIn {h}{type}"
 
+-- protected def AxEffects.searchFor
+
+/-- Build the lazy search structure (for use with `searchLCtx`)
+to populate the `SymContext` state from the local context.
+
+NOTE: some search might be performed eagerly. The resulting search structure
+is tied to the specific `SymM` state and local context, it's expected that
+neither of these change between construction and execution of the search. -/
+protected def searchFor : SearchLCtxForM SymM Unit := do
+  let c ← getThe SymContext
+  let initialState ← AxEffects.getInitialState
+
+  /- We start by doing an eager search for `h_run`, outside the `SearchLCxtForM`
+  monad. This is needed to instantiate the initial state -/
+  let runSteps ← mkFreshExprMVar (Expr.const ``Nat [])
+  let hRunType := h_run_type c.finalState runSteps initialState
+  let some hRun ← findLocalDeclOfType? hRunType
+    | throwNotFound hRunType
+
+  let runSteps? ← reflectNatLiteral? runSteps
+  if runSteps?.isNone then
+    trace[Tactic.sym] "failed to reflect {runSteps} \
+      (from {hRun.toExpr} : {hRun.type})"
+
+  modifyThe SymContext ({ · with
+    h_run := hRun.userName
+    runSteps?
+  })
+
+  /-
+  From here on out, we're building the lazy search patterns
+  -/
+  -- Find `h_program : <s>.program = <program>`
+  let program ← mkFreshExprMVar none
+  searchLCtxForOnce (h_program_type initialState program)
+    (whenNotFound := throwNotFound)
+    (whenFound := fun decl _ => do
+      -- Register the program proof
+      modifyThe AxEffects ({· with
+        programProof := decl.toExpr
+      })
+      -- Assert that `program` is a(n application of a) constant
+      let program := (← instantiateMVars program).getAppFn
+      let .const program _ := program
+        | throwError "Expected a constant, found:\n\t{program}"
+      -- Retrieve the programInfo from the environment
+      let some programInfo ← ProgramInfo.lookup? program
+        | throwError "Could not find program info for `{program}`.\n\
+            Did you remember to generate step theorems with:\n  \
+              #generateStepEqTheorems {program}"
+      modifyThe SymContext ({· with
+        programInfo
+      })
+    )
+
+  -- Find `h_pc : r .PC <s> = <pc>`
+  let pc ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
+  searchLCtxForOnce (h_pc_type initialState pc)
+    (whenNotFound := throwNotFound)
+    (whenFound := fun decl _ => do
+      let pc ← instantiateMVars pc
+      -- Set the field effect
+      AxEffects.setFieldEffect .PC {
+        value := pc
+        proof := decl.toExpr
+      }
+      -- Then, reflect the value
+      let pc ← withErrorContext decl.userName decl.type <|
+        reflectBitVecLiteral 64 pc
+      modifyThe SymContext ({ · with pc })
+    )
+
+  -- Find `h_err : r .ERR <s> = .None`, or add a new subgoal if it isn't found
+  searchLCtxForOnce (h_err_type initialState)
+    (whenFound := fun decl _ =>
+      AxEffects.setErrorProof decl.toExpr
+    )
+    (whenNotFound := fun _ => do
+      let errHyp ← mkFreshExprMVar (h_err_type initialState)
+      replaceMainGoal [← getMainGoal, errHyp.mvarId!]
+      AxEffects.setErrorProof errHyp
+    )
+
+  -- Find `h_sp : CheckSPAlignment <initialState>`.
+  searchLCtxForOnce (h_sp_type initialState)
+    (whenNotFound := traceNotFound `Tactic.sym)
+    -- ^^ Note that `h_sp` is optional, so there's no need to throw an error,
+    --    we merely add a message to the trace and move on
+    (whenFound := fun decl _ =>
+      modifyThe AxEffects ({ · with
+        stackAlignmentProof? := some decl.toExpr
+      })
+    )
+
+  /- TODO(@alexkeizer): search for any other hypotheses of the form
+      `r ?field <initialState> = _`, and record those.
+    Keeping in mind that we have to do this search AFTER `h_pc` or `h_err`, to
+    ensure those field-specific searches take priority.
+    Also, maybe for memory-reads as well? Or we can hold off on that untill
+    after the equality refactor  -/
+
+  return ()
+
 /-- Build a `SymContext` by searching the local context for hypotheses of the
 required types (up-to defeq) . The local context is modified to unfold the types
 to be syntactically equal to the expected type.
@@ -297,85 +431,14 @@ def fromLocalContext (state? : Option Name) : TacticM SymContext := do
       pure (Expr.fvar decl.fvarId)
     | none => mkFreshExprMVar (Expr.const ``ArmState [])
 
-  -- Find `h_run`
-  let finalState ← mkFreshExprMVar none
-  let runSteps ← mkFreshExprMVar (Expr.const ``Nat [])
-  let h_run ←
-    findLocalDeclOfTypeOrError <| h_run_type finalState runSteps stateExpr
+  -- We create a bogus initial context
+  let c ← SymContext.initial stateExpr
+  c.modify <| do
+    searchLCtx SymContext.searchFor
+
+    let thms ← (← readThe AxEffects).toSimpTheorems
+    modifyThe SymContext (·.addSimpTheorems thms)
 
-  -- Unwrap and reflect `runSteps`
-  let runSteps? ← do
-    let msg := m!"Reflecting: {runSteps}"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| do
-      let runSteps? ← reflectNatLiteral? runSteps
-      trace[Tactic.sym] "got: {runSteps?}"
-      pure runSteps?
-  let finalState ← instantiateMVars finalState
-
-  -- 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
-
-  -- Try to find `h_program`, and infer `program` from it
-  let ⟨h_program, program⟩ ← withErrorContext h_run.userName h_run.type <|
-    findProgramHyp stateExpr
-
-  -- Then, try to find `h_pc`
-  let pcE ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
-  let h_pc ← findLocalDeclOfTypeOrError <| h_pc_type stateExpr pcE
-
-  -- Unwrap and reflect `pc`
-  let pcE ← instantiateMVars pcE
-  let pc ← withErrorContext h_pc.userName h_pc.type <|
-    reflectBitVecLiteral 64 pcE
-
-  -- Attempt to find `h_err`, adding a new subgoal if it couldn't be found
-  let errHyp ← do
-    let h_err? ← findLocalDeclOfType? (h_err_type stateExpr)
-    match h_err? with
-    | some d => pure d.toExpr
-    | none => do
-        let errHyp ← mkFreshExprMVar (h_err_type stateExpr)
-        replaceMainGoal [← getMainGoal, errHyp.mvarId!]
-        pure errHyp
-
-  let h_sp?  ← findLocalDeclOfType? (h_sp_type stateExpr)
-  if h_sp?.isNone then
-    trace[Sym] "Could not find local hypothesis of type {h_sp_type stateExpr}"
-
-  -- Finally, retrieve the programInfo from the environment
-  let some programInfo ← ProgramInfo.lookup? program
-    | throwError "Could not find program info for `{program}`.
-        Did you remember to generate step theorems with:
-          #generateStepEqTheorems {program}"
-
-  -- Initialize the axiomatic hypotheses with hypotheses for the initial state
-  let axHyps := #[h_program, h_pc] ++ h_sp?.toArray
-  let (aggregateSimpCtx, aggregateSimprocs) ←
-    LNSymSimpContext
-      (config := {decide := true, failIfUnchanged := false})
-      (decls := axHyps)
-      (exprs := #[errHyp])
-      (noIndexAtArgs := false)
-
-  -- Build an initial AxEffects
-  let effects := AxEffects.initial stateExpr
-  let effects := { effects with
-    fields := effects.fields
-      |>.insert .PC { value := pcE, proof := h_pc.toExpr}
-      |>.insert .ERR { value := mkConst ``StateError.None, proof := errHyp}
-    programProof := h_program.toExpr
-    stackAlignmentProof? := h_sp?.map (·.toExpr)
-  }
-  let c : SymContext := {
-    finalState, runSteps?, pc,
-    h_run := h_run.userName,
-    h_sp? := (·.userName) <$> h_sp?,
-    programInfo,
-    effects,
-    aggregateSimpCtx, aggregateSimprocs
-  }
-  c.modify <|
     inferStatePrefixAndNumber
 where
   findLocalDeclOfType? (expectedType : Expr) : MetaM (Option LocalDecl) := do

From 2ea3ad573f1679a0a3f2a74b44046a65c4709fbb Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 16:22:42 -0500
Subject: [PATCH 10/14] fix: remember to update `h_sp?` as well

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

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index ddfa10d2..4dea5b4a 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -398,6 +398,9 @@ 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

From f7c72743b2d447bfa845e7a5173cf94d3c0b41f6 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 16:23:19 -0500
Subject: [PATCH 11/14] fix: instantiate finalState after matching h_run

---
 Tactics/Sym/Context.lean | 1 +
 1 file changed, 1 insertion(+)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 4dea5b4a..e6f897d5 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -332,6 +332,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
 
   modifyThe SymContext ({ · with
     h_run := hRun.userName
+    finalState := ←instantiateMVars c.finalState
     runSteps?
   })
 

From 0670d9fb0cb56bded3d461d4ec37f6b6979c4bf3 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 16:35:47 -0500
Subject: [PATCH 12/14] add option for type canonicalization to LCtxSearch

---
 Tactics/Sym/Context.lean    | 62 ++++++-------------------------------
 Tactics/Sym/LCtxSearch.lean | 28 ++++++++++++-----
 2 files changed, 31 insertions(+), 59 deletions(-)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index e6f897d5..cacf25d6 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -271,7 +271,6 @@ private def initial (state : Expr) : MetaM SymContext := do
     effects := AxEffects.initial state
   }
 
-
 /-- 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,
@@ -365,6 +364,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
   -- Find `h_pc : r .PC <s> = <pc>`
   let pc ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
   searchLCtxForOnce (h_pc_type initialState pc)
+    (changeType := true)
     (whenNotFound := throwNotFound)
     (whenFound := fun decl _ => do
       let pc ← instantiateMVars pc
@@ -381,6 +381,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
 
   -- Find `h_err : r .ERR <s> = .None`, or add a new subgoal if it isn't found
   searchLCtxForOnce (h_err_type initialState)
+    (changeType := true)
     (whenFound := fun decl _ =>
       AxEffects.setErrorProof decl.toExpr
     )
@@ -392,10 +393,11 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
 
   -- Find `h_sp : CheckSPAlignment <initialState>`.
   searchLCtxForOnce (h_sp_type initialState)
+    (changeType := true)
     (whenNotFound := traceNotFound `Tactic.sym)
     -- ^^ Note that `h_sp` is optional, so there's no need to throw an error,
     --    we merely add a message to the trace and move on
-    (whenFound := fun decl _ =>
+    (whenFound := fun decl _ => do
       modifyThe AxEffects ({ · with
         stackAlignmentProof? := some decl.toExpr
       })
@@ -422,7 +424,7 @@ 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
+  withTraceNode `Tactic.sym (fun _ => pure msg) <| withMainContext' do
   trace[Tactic.Sym] "state? := {state?}"
   let lctx ← getLCtx
 
@@ -440,55 +442,11 @@ def fromLocalContext (state? : Option Name) : TacticM SymContext := do
   c.modify <| do
     searchLCtx SymContext.searchFor
 
-    let thms ← (← readThe AxEffects).toSimpTheorems
-    modifyThe SymContext (·.addSimpTheorems thms)
-
-    inferStatePrefixAndNumber
-where
-  findLocalDeclOfType? (expectedType : Expr) : MetaM (Option LocalDecl) := do
-    let msg := m!"Searching for hypothesis of type: {expectedType}"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| do
-      let decl? ← _root_.findLocalDeclOfType? expectedType
-      trace[Tactic.sym] "Found: {(·.toExpr) <$> decl?}"
-      return decl?
-  findLocalDeclOfTypeOrError (expectedType : Expr) : MetaM LocalDecl := do
-    let msg := m!"Searching for hypothesis of type: {expectedType}"
-    withTraceNode `Tactic.sym (fun _ => pure msg) <| do
-      let decl ← _root_.findLocalDeclOfTypeOrError expectedType
-      trace[Tactic.sym] "Found: {decl.toExpr}"
-      return decl
-
-/-! ## Massaging the local context -/
-
-/-- 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
-described in the relevant docstrings.
-
-That is, (un)fold types which were definitionally, but not syntactically,
-equal to the expected shape. -/
-def canonicalizeHypothesisTypes : SymReaderM Unit := withMainContext' do
-  let c ← readThe SymContext
-  let lctx ← getLCtx
-  let mut goal ← getMainGoal
-  let state := c.effects.currentState
-
-  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_sp := c.h_sp? then
-    hyps := hyps.push (h_sp, h_sp_type state)
-
-  let mut hypIds ← hyps.mapM fun ⟨name, type⟩ => do
-    let some decl := lctx.findFromUserName? name
-      | throwError "Unknown local hypothesis `{name}`"
-    pure (decl.fvarId, type)
-
-  let errHyp ← AxEffects.getFieldM .ERR
-  if let Expr.fvar id := errHyp.proof then
-    hypIds := hypIds.push (id, h_err_type state)
-  for ⟨fvarId, type⟩ in hypIds do
-    goal ← goal.replaceLocalDeclDefEq fvarId type
-  replaceMainGoal [goal]
+    withMainContext' <| do
+      let thms ← (← readThe AxEffects).toSimpTheorems
+      modifyThe SymContext (·.addSimpTheorems thms)
+
+      inferStatePrefixAndNumber
 
 /-! ## Incrementing the context to the next state -/
 
diff --git a/Tactics/Sym/LCtxSearch.lean b/Tactics/Sym/LCtxSearch.lean
index 34e33d72..58d4c3ea 100644
--- a/Tactics/Sym/LCtxSearch.lean
+++ b/Tactics/Sym/LCtxSearch.lean
@@ -59,6 +59,8 @@ structure LCtxSearchState.Pattern where
   (as determined by the return value of `whenFound`)
   could be found in the local context -/
   whenNotFound : Expr → m Unit
+  /-- Whether to change the type of successful matches -/
+  changeType : Bool
   /-- How many times have we (successfully) found the pattern -/
   occurences : Nat := 0
   /-- Whether the pattern is active; is `isActive = false`,
@@ -93,6 +95,9 @@ variable {m}
   def-eq to the pattern, or if `whenFound` returned `skip` for all variables
   that were found. For convenience, we pass in the `expectedType` here as well.
   See `throwNotFound` for a convenient way to throw an error here.
+- If `changeType` is set to true, then we change the type of every successful
+  match (i.e., `whenFound` returns `continu` or `done`) to be exactly the
+  `expectedType`
 
 WARNING: Once a pattern is found for which `whenFound` returns `done`, that
 particular variable will not be matched for any other patterns.
@@ -103,30 +108,35 @@ def searchLCtxFor
     (expectedType : m Expr)
     (whenFound : LocalDecl → Expr → m LCtxSearchResult)
     (whenNotFound : Expr → m Unit := fun _ => pure ())
+    (changeType : Bool := false)
     : SearchLCtxForM m Unit := do
   let pattern := {
     -- Placeholder value, since we can't evaluate `m` inside of `LCtxSearchM`
     cachedExpectedType :=← expectedType
-    expectedType, whenFound, whenNotFound
+    expectedType, whenFound, whenNotFound, changeType
   }
   modify fun state => { state with
     patterns := state.patterns.push pattern
   }
 
 /-- A wrapper around `searchLCtxFor`, which is simplified for matching at most
-one occurence of `expectedType` -/
+one occurence of `expectedType`.
+
+See `searchLCtxFor` for an explanation of the arguments -/
 def searchLCtxForOnce
     (expectedType : Expr)
     (whenFound : LocalDecl → Expr → m Unit)
     (whenNotFound : Expr → m Unit := fun _ => pure ())
+    (changeType : Bool := false)
     : SearchLCtxForM m Unit := do
   searchLCtxFor (pure expectedType)
     (fun d e => do whenFound d e; return .done)
-    whenNotFound
+    whenNotFound changeType
 
 section Run
+open Elab.Tactic
 open Meta (isDefEq)
-variable [MonadLCtx m] [MonadLiftT MetaM m]
+variable [MonadLCtx m] [MonadLiftT MetaM m] [MonadLiftT TacticM m]
 
 /--
 Attempt to match `e` against the given pattern:
@@ -144,9 +154,13 @@ def LCtxSearchState.Pattern.match? (pat : Pattern m) (decl : LocalDecl) :
   else
     let cachedExpectedType ← pat.expectedType
     let res ← pat.whenFound decl pat.cachedExpectedType
-    let occurences := match res with
-      | .skip => pat.occurences
-      | .done | .continu => pat.occurences + 1
+    let mut occurences := pat.occurences
+    if res != .skip then
+      occurences := occurences + 1
+      if pat.changeType = true then
+        let goal ← getMainGoal
+        let goal ← goal.replaceLocalDeclDefEq decl.fvarId pat.cachedExpectedType
+        replaceMainGoal [goal]
     return some ({pat with cachedExpectedType, occurences}, res)
 
 /-- Search the local context for variables of certain types, in a single pass.

From e361ce1e876266bfc556d7ee8a7007cda36bbe16 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Wed, 25 Sep 2024 16:50:12 -0500
Subject: [PATCH 13/14] fix: remove reference to `canonicalizeHypothesisType`,
 rename `fromLocalContext` to `fromMainContext` to reflect that we now use
 `withMainContext`

---
 Tactics/Sym.lean         |  9 +++------
 Tactics/Sym/Context.lean | 11 ++++++-----
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 5a62f834..b0d5621f 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -244,7 +244,7 @@ elab "explode_step" h_step:term " at " state:term : tactic => withMainContext do
   let .fvar stateFVar := state
     | throwError "Expected fvar, found {state}"
   let stateDecl := (← getLCtx).get! stateFVar
-  let c ← SymContext.fromLocalContext (some stateDecl.userName)
+  let c ← SymContext.fromMainContext (some stateDecl.userName)
   let _ ← SymM.run c <| explodeStep hStep
 
 /--
@@ -403,11 +403,8 @@ elab "sym_n" whileTac?:(sym_while)? n:num s:(sym_at)? : tactic => do
         omega;
         ))
 
-  let c ← withMainContext <| SymContext.fromLocalContext s
-  SymM.run' c <| do
-    -- Context preparation
-    canonicalizeHypothesisTypes
-
+  let c ← SymContext.fromMainContext s
+  SymM.run' c <| withMainContext' <|  do
     -- Check pre-conditions
     assertStepTheoremsGenerated
     let n ← ensureAtMostRunSteps n.getNat
diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index cacf25d6..5ca4872b 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -415,14 +415,15 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
 
   return ()
 
-/-- Build a `SymContext` by searching the local context for hypotheses of the
-required types (up-to defeq) . The local context is modified to unfold the types
-to be syntactically equal to the expected type.
+/-- Build a `SymContext` by searching the local context of the main goal for
+hypotheses of the required types (up-to defeq).
+The local context is modified to unfold the types to be syntactically equal to
+the expected types.
 
 If an hypothesis `h_err : r <state> .ERR = None` is not found,
-we create a new subgoal of this type
+we create a new subgoal of this type.
 -/
-def fromLocalContext (state? : Option Name) : TacticM SymContext := do
+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
   trace[Tactic.Sym] "state? := {state?}"

From 7e6b84fde4d7b2cc82de42adef071a7bba16fd35 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 17:13:01 -0500
Subject: [PATCH 14/14] fix: base search on the current state, rather than the
 initial state

---
 Tactics/Sym/Context.lean | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/Tactics/Sym/Context.lean b/Tactics/Sym/Context.lean
index 5ca4872b..e38c1ee0 100644
--- a/Tactics/Sym/Context.lean
+++ b/Tactics/Sym/Context.lean
@@ -315,12 +315,12 @@ is tied to the specific `SymM` state and local context, it's expected that
 neither of these change between construction and execution of the search. -/
 protected def searchFor : SearchLCtxForM SymM Unit := do
   let c ← getThe SymContext
-  let initialState ← AxEffects.getInitialState
+  let currentState ← AxEffects.getCurrentState
 
   /- We start by doing an eager search for `h_run`, outside the `SearchLCxtForM`
   monad. This is needed to instantiate the initial state -/
   let runSteps ← mkFreshExprMVar (Expr.const ``Nat [])
-  let hRunType := h_run_type c.finalState runSteps initialState
+  let hRunType := h_run_type c.finalState runSteps currentState
   let some hRun ← findLocalDeclOfType? hRunType
     | throwNotFound hRunType
 
@@ -340,7 +340,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
   -/
   -- Find `h_program : <s>.program = <program>`
   let program ← mkFreshExprMVar none
-  searchLCtxForOnce (h_program_type initialState program)
+  searchLCtxForOnce (h_program_type currentState program)
     (whenNotFound := throwNotFound)
     (whenFound := fun decl _ => do
       -- Register the program proof
@@ -363,7 +363,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
 
   -- Find `h_pc : r .PC <s> = <pc>`
   let pc ← mkFreshExprMVar (← mkAppM ``BitVec #[toExpr 64])
-  searchLCtxForOnce (h_pc_type initialState pc)
+  searchLCtxForOnce (h_pc_type currentState pc)
     (changeType := true)
     (whenNotFound := throwNotFound)
     (whenFound := fun decl _ => do
@@ -380,19 +380,19 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
     )
 
   -- Find `h_err : r .ERR <s> = .None`, or add a new subgoal if it isn't found
-  searchLCtxForOnce (h_err_type initialState)
+  searchLCtxForOnce (h_err_type currentState)
     (changeType := true)
     (whenFound := fun decl _ =>
       AxEffects.setErrorProof decl.toExpr
     )
     (whenNotFound := fun _ => do
-      let errHyp ← mkFreshExprMVar (h_err_type initialState)
+      let errHyp ← mkFreshExprMVar (h_err_type currentState)
       replaceMainGoal [← getMainGoal, errHyp.mvarId!]
       AxEffects.setErrorProof errHyp
     )
 
-  -- Find `h_sp : CheckSPAlignment <initialState>`.
-  searchLCtxForOnce (h_sp_type initialState)
+  -- Find `h_sp : CheckSPAlignment <currentState>`.
+  searchLCtxForOnce (h_sp_type currentState)
     (changeType := true)
     (whenNotFound := traceNotFound `Tactic.sym)
     -- ^^ Note that `h_sp` is optional, so there's no need to throw an error,
@@ -407,7 +407,7 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
     )
 
   /- TODO(@alexkeizer): search for any other hypotheses of the form
-      `r ?field <initialState> = _`, and record those.
+      `r ?field <currentState> = _`, and record those.
     Keeping in mind that we have to do this search AFTER `h_pc` or `h_err`, to
     ensure those field-specific searches take priority.
     Also, maybe for memory-reads as well? Or we can hold off on that untill