chore: Give mem_omega only [h1, ... hn] control

In the next step, this will allow simp_mem to filter hypotheses it sends over to mem_omega.

This lays out the full scope of control we'll implement for `simp_mem`.

Arm/Memory/MemOmega.lean

@@ -18,6 +18,7 @@ import Lean
 import Lean.Meta.Tactic.Rewrite
 import Lean.Meta.Tactic.Rewrites
 import Lean.Elab.Tactic.Conv
+import Lean.Elab.Tactic.Simp
 import Lean.Elab.Tactic.Conv.Basic
 import Tactics.Simp
 import Tactics.BvOmegaBench
@@ -27,6 +28,22 @@ open Lean Meta Elab Tactic Memory
 
 namespace MemOmega
 
+/--
+A user given hypothesis for mem_omega, which we process as either a hypothesis (FVarId),
+or a term that is added into the user's context.
+-/
+inductive UserHyp
+| hyp : FVarId → UserHyp
+| expr : Expr → UserHyp
+
+
+namespace UserHyp
+  def toExpr : UserHyp → Expr
+  | .hyp fvarId => Expr.fvar fvarId
+  | .expr e => e
+end UserHyp
+
+
 structure Config where
   /--
     If true, then MemOmega will explode uses of pairwiseSeparate [mem₁, ... memₙ]
@@ -42,6 +59,11 @@ def Config.mkBang (c : Config) : Config :=
 structure Context where
   /-- User configurable options for `simp_mem`. -/
   cfg : Config
+  /--
+  If we are using `mem_omega only [...]`, then we will have `some` plus the hyps.
+  If we are using `mem_omega`, then we will get `none`.
+  -/
+  userHyps? : Option (Array UserHyp)
   /-- Cache of `bv_toNat` simp context. -/
   bvToNatSimpCtx : Simp.Context
   /-- Cache of `bv_toNat` simprocs. -/
@@ -50,7 +72,7 @@ structure Context where
 
 namespace Context
 
-def init (cfg : Config) : MetaM Context := do
+def init (cfg : Config) (userHyps? : Option (Array UserHyp)) : MetaM Context := do
   let (bvToNatSimpCtx, bvToNatSimprocs) ←
     LNSymSimpContext
       (config := {failIfUnchanged := false})
