refactor: cleanup popcount proof (#168)

### Description:

We've cleaned up the popcount proof by:
* extending `sym_aggregate` to also search for memory effect hypotheses
to simplify, and
* changing the semantics of `add_sub_immediate` to avoid
pattern-matching. It turns out that pattern-matching on the result of an
`if` (e.g., `let (a, b) := if x = 0 then (foo, bar) else (bar, foo)`)
get's reduced to a `Decidable.rec` by `simp only`, which then gets
stuck. By rephrasing it in terms of projections, we avoid this, and get
clean simplified semantics.
* In the same vein, we also add existing simplification rules for
`AddWithCarry` to the `bitvec_rules` simpset

### Testing:

`make all` succeeds.

### 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: Siddharth Bhat <[email protected]>
Co-authored-by: Shilpi Goel <[email protected]>

Proofs/Popcount32.lean

@@ -84,43 +84,28 @@ theorem popcount32_sym_meets_spec (s0 s_final : ArmState)
   (∀ (n : Nat) (addr : BitVec 64),
     mem_separate' addr n (r (.GPR 31) s0 - 16#64) 16 →
     s_final[addr, n] = s0[addr, n]) := by
-  -- Prelude
-  simp_all only [state_simp_rules, -h_run]
-  -- Symbolic Simulation
-  sym_n 27
-  -- Final Steps
-  -- Split all the Ands in the conclusion.
-  repeat' apply And.intro
-  · simp only [popcount32_spec,
-               fst_AddWithCarry_eq_add,
-               fst_AddWithCarry_eq_sub_neg]
-    simp only [popcount32_spec_rec]
+  simp_all only [state_simp_rules, -h_run] -- prelude
+  sym_n 27 -- Symbolic simulation
+  repeat' apply And.intro -- split conjunction.
+  · simp only [popcount32_spec, popcount32_spec_rec]
     bv_decide
   · sym_aggregate
-  · -- (TODO @alex) Let's do away with
-    -- ∀ (n : Nat) (addr : BitVec 64), read_mem_bytes n addr s₁ = read_mem_bytes n addr s₂
-    -- in favor of
-    -- s₁.mem = s₂.mem
-    -- as Sid said.
-    simp only [←Memory.mem_eq_iff_read_mem_bytes_eq] at *
+  · intro n addr h_separate
     simp only [memory_rules] at *
-    intro n addr h_separate
-
-    -- (TODO @alex/@bollu) Can we hope to make this shorter after the marriage
-    -- of `sym_n` and `simp_mem`?
-    simp_mem
-    sym_aggregate
-    simp only [*, bitvec_rules]
-    sym_aggregate
-    simp_mem
-    rfl
-  · apply Aligned_BitVecSub_64_4
+    repeat (simp_mem (config := { useOmegaToClose := false }); sym_aggregate)
+  · apply Aligned_BitVecSub_64_4 -- TODO(@bollu): automation
     · assumption
     · decide
   · apply Aligned_BitVecAdd_64_4
     · assumption
     · decide
 
+/--
+info: 'popcount32_sym_meets_spec' depends on axioms:
+[propext, Classical.choice, Lean.ofReduceBool, Quot.sound]
+-/
+#guard_msgs in #print axioms popcount32_sym_meets_spec
+
 -------------------------------------------------------------------------------
 
 /-! ## Tests for step theorem generation -/

Proofs/SHA512/SHA512Loop.lean

@@ -74,7 +74,7 @@ set_option debug.skipKernelTC true in
 -- set_option profiler true in
 -- set_option profiler.threshold 1 in
 set_option maxHeartbeats 0 in
--- set_option maxRecDepth 0 in
+set_option maxRecDepth 8000 in
 theorem sha512_block_armv8_loop_1block (si sf : ArmState)
   (h_N : N = 1#64)
   (h_si_prelude : sha512_prelude 0x126500#64 N SP CtxBase InputBase si)
@@ -104,7 +104,7 @@ theorem sha512_block_armv8_loop_1block (si sf : ArmState)
          state_simp_rules,
          bitvec_rules, minimal_theory]
   sym_aggregate
-  exact ⟨h_si_ktbl, h_si_separate⟩
+  assumption
   done
 
 end SHA512

Proofs/SHA512/SHA512Prelude.lean

@@ -161,6 +161,12 @@ private theorem add_eq_sub_16 (x : BitVec 64) :
 
 -- (TODO) Modifying this theorem is an exercise in patience because of
 -- user-interactivity delays. Let's report this.
+/-
+Adding the ability to aggregate memory effects has caused a need to increase `maxRecDepth`
+in this proof. This will hopefully go down, once we optimize `sym_aggregate`.
+-/
+set_option maxRecDepth 8000 in
+set_option linter.unusedVariables false in
 theorem sha512_block_armv8_prelude (s0 sf : ArmState)
   -- We fix the number of blocks to hash to 1.
   (h_N : N = 1#64)
@@ -242,7 +248,7 @@ theorem sha512_block_armv8_prelude (s0 sf : ArmState)
   · constructor
     · -- (TODO @bollu) Think about whether `simp_mem` should be powerful enough to solve this goal.
       -- Also, `mem_omega` name suggestion from Alex for the already souped up `simp_mem`.
-      simp_mem
+      simp_mem (config := { useOmegaToClose := false } )
       simp only [h_s0_ctx_base, Nat.sub_self, minimal_theory, bitvec_rules]
     · constructor
       · -- (FIXME @bollu) simp_mem doesn't make progress here. :-(

Tactics/Aggregate.lean

@@ -1,7 +1,7 @@
 /-
 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
+Author(s): Alex Keizer, Siddharth Bhat
 -/
 import Lean
 import Tactics.Common
@@ -68,12 +68,18 @@ elab "sym_aggregate" simpConfig?:(config)? loc?:(location)? : tactic => withMain
     let lctx ← getLCtx
     -- We keep `expectedRead`/`expectedAlign` as monadic values,
     -- so that we get new metavariables for each localdecl we check
-    let expectedRead := do
+    let expectedRead : MetaM Expr := do
       let fld ← mkFreshExprMVar (mkConst ``StateField)
       let state ← mkFreshExprMVar mkArmState
       let rhs ← mkFreshExprMVar none
       mkEq (mkApp2 (mkConst ``r) fld state) rhs
-    let expectedAlign := do
+    let expectedReadMem : MetaM Expr := do
+      let n ← mkFreshExprMVar (mkConst ``Nat)
+      let addr ← mkFreshExprMVar (mkApp (mkConst ``BitVec) (toExpr 64))
+      let mem ← mkFreshExprMVar (mkConst ``Memory)
+      let rhs ← mkFreshExprMVar none
+      mkEq (mkApp3 (mkConst ``Memory.read_bytes) n addr mem) rhs
+    let expectedAlign : MetaM Expr := do
       let state ← mkFreshExprMVar mkArmState
       return mkApp (mkConst ``CheckSPAlignment) state
 
@@ -84,12 +90,16 @@ elab "sym_aggregate" simpConfig?:(config)? loc?:(location)? : tactic => withMain
             trace[Tactic.sym] "checking {decl.toExpr} with type {type}"
             let expectedRead ← expectedRead
             let expectedAlign ← expectedAlign
+            let expectedReadMem ← expectedReadMem
             if ← isDefEq type expectedRead then
               trace[Tactic.sym] "{Lean.checkEmoji} match for {expectedRead}"
               return axHyps.push decl
             else if ← isDefEq type expectedAlign then
               trace[Tactic.sym] "{Lean.checkEmoji} match for {expectedAlign}"
               return axHyps.push decl
+            else if ← isDefEq type expectedReadMem then
+              trace[Tactic.sym] "{Lean.checkEmoji} match for {expectedReadMem}"
+              return axHyps.push decl
             else
               trace[Tactic.sym] "{Lean.crossEmoji} no match"
               return axHyps