chore: Split out the finishing tactic aspect of simp_mem into mem_omega [2/?]

Arm/Memory/Common.lean

@@ -227,6 +227,9 @@ inductive Hypothesis
 | pairwiseSeparate (proof : MemPairwiseSeparateProof e)
 | read_eq (proof : ReadBytesEqProof)
 
+def Hypothesis.isPairwiseSeparate : Hypothesis → Bool
+| .pairwiseSeparate _ => true
+| _ => false
 
 def Hypothesis.proof : Hypothesis → Expr
 | .pairwiseSeparate proof  => proof.h
@@ -245,17 +248,35 @@ instance : ToMessageData Hypothesis where
 
 /-- create a trace node in trace class (i.e. `set_option traceClass true`),
 with header `header`, whose default collapsed state is `collapsed`. -/
-def TacticM.withTraceNode' (header : MessageData) (k : TacticM α)
+def TacticM.withTraceNode'
+    [Monad m]
+    [MonadTrace m]
+    [MonadLiftT IO m]
+    [MonadRef m]
+    [AddMessageContext m]
+    [MonadOptions m]
+    [always : MonadAlwaysExcept ε m]
+    [MonadLiftT BaseIO m]
+    (header : MessageData) (k : m α)
     (collapsed : Bool := false)
-    (traceClass : Name := `simp_mem.info) : TacticM α :=
+    (traceClass : Name := `simp_mem.info) : m α :=
   Lean.withTraceNode traceClass (fun _ => return header) k (collapsed := collapsed)
 
 /-- Create a trace note that folds `header` with `(NOTE: can be large)`,
 and prints `msg` under such a trace node. Collapsed by default.
 -/
-def TacticM.traceLargeMsg (header : MessageData) (msg : MessageData) : TacticM Unit := do
+def TacticM.traceLargeMsg
+    [Monad m]
+    [MonadTrace m]
+    [MonadLiftT IO m]
+    [MonadRef m]
+    [AddMessageContext m]
+    [MonadOptions m]
+    [always : MonadAlwaysExcept ε m]
+    [MonadLiftT BaseIO m]
+    (header : MessageData)
+    (msg : MessageData) : m Unit := do
     withTraceNode' m!"{header} (NOTE: can be large)" (collapsed := true) do
-      -- | TODO: change trace class?
       trace[simp_mem.info] msg
 
 