@@ -59,7 +81,7 @@ def init (cfg : Config) : MetaM Context := do
       -- (thms := #[``mem_legal'.iff_omega, ``mem_subset'.iff_omega, ``mem_separate'.iff_omega])
       (simp_attrs := #[`bv_toNat])
       (useDefaultSimprocs := false)
-  return {cfg, bvToNatSimpCtx, bvToNatSimprocs}
+  return {cfg, bvToNatSimpCtx, bvToNatSimprocs, userHyps? }
 end Context
 
 abbrev MemOmegaM := (ReaderT Context MetaM)
@@ -68,7 +90,30 @@ namespace MemOmegaM
   def run (ctx : Context) (x : MemOmegaM α) : MetaM α := ReaderT.run x ctx
 end MemOmegaM
 
-def memOmega (g : MVarId) : MemOmegaM Unit := do
+/--
+Given the user hypotheses, build a more focusedd MVarId that contains only those hypotheses.
+This makes `omega` focus only on those hypotheses, since omega by default crawls the entire goal state.
+
+This is arguably a workaround to having to plumb the hypotheses through the full layers of code, but it works,
+and should be a cheap solution.
+-/
+def mkGoalWithOnlyUserHyps (g : MVarId) (userHyps? : Option (Array UserHyp)) : MetaM <| MVarId :=
+  match userHyps? with
+  | none => pure g
+  | some userHyps => do
+    g.withContext do 
+      let mut keepHyps : Std.HashSet FVarId := ∅
+      for h in userHyps do
+        if let .hyp fvar := h then 
+          keepHyps := keepHyps.insert fvar
+      let hyps ← g.getNondepPropHyps
+      let mut g := g
+      for h in hyps do
+        if !keepHyps.contains h then
+          g ← g.clear h
+      return g
+
+def memOmega' (g : MVarId) : MemOmegaM Unit := do
   g.withContext do
     /- We need to explode all pairwise separate hyps -/
     let rawHyps ← getLocalHyps
@@ -96,36 +141,107 @@ def memOmega (g : MVarId) : MemOmegaM Unit := do
         catch e =>
           trace[simp_mem.info]  "{crossEmoji} `omega` failed with error:\n{e.toMessageData}"
 
+def memOmega (g : MVarId) : MemOmegaM Unit := do
+    let g ← mkGoalWithOnlyUserHyps g (← readThe Context).userHyps? 
+    g.withContext do
+      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 g (← readThe Context).bvToNatSimpCtx (← readThe Context).bvToNatSimprocs hyps then
+        return ()
+      else
+        -- in the bad case, just rip through everything.
+        let (_, g) ← Hypothesis.addOmegaFactsOfHyps g hyps.toList #[]
+
+        TacticM.withTraceNode' m!"Reducion to omega" do
+          try
+            TacticM.traceLargeMsg m!"goal (Note: can be large)"  m!"{g}"
+            omega g (← 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}"
+
+-- syntax memOmegaRule := term
+
 /--
 Allow elaboration of `MemOmegaConfig` arguments to tactics.
 -/
 declare_config_elab elabMemOmegaConfig MemOmega.Config
 
+syntax memOmegaWith := &"with" "[" withoutPosition(term,*,?) "]"
+
 /--
-Implement the `mem_omega` tactic, which unfolds information about memory
-in terms of
+Implement the `mem_omega` tactic, which unfolds information about memory and then closes the goal state using `omega`.
 -/
-syntax (name := mem_omega) "mem_omega" (Lean.Parser.Tactic.config)? : tactic
+syntax (name := mem_omega) "mem_omega" (Lean.Parser.Tactic.config)? (memOmegaWith)?  : tactic
 
 /--
 Implement the `mem_omega` tactic frontend.
 -/
-syntax (name := mem_omega_bang) "mem_omega!" (Lean.Parser.Tactic.config)? : tactic
+syntax (name := mem_omega_bang) "mem_omega!" (memOmegaWith)?  : tactic
+
+-- /-- Since we have
+-- syntax memOmegaRule := term
+-- 
+-- it is safe to coerce a UserHyp into a `term`.
+-- -/
+-- def UserHyp.toTermSyntax (t : TSyntax ``MemOmega.memOmegaRule) : TSyntax `term :=
+--   TSyntax.mk t.raw
+
+
+
+/--
+build a `UserHyp` from the raw syntax.
+This supports using fars, using CDot notation to partially apply theorems, and to use terms.
 
+Adapted from Lean.Elab.Tactic.Rw.WithRWRulesSeq, Lean.Elab.Tactic.Simp.resolveSimpIdTheorem, Lean.Elab.Tactic.Simp.addSimpTheorem
+-/
+def UserHyp.ofSyntax (t : TSyntax `term) : TacticM UserHyp := do
+  -- See if we can interpret `id` as a hypothesis first.
+  if let .some fvarId ← optional <| getFVarId t then
+    return .hyp fvarId
+  else if let some e ← Term.elabCDotFunctionAlias? t then
+    return .expr e
+  else 
+    let e ← Term.elabTerm t none
+    Term.synthesizeSyntheticMVars (postpone := .no) (ignoreStuckTC := true)
+    let e ← instantiateMVars e
+    let e := e.eta
+    if e.hasMVar then
+      throwErrorAt t "found metavariables when elaborating rule, giving up."
+    return .expr e
+
+
+-- Adapted from WithRWRulesSeq.
+def elabMemOmegaWith : TSyntax ``MemOmega.memOmegaWith → TacticM (Array UserHyp) 
+| `(memOmegaWith| with [ $[ $rules],* ]) =>  do
+    rules.mapM UserHyp.ofSyntax
+| _ => throwUnsupportedSyntax
+
+open Lean.Parser.Tactic in
 @[tactic mem_omega]
-def evalMemOmega : Tactic := fun
-  | `(tactic| mem_omega $[$cfg]?) => do
+def evalMemOmega : Tactic := fun 
+  | `(tactic| mem_omega $[$cfg:config]? $[$v:memOmegaWith]?) => do
     let cfg ← elabMemOmegaConfig (mkOptionalNode cfg)
+    let memOmegaRules? := ← v.mapM elabMemOmegaWith
     liftMetaFinishingTactic fun g => do
-      memOmega g |>.run (← Context.init cfg)
+      memOmega g |>.run (← Context.init cfg memOmegaRules?)
   | _ => throwUnsupportedSyntax
 
 @[tactic mem_omega_bang]
 def evalMemOmegaBang : Tactic := fun
   | `(tactic| mem_omega! $[$cfg]?) => do
     let cfg ← elabMemOmegaConfig (mkOptionalNode cfg)
     liftMetaFinishingTactic fun g => do
-      memOmega g |>.run (← Context.init cfg.mkBang)
+      memOmega g |>.run (← Context.init cfg.mkBang .none)
   | _ => throwUnsupportedSyntax
 
 end MemOmega

Arm/Memory/SeparateAutomation.lean

@@ -24,6 +24,8 @@ import Tactics.Simp
 import Tactics.BvOmegaBench
 import Arm.Memory.Common
 import Arm.Memory.MemOmega
+import Lean.Elab.Tactic.Location
+import Init.Tactics
 
 open Lean Meta Elab Tactic Memory
 
@@ -382,10 +384,53 @@ Allow elaboration of `SimpMemConfig` arguments to tactics.
 -/
 declare_config_elab elabSimpMemConfig SeparateAutomation.SimpMemConfig
 
+
+/-
+This allows users to supply a list of hypotheses that simp_mem should use.
+Modeled after `rwRule`.
+-/
+syntax simpMemRule := term
+
+/-
+The kind of simplification that must be performed. If we are told
+that we must simplify a separation, a subset, or a read of a write,
+we perform this kind of simplification.
+-/
+syntax simpMemSimplificationKind := "⟂" <|> "⊂w" <|> "⊂r" (term)?
+
+
+open Lean.Parser.Tactic in
 /--
-Implement the simp_mem tactic frontend.
+The simp_mem tactic allows simplifying expressions of the form `Memory.read_bytes rbase rn (mem')`.
+`simp_mem` attempts to discover the result of the expression by various heuristics,
+which can be controlled by the end user:
+
+- (a) If `mem' = Memory.write_bytes wbase wn mem` and we know that `(rbase, rn) ⟂ (wbase, wn)`, then we simplify to `mem.read (rbase, rn)`.
+- (b) If `mem' = Memory.write_bytes wbase wn wval mem` and we kow that `(rbase, rn) ⊆ (wbase, wn)`, then we simplify to `wval.extract (rbase, rn) (wbase, wn)`.
+- (c) If we have a hypothesis `hr' : mem'.read_bytes  rbase' rn' = rval`, and we know that `(rbase, rn) ⊆ (rbase', rn')`, then we simplify to `rval.extract (rbase, rn) (rbase', rn')`.
+
+These simplifications are performed by reducing the problem to a problem that can be solved by a decision procedure (`omega`) to establish
+which hypotheses are at play. `simp_mem` can be controlled along multiple axes:
+
+1. The hypotheses that `simp_mem` will pass along to the decision procedure to discover overlapping reads (like `hr'`),
+   and hypotheses to establish memory (non-)interference, such as `(rbase, rn) ⟂ (wbase, wn)`.
+  + simp_mem using []: try to perform the rewrite using no hypotheses.
+  + simp_mem using [h₁, h₂]: try to perform the rewrite using h₁, h₂, as hypotheses.
+
+2. The kind of rewrite that simp_mem should apply. By default, it explores all possible choices, which might be expensive due to repeated calls to the decision
+   procedure. The user can describe which of (a), (b), (c) above happen:
+   + `simp_mem ⟂` : Only simplify when read is disjoint from write.
+   + `simp_mem ⊂w` : Only simplify when read overlaps the write.
+   + `simp_mem ⊂r hr` : Simplify when read overlaps with a known read `hr : mem.read_bytes baseaddr' n' = val`.
+                        This is useful for static information such as lookup tables that are at a fixed location and never modified.
+   + `simp_mem ⊂r` : Simplify when read overlaps with known read from hypothesis list.
+
+3. The targets where the rewrite must be applied. (This needs some thought: does this even make sense?)
+   + `simp_mem at ⊢`
+   + `simp_mem at h₁, h₂, ⊢`
+
 -/
-syntax (name := simp_mem) "simp_mem" (Lean.Parser.Tactic.config)? : tactic
+syntax (name := simp_mem) "simp_mem" (Lean.Parser.Tactic.config)? (simpMemSimplificationKind)? ("using" "[" withoutPosition(simpMemRule,*,?) "]")? (location)? : tactic
 
 @[tactic simp_mem]
 def evalSimpMem : Tactic := fun