From 430db9098fe2baf9655433b535b76ebb1875dca8 Mon Sep 17 00:00:00 2001
From: Alex Keizer <alex@keizer.dev>
Date: Thu, 26 Sep 2024 12:55:36 -0500
Subject: [PATCH] feat: implement persistent `AxEffects` updates (#182)

### Description:

Stacked on #181.

This is the first step towards proper effect aggregation. Crucially, we
don't use these updates in `sym_n` yet, as we need to integrate some
form of CSE to make it actually usable

### 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.
---
 Tactics/Reflect/AxEffects.lean | 40 ++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/Tactics/Reflect/AxEffects.lean b/Tactics/Reflect/AxEffects.lean
index 69cb5238..a0653c11 100644
--- a/Tactics/Reflect/AxEffects.lean
+++ b/Tactics/Reflect/AxEffects.lean
@@ -420,19 +420,12 @@ return an `AxEffects` where `s` is the intial state, and `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 fromExpr (e : Expr) : MetaM AxEffects := do
-  let msg := m!"Building 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 ← fromExpr e
-        eff.update_write_mem n addr val
+def fromExpr (e : Expr) : MetaM AxEffects := do
+  let s0 ← mkFreshExprMVar mkArmState
+  let eff := initial s0
+  let eff ← eff.updateWithExpr e
+  return { eff with initialState := ← instantiateMVars eff.initialState}
 
-    | w field value e =>
-        let eff ← fromExpr e
-        eff.update_w field value
-
-    | _ =>
-        return initial e
 
 /-- Given a proof `eq : s = <currentState>`,
 set `s` to be the new `currentState`, and update all proofs accordingly -/
@@ -464,25 +457,34 @@ def adjustCurrentStateWithEq (eff : AxEffects) (s eq : Expr) :
       currentState, fields, nonEffectProof, memoryEffectProof, programProof
     }
 
-/-- Given a proof `eq : ?s = <sequence of w/write_mem to ?s0>`,
+/-- Given a proof `eq : ?s = <sequence of w/write_mem to eff.currentState>`,
 where `?s` and `?s0` are arbitrary `ArmState`s,
-return an `AxEffect` with `?s0` as the initial state,
-the rhs of the equality as the current state, and
-`eq` stored as `currentStateEq`,
-so that `?s` is the public-facing current state -/
-def fromEq (eq : Expr) : MetaM AxEffects :=
+return an `AxEffect` with the rhs of the equality as the current state,
+and the (non-)effects updated accordingly -/
+def updateWithEq (eff : AxEffects) (eq : Expr) : MetaM AxEffects :=
   let msg := m!"Building effects with equality: {eq}"
   withTraceNode `Tactic.sym (fun _ => pure msg) <| do
     let s ← mkFreshExprMVar mkArmState
     let rhs ← mkFreshExprMVar mkArmState
     assertHasType eq <| mkEqArmState s rhs
 
-    let eff ← fromExpr (← instantiateMVars rhs)
+    let eff ← eff.updateWithExpr (← instantiateMVars rhs)
     let eff ← eff.adjustCurrentStateWithEq s eq
     withTraceNode `Tactic.sym (fun _ => pure "new state") do
       trace[Tactic.sym] "{eff}"
     return eff
 
+/-- Given a proof `eq : ?s = <sequence of w/write_mem to ?s0>`,
+where `?s` and `?s0` are arbitrary `ArmState`s,
+return an `AxEffect` with `?s0` as the initial state,
+the rhs of the equality as the current state,
+and the (non-)effects updated accordingly -/
+def fromEq (eq : Expr) : MetaM AxEffects := do
+  let s0 ← mkFreshExprMVar mkArmState
+  let eff := initial s0
+  let eff ← eff.updateWithEq eq
+  return { eff with initialState := ← instantiateMVars eff.initialState}
+
 /-- Given an equality `eq : ?currentProgram = ?newProgram`,
 where `currentProgram` is the rhs of `eff.programProof`,
 apply transitivity to make `newProgram` the rhs of `programProof`