@@ -270,15 +291,14 @@ def omega (bvToNatSimpCtx : Simp.Context) (bvToNatSimprocs : Array Simp.Simprocs
       trace[simp_mem.info] "{goal}"
     -- @bollu: TODO: understand what precisely we are recovering from.
     withoutRecover do
-    evalTactic (← `(tactic| bv_omega_bench))
+      evalTactic (← `(tactic| bv_omega_bench))
 
 /-
 Introduce a new definition into the local context, simplify it using `simp`,
 and return the FVarId of the new definition in the goal.
 -/
 def simpAndIntroDef (name : String) (hdefVal : Expr) : TacticM FVarId  := do
   withMainContext do
-
     let name ← mkFreshUserName <| .mkSimple name
     let goal ← getMainGoal
     let hdefTy ← inferType hdefVal
@@ -301,7 +321,7 @@ def simpAndIntroDef (name : String) (hdefVal : Expr) : TacticM FVarId  := do
     let hdefVal ← simpResult.mkCast hdefVal
     let hdefTy ← inferType hdefVal
 
-    let goal ← goal.define name hdefTy hdefVal
+    let goal ← goal.assert name hdefTy hdefVal
     let (fvar, goal) ← goal.intro1P
     replaceMainGoal [goal]
     return fvar
@@ -408,7 +428,7 @@ partial def MemPairwiseSeparateProof.ofExpr? (e : Expr) : Option MemPairwiseSepa
 /-- Match an expression `h` to see if it's a useful hypothesis. -/
 def hypothesisOfExpr (h : Expr) (hyps : Array Hypothesis) : MetaM (Array Hypothesis) := do
   let ht ← inferType h
-  trace[simp_mem.info] "{processingEmoji} Processing '{h}' : '{toString ht}'"
+  trace[simp_mem.info] "{processingEmoji} Processing '{h}' : '{ht}'"
   if let .some sep := MemSeparateProp.ofExpr? ht then
     let proof : MemSeparateProof sep := ⟨h⟩
     let hyps := hyps.push (.separate proof)
@@ -749,4 +769,32 @@ def proveWithOmega?  {α : Type} [ToMessageData α] [OmegaReducible α] (e : α)
     return none
   end ReductionToOmega
 
+/--
+simplify the goal state, closing legality, subset, and separation goals,
+and simplifying all other expressions. return `true` if goal has been closed, and `false` otherwise.
+-/
+partial def closeMemSideCondition (g : MVarId)
+    (bvToNatSimpCtx : Simp.Context) (bvToNatSimprocs : Array Simp.Simprocs)
+    (hyps : Array Memory.Hypothesis) : TacticM Bool := do
+  g.withContext do
+    trace[simp_mem.info] "{processingEmoji} Matching on ⊢ {← g.getType}"
+    let gt ← g.getType
+    if let .some e := MemLegalProp.ofExpr? gt then
+      TacticM.withTraceNode' m!"Matched on ⊢ {e}. Proving..." do
+        if let .some proof ← proveWithOmega? e bvToNatSimpCtx bvToNatSimprocs hyps then
+          g.assign proof.h
+    if let .some e := MemSubsetProp.ofExpr? gt then
+      TacticM.withTraceNode' m!"Matched on ⊢ {e}. Proving..." do
+        if let .some proof ← proveWithOmega? e bvToNatSimpCtx bvToNatSimprocs hyps then
+          g.assign proof.h
+    if let .some e := MemSeparateProp.ofExpr? gt then
+      TacticM.withTraceNode' m!"Matched on ⊢ {e}. Proving..." do
+        if let .some proof ← proveWithOmega? e bvToNatSimpCtx bvToNatSimprocs hyps then
+          g.assign proof.h
+  return ← g.isAssigned
+
+
+
+
+
 end Tactic.Memory

Arm/Memory/MemOmega.lean

@@ -0,0 +1,137 @@
+/-
+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): Siddharth Bhat
+
+In this file, we define proof automation for separation conditions of memory.
+
+References:
+- https://github.com/leanprover/lean4/blob/240ebff549a2cf557f9abe9568f5de885f13e50d/src/Lean/Elab/Tactic/Omega/OmegaM.lean
+- https://github.com/leanprover/lean4/blob/240ebff549a2cf557f9abe9568f5de885f13e50d/src/Lean/Elab/Tactic/Omega/Frontend.lean
+-/
+import Arm
+import Arm.Memory.MemoryProofs
+import Arm.BitVec
+import Arm.Memory.Attr
+import Arm.Memory.AddressNormalization
+import Lean
+import Lean.Meta.Tactic.Rewrite
+import Lean.Meta.Tactic.Rewrites
+import Lean.Elab.Tactic.Conv
+import Lean.Elab.Tactic.Conv.Basic
+import Tactics.Simp
+import Tactics.BvOmegaBench
+import Arm.Memory.Common
+
+open Lean Meta Elab Tactic Memory
+
+namespace MemOmega
+
+structure Config where
+  /--
+    If true, then MemOmega will explode uses of pairwiseSeparate [mem₁, ... memₙ]
+    into O(n^2) separation conditions.
+  -/
+  explodePairwiseSeparate : Bool := false
+
+/-- Edit the config for mem_omega! -/
+def Config.mkBang (c : Config) : Config :=
+  { c with explodePairwiseSeparate := true }
+
+/-- Context for the `SimpMemM` monad, containing the user configurable options. -/
+structure Context where
+  /-- User configurable options for `simp_mem`. -/
+  cfg : Config
+  /-- Cache of `bv_toNat` simp context. -/
+  bvToNatSimpCtx : Simp.Context
+  /-- Cache of `bv_toNat` simprocs. -/
+  bvToNatSimprocs : Array Simp.Simprocs
+
+
+namespace Context
+
+def init (cfg : Config) : MetaM Context := do
+  let (bvToNatSimpCtx, bvToNatSimprocs) ←
+    LNSymSimpContext
+      (config := {failIfUnchanged := false})
+      -- Also use `mem_{legal', subset', separate'}.iff_omega to unfold definitions that
+      -- occur inside compound expressions, such as (mem_subset' .. ∨ mem_subset' ..)
+      -- (thms := #[``mem_legal'.iff_omega, ``mem_subset'.iff_omega, ``mem_separate'.iff_omega])
+      (simp_attrs := #[`bv_toNat])
+      (useDefaultSimprocs := false)
+  return {cfg, bvToNatSimpCtx, bvToNatSimprocs}
+end Context
+
+abbrev MemOmegaM := (ReaderT Context TacticM)
+
+namespace MemOmegaM
+
+  def run (ctx : Context) (x : MemOmegaM α) : TacticM α := ReaderT.run x ctx
+
+end MemOmegaM
+
+def memOmegaTac : MemOmegaM Unit := do
+  let g ← getMainGoal
+  g.withContext do
+    /- We need to explode all pairwise separate hyps -/
+    let rawHyps ← getLocalHyps
+    let mut hyps := #[]
+    -- extract out structed values for all hyps.
+    for h in rawHyps do
+      hyps ← hypothesisOfExpr h hyps
+
+    -- only enable pairwise constraints if it is enabled.
+    let isPairwiseEnabled := (← readThe Context).cfg.explodePairwiseSeparate
+    hyps := hyps.filter (!·.isPairwiseSeparate || isPairwiseEnabled)
+
+    -- used specialized procedure that doesn't unfold everything for the easy case.
+    if ← closeMemSideCondition (← getMainGoal) (← readThe Context).bvToNatSimpCtx (← readThe Context).bvToNatSimprocs hyps then
+      return ()
+    else
+      -- in the bad case, just rip through everything.
+      -- let _ ← Hypothesis.addOmegaFactsOfHyps (hyps.toList.filter (fun h => h.isPairwiseSeparate)) #[]
+      let _ ← Hypothesis.addOmegaFactsOfHyps hyps.toList #[]
+
+      TacticM.withTraceNode' m!"Reducion to omega" do
+        try
+          TacticM.traceLargeMsg m!"goal (Note: can be large)"  m!"{← getMainGoal}"
+          omega (← readThe Context).bvToNatSimpCtx (← readThe Context).bvToNatSimprocs
+          trace[simp_mem.info] "{checkEmoji} `omega` succeeded."
+        catch e =>
+          trace[simp_mem.info]  "{crossEmoji} `omega` failed with error:\n{e.toMessageData}"
+
+
+/--
+Allow elaboration of `MemOmegaConfig` arguments to tactics.
+-/
+declare_config_elab elabMemOmegaConfig MemOmega.Config
+
+/--
+The `mem_omega` tactic is a finishing tactic which is used to dispatch memory side conditions.
+Broadly, the algorithm works as follows:
+- It scans the set of hypotheses for `mem_separate`, `mem_subset`, and `mem_legal` hypotheses, and turns them into `omega` based information.
+- It calls `omega` as a finishing tactic to close the current goal state.
+- Cruicially, it **does not unfold** `pairwiseSeparate` constraints. We expect the user to do so. If they want `pairwiseSeparate` unfolded, then please use `mem_omega!`.
+-/
+syntax (name := mem_omega) "mem_omega" (Lean.Parser.Tactic.config)? : tactic
+
+/--
+The `mem_omega!` tactic is a finishing tactic, that is a more aggressive variant of `mem_omega`.
+-/
+syntax (name := mem_omega_bang) "mem_omega!" (Lean.Parser.Tactic.config)? : tactic
+
+@[tactic mem_omega]
+def evalMemOmega : Tactic := fun
+  | `(tactic| mem_omega $[$cfg]?) => do
+    let cfg ← elabMemOmegaConfig (mkOptionalNode cfg)
+    memOmegaTac.run (← Context.init cfg)
+  | _ => throwUnsupportedSyntax
+
+@[tactic mem_omega_bang]
+def evalMemOmegaBang : Tactic := fun
+  | `(tactic| mem_omega! $[$cfg]?) => do
+    let cfg ← elabMemOmegaConfig (mkOptionalNode cfg)
+    memOmegaTac.run (← Context.init cfg.mkBang)
+  | _ => throwUnsupportedSyntax
+
+end MemOmega

