BLOCKED: refactor: replace rewrite with mkEqNDRec

Proofs/Popcount32.lean

@@ -70,7 +70,6 @@ def popcount32_program : Program :=
 
 #genStepEqTheorems popcount32_program
 
-set_option trace.simp_mem.info true in
 theorem popcount32_sym_meets_spec (s0 sf : ArmState)
   (h_s0_pc         : read_pc s0 = 0x4005b4#64)
   (h_s0_program    : s0.program = popcount32_program)

Tactics/Common.lean

@@ -277,6 +277,35 @@ def Lean.Expr.eqReadField? (e : Expr) : Option (Expr × Expr × Expr) := do
     | none
   some (field, state, value)
 
+/-- Return `ArmState.program <state> = <program>` -/
+def mkEqProgram (state program : Expr) : Expr :=
+  mkApp3 (.const ``Eq [1]) (mkConst ``Program)
+    (mkApp (mkConst ``ArmState.program) state)
+    program
+
+/-- Return `BitVec n` given an expression `n : Nat` -/
+@[inline] def mkBitVec (n : Expr) : Expr :=
+  mkApp (mkConst ``BitVec) n
+
+/-- Return `x = y`, given expressions `x, y : BitVec <n>` -/
+def mkEqBitVec (n x y : Expr) : Expr :=
+  mkApp3 (.const ``Eq [1]) (mkBitVec n) x y
+
+/-- Return `read_mem_bytes <n> <addr> <state>` -/
+def mkReadMemBytes (n addr state : Expr) : Expr :=
+  mkApp3 (mkConst ``read_mem_bytes) n addr state
+
+/-- Return `read_mem_bytes <n> <addr> <state> = <value>`, given expressions
+`n : Nat`, `addr : BitVec 64`, `state : ArmState` and `value : BitVec (n*8)` -/
+def mkEqReadMemBytes (n addr state value : Expr) : Expr :=
+  let n8 := mkNatMul n (toExpr 8)
+  mkEqBitVec n8 (mkReadMemBytes n addr state) value
+
+-- def mkForallReadMemBytesEqReadMemBytes (leftState rightState : Expr) : Expr :=
+  -- TODO
+
+-- def mkForallEqReadMem
+
 /-! ## Tracing helpers -/
 
 def traceHeartbeats (cls : Name) (header : Option String := none) :

Tactics/Sym/AxEffects.lean

@@ -70,7 +70,9 @@ structure AxEffects where
       = read_mem_bytes n addr <memoryEffect>
     ``` -/
   memoryEffectProof : Expr
-  /-- A proof that `<currentState>.program = <initialState>.program` -/
+  /-- The program of the current state, see `programProof` -/
+  program : Expr
+  /-- A proof that `<currentState>.program = <program>` -/
   programProof : Expr
   /-- An optional proof of `CheckSPAlignment <currentState>`.
 
@@ -141,6 +143,8 @@ def initial (state : Expr) : AxEffects where
         mkApp2 (.const ``Eq.refl [1])
           (mkApp (mkConst ``BitVec) <| mkNatMul (.bvar 1) (toExpr 8))
           (mkApp3 (mkConst ``read_mem_bytes) (.bvar 1) (.bvar 0) state)
+  program := -- `ArmState.program <initialState>`
+    mkApp (mkConst ``ArmState.program) state
   programProof      :=
     -- `rfl`
     mkAppN (.const ``Eq.refl [1]) #[
@@ -487,11 +491,18 @@ def fromExpr (e : Expr) : MetaM AxEffects := do
   let eff ← eff.updateWithExpr e
   return { eff with initialState := ← instantiateMVars eff.initialState}
 
+
+/-- The expression `fun s => r <field> s = <value>` -/
+protected def fieldMotive (field value : Expr) : Expr :=
+  let body := mkEqReadField field (.bvar 0) value
+  Expr.lam `s mkArmState body (.default)
+open AxEffects (fieldMotive)
+
 /-- Given a proof `eq : s = <currentState>`,
 set `s` to be the new `currentState`, and update all proofs accordingly -/
 def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
     MetaM AxEffects := do
-  withTraceNode m!"adjustCurrentStateWithEq" (tag := "adjustCurrentStateWithEq") do
+  Sym.withTraceNode m!"adjustCurrentStateWithEq" (tag := "adjustCurrentStateWithEq") do
     trace[Tactic.sym] "rewriting along {eq}"
     eff.traceCurrentState
 
@@ -503,19 +514,47 @@ def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
     let fields ← eff.fields.toList.mapM fun (field, fieldEff) => do
       withTraceNode m!"rewriting field {field}" (tag := "rewriteField") do
         trace[Tactic.sym] "original proof: {fieldEff.proof}"
-        let proof ← rewriteType fieldEff.proof eq
+        let motive := fieldMotive (toExpr field) fieldEff.value
+        let proof ← mkEqNDRec motive fieldEff.proof eq
         trace[Tactic.sym] "new proof: {proof}"
         pure (field, {fieldEff with proof})
     let fields := .ofList fields
 
-    withTraceNode m!"rewriting other proofs" (tag := "rewriteMisc") <| do
-      let nonEffectProof ← rewriteType eff.nonEffectProof eq
-      let memoryEffectProof ← rewriteType eff.memoryEffectProof eq
-      -- ^^ TODO: what happens if `memoryEffect` is the same as `currentState`?
-      --    Presumably, we would *not* want to encapsulate `memoryEffect` here
-      let programProof ← rewriteType eff.programProof eq
-      let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM
-        (rewriteType · eq)
+    Sym.withTraceNode m!"rewriting other proofs" (tag := "rewriteMisc") <| do
+      let nonEffectProof ← lambdaTelescope eff.nonEffectProof fun args proof => do
+        let f := args[0]!
+        let motive ← -- `fun s => r <f> s = r <f> <eff.initialState>`
+          withLocalDeclD `s mkArmState <| fun s =>
+            let eq := mkEqStateValue f
+              (mkApp2 (mkConst ``r) f s)
+              (mkApp2 (mkConst ``r) f eff.initialState)
+            mkLambdaFVars #[s] eq
+        mkLambdaFVars args <|← mkEqNDRec motive proof eq
+
+      let memoryMotive : Expr :=
+        Expr.lam `s mkArmState (binderInfo := .default) <|
+        Expr.forallE `n (mkConst ``Nat) (binderInfo := .default) <|
+        Expr.forallE `addr (mkBitVec <| toExpr 64) (binderInfo := .default) <|
+          let s := Expr.bvar 2
+          let n := Expr.bvar 1
+          let addr := Expr.bvar 0
+          let lhs := mkReadMemBytes n addr s
+          let rhs := mkReadMemBytes n addr eff.memoryEffect
+          mkEqBitVec (mkNatMul n (toExpr 8)) lhs rhs
+      let memoryEffectProof ← mkEqNDRec memoryMotive eff.memoryEffectProof eq
+
+      let programMotive : Expr :=
+        Expr.lam `s mkArmState (binderInfo := .default) <|
+          let s := Expr.bvar 0
+          mkEqProgram s eff.program
+      let programProof ← mkEqNDRec programMotive eff.programProof eq
+
+      let stackAlignmentProof? ← eff.stackAlignmentProof?.mapM fun proof => do
+        let motive :=
+          Expr.lam `s mkArmState (binderInfo := .default) <|
+            let s := Expr.bvar 0
+            mkApp (mkConst ``CheckSPAlignment) s
+        mkEqNDRec motive proof eq
 
       return { eff with
         currentState, fields, nonEffectProof, memoryEffectProof, programProof,

Tactics/Sym/Context.lean

@@ -333,12 +333,14 @@ protected def searchFor : SearchLCtxForM SymM Unit := do
   searchLCtxForOnce (h_program_type currentState program)
     (whenNotFound := throwNotFound)
     (whenFound := fun decl _ => do
+      let program ← instantiateMVars program
       -- Register the program proof
       modifyThe AxEffects ({· with
+        program
         programProof := decl.toExpr
       })
       -- Assert that `program` is a(n application of a) constant
-      let program := (← instantiateMVars program).getAppFn
+      let program := program.getAppFn
       let .const program _ := program
         | throwError "Expected a constant, found:\n\t{program}"
       -- Retrieve the programInfo from the environment