From 6b0d380d18be18681b03f8e0e56b47c225f209f0 Mon Sep 17 00:00:00 2001
From: Shilpi Goel <shigoel@gmail.com>
Date: Fri, 1 Mar 2024 16:16:21 -0600
Subject: [PATCH] Add a simp attribute to control the lemmas used to simplify
 terms involving state accessors and updaters.

---
 Arm/Attr.lean                                 | 21 ++++++
 Arm/Insts/DPSFP/Advanced_simd_three_same.lean | 11 +--
 Arm/Map.lean                                  | 25 +++++--
 Arm/Memory.lean                               | 17 ++---
 Arm/MemoryProofs.lean                         | 29 +++++---
 Arm/State.lean                                | 69 +++++++++++--------
 Proofs/Sha512_block_armv8.lean                | 14 +++-
 Tactics/Sym.lean                              |  8 +--
 lean-toolchain                                |  2 +-
 9 files changed, 137 insertions(+), 59 deletions(-)
 create mode 100644 Arm/Attr.lean

diff --git a/Arm/Attr.lean b/Arm/Attr.lean
new file mode 100644
index 00000000..4de079b3
--- /dev/null
+++ b/Arm/Attr.lean
@@ -0,0 +1,21 @@
+import Lean
+
+-- Non-interference lemmas for simplifying terms involving state
+-- accessors and updaters.
+register_simp_attr state_simp_rules
+
+syntax "state_simp" : tactic
+macro_rules
+  | `(tactic| state_simp) => `(tactic| simp only [state_simp_rules])
+
+syntax "state_simp?" : tactic
+macro_rules
+  | `(tactic| state_simp?) => `(tactic| simp? only [state_simp_rules])
+
+syntax "state_simp_all" : tactic
+macro_rules
+  | `(tactic| state_simp_all) => `(tactic| simp_all only [state_simp_rules])
+
+syntax "state_simp_all?" : tactic
+macro_rules
+  | `(tactic| state_simp_all?) => `(tactic| simp_all? only [state_simp_rules])
diff --git a/Arm/Insts/DPSFP/Advanced_simd_three_same.lean b/Arm/Insts/DPSFP/Advanced_simd_three_same.lean
index 88520c06..3d9bf265 100644
--- a/Arm/Insts/DPSFP/Advanced_simd_three_same.lean
+++ b/Arm/Insts/DPSFP/Advanced_simd_three_same.lean
@@ -113,11 +113,14 @@ theorem pc_of_exec_advanced_simd_three_same
   -- (r StateField.PC s) + 4#64 -- TODO: How do I use + here?
   (BitVec.add (r StateField.PC s) 4#64) := by
   simp_all!
-  simp [exec_advanced_simd_three_same, exec_binary_vector, exec_logic_vector]
+  simp only [exec_advanced_simd_three_same, exec_binary_vector, 
+             Bool.and_eq_true, beq_iff_eq, binary_vector_op,
+             ofNat_eq_ofNat, zero_eq, exec_logic_vector, 
+             logic_vector_op]
   split
-  · split <;> simp
-  · simp
-  · simp
+  · split <;> state_simp
+  · state_simp
+  · state_simp
 
 ----------------------------------------------------------------------
 
diff --git a/Arm/Map.lean b/Arm/Map.lean
index 21194148..ff18287e 100644
--- a/Arm/Map.lean
+++ b/Arm/Map.lean
@@ -115,8 +115,17 @@ def Map.size (m : Map α β) : Nat :=
 @[simp] theorem Map.size_erase_le [DecidableEq α] (m : Map α β) (a : α) : (m.erase a).size ≤ m.size := by
   induction m <;> simp [erase, size] at *
   split
-  next => omega
-  next => simp; omega
+  next => 
+    -- (FIXME) This could be discharged by omega in
+    -- leanprover/lean4:nightly-2024-02-24, but not in
+    -- leanprover/lean4:nightly-2024-03-01.
+    exact Nat.le_succ_of_le (by assumption)
+  next => 
+    simp; 
+    -- (FIXME) This could be discharged by omega in
+    -- leanprover/lean4:nightly-2024-02-24, but not in
+    -- leanprover/lean4:nightly-2024-03-01.  
+    exact Nat.succ_le_succ (by assumption)
 
 @[simp] theorem Map.size_erase_eq [DecidableEq α] (m : Map α β) (a : α) : m.contains a = false → (m.erase a).size = m.size := by
   induction m <;> simp [erase, size] at *
@@ -127,5 +136,13 @@ def Map.size (m : Map α β) : Nat :=
   induction m <;> simp [erase, size, contains, find?] at *
   next head tail ih =>
   split
-  next => have := Map.size_erase_le tail a; omega
-  next he => simp [he] at h; simp [h] at ih; simp; omega
+  next => have := Map.size_erase_le tail a; 
+          -- (FIXME) This could be discharged by omega in
+          -- leanprover/lean4:nightly-2024-02-24, but not in
+          -- leanprover/lean4:nightly-2024-03-01.
+          exact Nat.lt_succ_of_le this
+  next he => simp [he] at h; simp [h] at ih; simp; 
+          -- (FIXME) This could be discharged by omega in
+          -- leanprover/lean4:nightly-2024-02-24, but not in
+          -- leanprover/lean4:nightly-2024-03-01.
+             exact Nat.succ_lt_succ ih
diff --git a/Arm/Memory.lean b/Arm/Memory.lean
index 5fcbe25b..93bff238 100644
--- a/Arm/Memory.lean
+++ b/Arm/Memory.lean
@@ -54,13 +54,13 @@ theorem r_of_write_mem : r fld (write_mem addr val s) = r fld s := by
   unfold write_mem
   split <;> simp
 
-@[simp]
+@[state_simp_rules]
 theorem r_of_write_mem_bytes :
   r fld (write_mem_bytes n addr val s) = r fld s := by
   induction n generalizing addr s
   case succ =>
     rename_i n n_ih
-    unfold write_mem_bytes; simp
+    unfold write_mem_bytes; simp only
     rw [n_ih, r_of_write_mem]
   case zero => rfl
   done
@@ -70,14 +70,14 @@ theorem fetch_inst_of_write_mem :
   unfold fetch_inst write_mem
   simp
 
-@[simp]
+@[state_simp_rules]
 theorem fetch_inst_of_write_mem_bytes :
   fetch_inst addr1 (write_mem_bytes n addr2 val s) = fetch_inst addr1 s := by
   induction n generalizing addr2 s
   case zero => rfl
   case succ =>
     rename_i n n_ih
-    unfold write_mem_bytes; simp
+    unfold write_mem_bytes; simp only
     rw [n_ih, fetch_inst_of_write_mem]
   done
 
@@ -88,26 +88,27 @@ theorem read_mem_of_w :
   unfold write_base_pc write_base_flag write_base_error
   split <;> simp
 
-@[simp]
+@[state_simp_rules]
 theorem read_mem_bytes_of_w :
   read_mem_bytes n addr (w fld v s) = read_mem_bytes n addr s := by
   induction n generalizing addr s
   case zero => rfl
   case succ =>
     rename_i n n_ih
-    unfold read_mem_bytes; simp [read_mem_of_w]
+    unfold read_mem_bytes; simp only [read_mem_of_w]
     rw [n_ih]
   done
 
+@[state_simp_rules]
 theorem write_mem_bytes_program {n : Nat} (addr : BitVec 64) (bytes : BitVec (n * 8)):
     (write_mem_bytes n addr bytes s).program = s.program := by
   intros
   induction n generalizing addr s
   · simp [write_mem_bytes]
   · rename_i n h_n
-    simp [write_mem_bytes]
+    simp only [write_mem_bytes]
     rw [h_n]
-    simp [write_mem]
+    simp only [write_mem]
 
 ---- Memory RoW/WoW lemmas ----
 
diff --git a/Arm/MemoryProofs.lean b/Arm/MemoryProofs.lean
index 96e45eb9..2a0bf7b0 100644
--- a/Arm/MemoryProofs.lean
+++ b/Arm/MemoryProofs.lean
@@ -82,7 +82,7 @@ theorem append_byte_of_extract_rest_same_cast (n : Nat) (v : BitVec ((n + 1) * 8
   · omega
   done
 
-@[simp]
+@[state_simp_rules]
 theorem read_mem_bytes_of_write_mem_bytes_same (hn1 : n <= 2^64) :
   read_mem_bytes n addr (write_mem_bytes n addr v s) = v := by
   by_cases hn0 : n = 0
@@ -127,7 +127,7 @@ theorem read_mem_bytes_of_write_mem_bytes_same (hn1 : n <= 2^64) :
 ----------------------------------------------------------------------
 -- Key theorem: read_mem_bytes_of_write_mem_bytes_different
 
-@[simp]
+@[state_simp_rules]
 theorem read_mem_bytes_of_write_mem_bytes_different
   (hn1 : n1 <= 2^64) (hn2 : n2 <= 2^64)
   (h : mem_separate addr1 (addr1 + (n1 - 1)#64) addr2 (addr2 + (n2 - 1)#64)) :
@@ -208,7 +208,7 @@ theorem write_mem_of_write_mem_bytes_commute
       · omega
   done
 
-@[simp]
+@[state_simp_rules]
 theorem write_mem_bytes_of_write_mem_bytes_commute
   (h1 : n1 <= 2^64) (h2 : n2 <= 2^64)
   (h3 : mem_separate addr2 (addr2 + (n2 - 1)#64) addr1 (addr1 + (n1 - 1)#64)) :
@@ -247,7 +247,7 @@ theorem write_mem_bytes_of_write_mem_bytes_commute
 -- Key theorems: write_mem_bytes_of_write_mem_bytes_shadow_same_region
 -- and write_mem_bytes_of_write_mem_bytes_shadow_general
 
-@[simp]
+@[state_simp_rules]
 theorem write_mem_bytes_of_write_mem_bytes_shadow_same_region
   (h : n <= 2^64) :
   write_mem_bytes n addr val2 (write_mem_bytes n addr val1 s) =
@@ -473,7 +473,7 @@ private theorem write_mem_bytes_of_write_mem_bytes_shadow_general_n2_eq
           · omega
         · exact h₁
 
-@[simp]
+@[state_simp_rules]
 theorem write_mem_bytes_of_write_mem_bytes_shadow_general
   (h1u : n1 <= 2^64) (h2l : 0 < n2) (h2u : n2 <= 2^64)
   (h3 : mem_subset addr1 (addr1 + (n1 - 1)#64) addr2 (addr2 + (n2 - 1)#64)) :
@@ -803,7 +803,7 @@ private theorem read_mem_bytes_of_write_mem_bytes_subset_n2_lt
     (extractLsb ((((addr2 - addr1).toNat + n2) * 8) - 1) ((addr2 - addr1).toNat * 8) val) := by
   induction n2, h2 using Nat.le_induction generalizing addr1 addr2 val s
   case base =>
-    simp only [Nat.reduceSucc, Nat.succ_sub_succ_eq_sub, 
+    simp only [Nat.reduceSucc, Nat.succ_sub_succ_eq_sub,
                Nat.sub_self, BitVec.add_zero] at h4
     simp_all only [read_mem_bytes, BitVec.cast_eq]
     have h' : (BitVec.toNat (addr2 - addr1) + 1) * 8 - 1 - BitVec.toNat (addr2 - addr1) * 8 + 1 = 8 := by
@@ -832,7 +832,7 @@ private theorem read_mem_bytes_of_write_mem_bytes_subset_n2_lt
         · have l0 := @mem_subset_trans (addr2 + 1#64) (addr2 + n#64) addr2 (addr2 + n#64)
                    addr1 (addr1 + (n1 - 1)#64)
           simp only [h4] at l0
-          rw [first_addresses_add_one_is_subset_of_region_general 
+          rw [first_addresses_add_one_is_subset_of_region_general
               (by omega) (by omega) (by omega)] at l0
           · simp_all only [Nat.succ_sub_succ_eq_sub, Nat.sub_zero, forall_const]
           · simp only [mem_subset_refl]
@@ -939,7 +939,7 @@ private theorem read_mem_bytes_of_write_mem_bytes_subset_n2_eq_alt
       simp [my_pow_2_gt_zero]
     · unfold my_pow; decide
 
-@[simp]
+@[state_simp_rules]
 theorem read_mem_bytes_of_write_mem_bytes_subset
   (h0 : 0 < n1) (h1 : n1 <= 2^64) (h2 : 0 < n2) (h3 : n2 <= 2^64)
   (h4 : mem_subset addr2 (addr2 + (n2 - 1)#64) addr1 (addr1 + (n1 - 1)#64))
@@ -1009,7 +1009,7 @@ private theorem extract_byte_of_read_mem_bytes_succ (n : Nat) :
   rw [l0, Nat.mod_eq_of_lt y.isLt]
   done
 
-@[simp]
+@[state_simp_rules]
 theorem write_mem_bytes_irrelevant :
   write_mem_bytes n addr (read_mem_bytes n addr s) s = s := by
   induction n generalizing addr s
@@ -1028,6 +1028,17 @@ theorem write_mem_bytes_irrelevant :
     exact n_ih'
     done
 
+-- set_option pp.deepTerms false in
+-- set_option pp.deepTerms.threshold 1000 in
+-- theorem write_mem_bytes_irrelevant :
+--   write_mem_bytes n addr (read_mem_bytes n addr s) s = s := by
+--   induction n generalizing addr s
+--   case zero => simp only [write_mem_bytes]
+--   case succ =>
+--     rename_i n n_ih
+--     simp only [read_mem_bytes, write_mem_bytes]
+--     sorry
+
 ----------------------------------------------------------------------
 
 end MemoryProofs
diff --git a/Arm/State.lean b/Arm/State.lean
index 094f0ef1..e3582ce1 100644
--- a/Arm/State.lean
+++ b/Arm/State.lean
@@ -6,6 +6,7 @@ Author(s): Shilpi Goel
 import Lean.Data.Format
 import Arm.BitVec
 import Arm.Map
+import Arm.Attr
 
 ------------------------------------------------------------------------------
 ------------------------------------------------------------------------------
@@ -83,6 +84,13 @@ inductive StateError where
   | Other (e : String)         : StateError
 deriving DecidableEq, Repr
 
+-- Injective Lemmas for StateError
+attribute [state_simp_rules] StateError.NotFound.injEq
+attribute [state_simp_rules] StateError.Unimplemented.injEq
+attribute [state_simp_rules] StateError.Illegal.injEq
+attribute [state_simp_rules] StateError.Fault.injEq
+attribute [state_simp_rules] StateError.Other.injEq
+
 -- PFlag (Process State's Flags)
 inductive PFlag where
   | N : PFlag
@@ -209,6 +217,11 @@ inductive StateField where
   | ERR    : StateField
 deriving DecidableEq, Repr
 
+-- Injective Lemmas for StateField
+attribute [state_simp_rules] StateField.GPR.injEq
+attribute [state_simp_rules] StateField.SFP.injEq
+attribute [state_simp_rules] StateField.FLAG.injEq
+
 def state_value (fld : StateField) : Type :=
   open StateField in
   match fld with
@@ -238,7 +251,7 @@ def w (fld : StateField) (v : (state_value fld)) (s : ArmState) : ArmState :=
   | FLAG i => write_base_flag i v s
   | ERR    => write_base_error v s
 
-@[simp]
+@[state_simp_rules]
 theorem r_of_w_same (fld : StateField) (v : (state_value fld)) (s : ArmState) :
   r fld (w fld v s) = v := by
   unfold r w
@@ -249,7 +262,7 @@ theorem r_of_w_same (fld : StateField) (v : (state_value fld)) (s : ArmState) :
   unfold read_base_error write_base_error
   split <;> split <;> simp_all!
 
-@[simp]
+@[state_simp_rules]
 theorem r_of_w_different (fld1 fld2 : StateField) (v : (state_value fld2)) (s : ArmState)
   (h : fld1 ≠ fld2) :
   r fld1 (w fld2 v s) = r fld1 s := by
@@ -262,7 +275,7 @@ theorem r_of_w_different (fld1 fld2 : StateField) (v : (state_value fld2)) (s :
   simp_all!
   split <;> split <;> simp_all!
 
-@[simp]
+@[state_simp_rules]
 theorem w_of_w_shadow (fld : StateField) (v1 v2 : (state_value fld)) (s : ArmState) :
   w fld v2 (w fld v1 s) = w fld v2 s := by
   unfold w
@@ -273,7 +286,7 @@ theorem w_of_w_shadow (fld : StateField) (v1 v2 : (state_value fld)) (s : ArmSta
   unfold write_base_error
   split <;> simp
 
-@[simp]
+@[state_simp_rules]
 theorem w_irrelevant (fld : StateField) (v1 v2 : (state_value fld)) (s : ArmState) :
   w fld (r fld s) s = s := by
   unfold r w
@@ -284,7 +297,7 @@ theorem w_irrelevant (fld : StateField) (v1 v2 : (state_value fld)) (s : ArmStat
   unfold read_base_error write_base_error
   split <;> simp
 
-@[simp]
+@[state_simp_rules]
 theorem fetch_inst_of_w (addr : BitVec 64) (fld : StateField) (val : (state_value fld)) (s : ArmState) :
   fetch_inst addr (w fld val s) = fetch_inst addr s := by
   unfold fetch_inst w
@@ -309,7 +322,7 @@ theorem w_program (sf : StateField) (v : state_value sf) (s : ArmState):
 -- The following functions are defined in terms of r and w, but may be
 -- simpler to use.
 
-@[simp]
+@[state_simp_rules]
 def read_gpr (width : Nat) (idx : BitVec 5) (s : ArmState)
   : BitVec width :=
     let val := r (StateField.GPR idx) s
@@ -317,7 +330,7 @@ def read_gpr (width : Nat) (idx : BitVec 5) (s : ArmState)
 
 -- Use read_gpr_zr when register 31 is mapped to the zero register ZR,
 -- instead of the default (Stack pointer).
-@[simp]
+@[state_simp_rules]
 def read_gpr_zr (width : Nat) (idx : BitVec 5) (s : ArmState)
   : BitVec width :=
   if idx ≠ 31#5 then
@@ -328,7 +341,7 @@ def read_gpr_zr (width : Nat) (idx : BitVec 5) (s : ArmState)
 -- In practice, we only ever access the low 32 bits or the full 64
 -- bits of these registers in Arm. When we write 32 bits to these
 -- registers, the upper 32 bits are zeroed out.
-@[simp]
+@[state_simp_rules]
 def write_gpr (width : Nat) (idx : BitVec 5) (val : BitVec width) (s : ArmState)
   : ArmState :=
     let val := BitVec.zeroExtend 64 val
@@ -336,71 +349,71 @@ def write_gpr (width : Nat) (idx : BitVec 5) (val : BitVec width) (s : ArmState)
 
 -- Use write_gpr_zr when register 31 is mapped to the zero register
 -- ZR, instead of the default (Stack pointer).
-@[simp]
+@[state_simp_rules]
 def write_gpr_zr (n : Nat) (idx : BitVec 5) (val : BitVec n) (s : ArmState)
   : ArmState :=
   if idx ≠ 31#5 then
     write_gpr n idx val s
   else
     s
--- read_gpr and write_gpr are tagged with @[simp], which let us solve
+-- read_gpr and write_gpr are tagged with @[state_simp_rules], which let us solve
 -- the following using just simp, write_gpr, read_gpr, r_of_w_same
 -- (see simp?).
 example (n : Nat) (idx : BitVec 5) (val : BitVec n) (s : ArmState) :
   read_gpr n idx (write_gpr n idx val s) =
   BitVec.zeroExtend n (BitVec.zeroExtend 64 val) := by
-  simp
+  state_simp
 
-@[simp]
+@[state_simp_rules]
 def read_sfp (width : Nat) (idx : BitVec 5) (s : ArmState) : BitVec width :=
   let val := r (StateField.SFP idx) s
   BitVec.zeroExtend width val
 
 -- Write `val` to the `idx`-th SFP, zeroing the upper bits, if
 -- applicable.
-@[simp]
+@[state_simp_rules]
 def write_sfp (n : Nat) (idx : BitVec 5) (val : BitVec n) (s : ArmState) : ArmState :=
    let val := BitVec.zeroExtend 128 val
    w (StateField.SFP idx) val s
 
-@[simp]
+@[state_simp_rules]
 def read_pc (s : ArmState) : BitVec 64 :=
   r StateField.PC s
 
-@[simp]
+@[state_simp_rules]
 def write_pc (v : BitVec 64) (s : ArmState) : ArmState :=
   w StateField.PC v s
 
-@[simp]
+@[state_simp_rules]
 def read_flag (flag : PFlag) (s : ArmState) : BitVec 1 :=
   r (StateField.FLAG flag) s
 
-@[simp]
+@[state_simp_rules]
 def write_flag (flag : PFlag) (val : BitVec 1) (s : ArmState) : ArmState :=
   w (StateField.FLAG flag) val s
 
-@[simp]
+@[state_simp_rules]
 def read_pstate (s : ArmState) : PState :=
   fun p => read_flag p s
 
-@[simp]
+@[state_simp_rules]
 def write_pstate (pstate : PState) (s : ArmState) : ArmState :=
    { s with pstate := pstate }
 
 -- (FIXME) Define in terms of write_flag so that we see checkpoints in
 -- terms of the w function.
-@[simp]
+@[state_simp_rules]
 def make_pstate (n z c v : BitVec 1) : PState :=
   fun (p : PFlag) =>
     open PFlag in
     match p with
       | N => n | Z => z | C => c | V => v
 
-@[simp]
+@[state_simp_rules]
 def read_err (s : ArmState) : StateError :=
   r StateField.ERR s
 
-@[simp]
+@[state_simp_rules]
 def write_err (v : StateError) (s : ArmState) : ArmState :=
   w StateField.ERR v s
 
@@ -450,19 +463,21 @@ end Load_program_and_fetch_inst
 
 ----------------------------------------------------------------------
 
+-- Adding some basic simp lemmas to `state_simp_rules`:
+attribute [state_simp_rules] ne_eq
+attribute [state_simp_rules] not_false_eq_true
+
 example :
   read_flag flag (write_flag flag val s) = val := by
-  simp
+  state_simp
 
 example (h : flag1 ≠ flag2) :
   read_flag flag1 (write_flag flag2 val s) = read_flag flag1 s := by
-  simp [*] at *
+  state_simp_all
 
 example :
   read_gpr width idx (write_flag flag2 val s) = read_gpr width idx s := by
-  simp
-
--- #help tactic simp
+  state_simp
 
 end State
 
diff --git a/Proofs/Sha512_block_armv8.lean b/Proofs/Sha512_block_armv8.lean
index cb34c966..9bc848f2 100644
--- a/Proofs/Sha512_block_armv8.lean
+++ b/Proofs/Sha512_block_armv8.lean
@@ -57,7 +57,16 @@ theorem sha512_program_test_1_sym (s0 s_final : ArmState)
   (h_s0_ok : read_err s0 = StateError.None)
   (h_run : s_final = run 4 s0) :
   read_err s_final = StateError.None := by
-  iterate 4 (sym1 [h_s0_program])
+  -- iterate 4 (sym1 [h_s0_program])
+  -- sym1 [h_s0_program]
+  (try simp_all (config := {decide := true, ground := true}) only [state_simp_rules]);
+  unfold run;
+  simp_all only [stepi, state_simp_rules];
+  (try rw [fetch_inst_from_program h_s0_program]);
+  -- (try simp (config := {ground := true}) only);
+  -- simp only [exec_inst, state_simp_rules];
+  -- (try simp_all (config := {decide := true, ground := true}) only [state_simp_rules])
+  sorry
   -- unfold read_pc at h_s0_pc
   -- sym_n 4 0x126538 sha512_program_test_1
   -- rw [h_run,h_s4_ok]
@@ -86,7 +95,8 @@ theorem sha512_program_test_2_sym (s0 s_final : ArmState)
   read_err s_final = StateError.None := by
   -- TODO: use sym_n. Using it causes an error at simp due to the
   -- large formula size.
-  iterate 4 (sym1 [h_s0_program])
+  -- iterate 4 (sym1 [h_s0_program])
+  sorry
 
 ----------------------------------------------------------------------
 
diff --git a/Tactics/Sym.lean b/Tactics/Sym.lean
index 35699103..259c47d5 100644
--- a/Tactics/Sym.lean
+++ b/Tactics/Sym.lean
@@ -16,18 +16,18 @@ syntax "sym1" "[" term "]" : tactic
 macro_rules
   | `(tactic| sym1 [$h_program:term]) =>
     `(tactic|
-      (try simp_all (config := {decide := true, ground := true}));
+      (try simp_all (config := {decide := true, ground := true}) only [state_simp_rules]);
       unfold run;
-      simp_all [stepi];
+      simp_all only [stepi, state_simp_rules];
       (try rw [fetch_inst_from_program $h_program]);
       (try simp (config := {decide := true, ground := true}) only);
       -- After exec_inst is opened up, the exec functions of the
       -- instructions which are tagged with simp will also open up
       -- here.
-      simp [exec_inst];
+      simp only [exec_inst, state_simp_rules];
       -- (try simp_all (config := {decide := true, ground := true}) only);
       -- (try simp only [ne_eq, r_of_w_different, r_of_w_same, w_of_w_shadow, w_irrelevant])
-      (try simp_all (config := {decide := true, ground := true})))
+      (try simp_all (config := {decide := true, ground := true}) only [state_simp_rules]))
 
 theorem run_onestep (s s': ArmState) (n : Nat) (h_nonneg : 0 < n):
   (s' = run n s) ↔ ∃ s'', s'' = stepi s ∧ s' = run (n-1) s'' := by
diff --git a/lean-toolchain b/lean-toolchain
index d71105d8..f16b484e 100644
--- a/lean-toolchain
+++ b/lean-toolchain
@@ -1 +1 @@
-leanprover/lean4:nightly-2024-02-24
+leanprover/lean4:nightly-2024-03-01