Arm/Memory/Separate.lean

@@ -230,12 +230,12 @@ theorem mem_legal'.omega_def (h : mem_legal' a n) : a.toNat + n ≤ 2^64 := h
 
 /-- The linear constraint is equivalent to `mem_legal'`. -/
 theorem mem_legal'.iff_omega (a : BitVec 64) (n : Nat) :
-    (a.toNat + n ≤ 2^64) ↔ mem_legal' a n := by
+    mem_legal' a n ↔ (a.toNat + n ≤ 2^64) := by
   constructor
-  · intros h
-    apply mem_legal'.of_omega h
   · intros h
     apply h.omega_def
+  · intros h
+    apply mem_legal'.of_omega h
 
 instance : Decidable (mem_legal' a n) :=
   if h : a.toNat + n ≤ 2^64 then
@@ -341,14 +341,15 @@ theorem mem_separate'.of_omega
 
 /-- The linear constraint is equivalent to `mem_separate'`. -/
 theorem mem_separate'.iff_omega (a : BitVec 64) (an : Nat) (b : BitVec 64) (bn : Nat) :
+    mem_separate' a an b bn ↔
     (a.toNat + an ≤ 2^64 ∧
     b.toNat + bn ≤ 2^64 ∧
-    (a.toNat + an ≤ b.toNat ∨ a.toNat ≥ b.toNat + bn)) ↔ mem_separate' a an b bn := by
+    (a.toNat + an ≤ b.toNat ∨ a.toNat ≥ b.toNat + bn)) := by
   constructor
