feat: search for reads of *any* register when building an initial `Sy…

…mContext` (#200)

### Description:

Stacked on #189, we use the more powerful search capabilities to include
hypotheses that describe reads from *any* `StateField` into the initial
AxEffects.

### Testing:

What tests have been run? Did `make all` succeed for your changes? Was
conformance testing successful on an Aarch64 machine? Yes

### License:

By submitting this pull request, I confirm that my contribution is
made under the terms of the Apache 2.0 license.

---------

Co-authored-by: Shilpi Goel <[email protected]>

Benchmarks/SHA512_150.lean

@@ -12,6 +12,5 @@ open Benchmarks
 benchmark sha512_150_instructions : SHA512Bench 150 := fun s0 _ h => by
   intros
   sym_n 150
-  simp only [h, bitvec_rules]
   · exact (sorry : Aligned ..)
   done

Benchmarks/SHA512_225.lean

@@ -12,6 +12,5 @@ open Benchmarks
 benchmark sha512_225_instructions : SHA512Bench 225 := fun s0 _ h => by
   intros
   sym_n 225
-  simp only [h, bitvec_rules]
   · exact (sorry : Aligned ..)
   done

Benchmarks/SHA512_400.lean

@@ -12,6 +12,5 @@ open Benchmarks
 benchmark sha512_400_instructions : SHA512Bench 400 := fun s0 _ h => by
   intros
   sym_n 400
-  simp only [h, bitvec_rules]
   · exact (sorry : Aligned ..)
   done

Tactics/Common.lean

@@ -150,7 +150,7 @@ def reflectPFLag (e : Expr) : MetaM PFlag :=
 
 /-- Reflect a concrete `StateField` -/
 def reflectStateField (e : Expr) : MetaM StateField :=
-  match_expr e with
+  match_expr e.consumeMData with
     | StateField.GPR x  => StateField.GPR <$> reflectBitVecLiteral _ x
     | StateField.SFP x  => StateField.SFP <$> reflectBitVecLiteral _ x
     | StateField.PC     => pure StateField.PC
@@ -259,6 +259,24 @@ def mkEqArmState (x y : Expr) : Expr :=
 def mkEqReflArmState (x : Expr) : Expr :=
   mkApp2 (.const ``Eq.refl [1]) mkArmState x
 
+/-- Return `x = y` given expressions `x, y : StateField <field>` -/
+def mkEqStateValue (field x y : Expr) : Expr :=
+  let ty := mkApp (mkConst ``state_value) field
+  mkApp3 (.const ``Eq [1]) ty x y
+
+/-- Return `r <field> <state> = <value>` -/
+def mkEqReadField (field state value : Expr) : Expr :=
+  let r := mkApp2 (mkConst ``r) field state
+  mkEqStateValue field r value
+
+/-- If expression `e` is `r ?field ?state = ?value`, return
+`some (field, state, value)`, else return `none` -/
+def Lean.Expr.eqReadField? (e : Expr) : Option (Expr × Expr × Expr) := do
+  let (_ty, lhs, value) ← e.eq?
+  let_expr r field state := lhs
+    | none
+  some (field, state, value)
+
 /-! ## Tracing helpers -/
 
 def traceHeartbeats (cls : Name) (header : Option String := none) :

Tactics/Sym/Context.lean

@@ -305,7 +305,6 @@ 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.
@@ -406,13 +405,29 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
       })
     )
 
-  /- TODO(@alexkeizer): search for any other hypotheses of the form
-      `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
-    after the equality refactor  -/
+  -- Find `r ?field currentState = ?value`
+  -- NOTE: this HAS to come after the search for specific fields, like `h_pc`,
+  --       or `h_err`, to ensure those take priority and the special handling
+  --       of those fields gets applied.
+  searchLCtxFor
+    (expectedType := do
+      let field ← mkFreshExprMVar (mkConst ``StateField)
+      let value ← mkFreshExprMVar none
+      return mkEqReadField field currentState value
+    )
+    (whenFound := fun decl ty => do
+      let some (field, _state, value) := ty.eqReadField?
+        | throwError "internal error: unexpected type:\n  {ty}"
 
+      let field ← reflectStateField (← instantiateMVars field)
+      AxEffects.setFieldEffect field {
+        value := ←instantiateMVars value,
+        proof := decl.toExpr
+      }
+      return .continu
+    )
+  /- TODO(@alexkeizer): Should we search for memory as well?
+     Probably we can only do so after the memoryProof refactor -/
   return ()
 
 /-- Build a `SymContext` by searching the local context of the main goal for