-  · intros h
-    apply mem_separate'.of_omega h
   · intros h
     apply h.omega_def
+  · intros h
+    apply mem_separate'.of_omega h
 
 instance : Decidable (mem_separate' a an b bn) :=
   if h : (a.toNat + an ≤ 2^64 ∧ b.toNat + bn ≤ 2^64 ∧ (a.toNat + an ≤ b.toNat ∨ a.toNat ≥ b.toNat + bn)) then
@@ -430,15 +431,16 @@ constructor
 
 /-- The linear constraint is equivalent to `mem_subset'`. -/
 theorem mem_subset'.iff_omega (a : BitVec 64) (an : Nat) (b : BitVec 64) (bn : Nat) :
+    mem_subset' a an b bn  ↔
     (a.toNat + an ≤ 2^64 ∧
     b.toNat + bn ≤ 2^64 ∧
     b.toNat ≤ a.toNat ∧
-    a.toNat + an ≤ b.toNat + bn) ↔ mem_subset' a an b bn := by
+    a.toNat + an ≤ b.toNat + bn) := by
   constructor
-  · intros h
-    apply mem_subset'.of_omega h
   · intros h
     apply h.omega_def
+  · intros h
+    apply mem_subset'.of_omega h
 
 instance : Decidable (mem_subset' a an b bn) :=
   if h : (a.toNat + an ≤ 2^64 ∧ b.toNat + bn ≤ 2^64 ∧ b.toNat ≤ a.toNat ∧ a.toNat + an ≤ b.toNat + bn) then
@@ -687,4 +689,21 @@ def Memory.Region.separate'_of_pairwiseSeprate_of_mem_of_mem
         · simp only [List.get?_eq_getElem?, ha]
         · simp only [List.get?_eq_getElem?, hb]
 
+
+/--
+Convenient API to extract out a separate proof from a Memory.Region.pairwiseSeparate.
+Proof obligations are given by `decide` to allow the API to be used as follows:
+
+let h := mem.pairwiseSeparate [(a, na), (b, nb), (c, nc)]
+let h01 := h.get 0 1
+-/
+def Memory.Region.pairwiseSeparate.get
+    (h : Memory.Region.pairwiseSeparate mems)
+    (i j : Nat)
+    (hi : i < mems.length := by simp)
+    (hj : j < mems.length := by simp)
+    (hij : i ≠ j := by omega) :
+    mem_separate' mems[i].fst mems[i].snd mems[j].fst mems[j].snd := by
+  apply Memory.Region.separate'_of_pairwiseSeprate_of_mem_of_mem h i j hij <;> simp
+
 end NewDefinitions