feat: Add lemmas about setArgumentValues? and the equation lemmas

Veir/Dominance.lean

@@ -174,3 +174,32 @@ it dominates the program point before the operation, or it is a result of the op
 axiom WfIRContext.Dom.value_dominatesIp_after_iff (ctxDom : ctx.Dom) :
   value.dominatesIp (InsertPoint.after op ctx.raw block blockIsParent opInBounds) ctx ↔
   value.dominatesIp (InsertPoint.before op) ctx ∨ value ∈ op.getResults! ctx.raw
+
+/-- A value dominating the entry of a successor block either already dominates the predecessor's
+end, or it is one of the successor's own block arguments. -/
+axiom WfIRContext.Dom.value_dominatesIp_successor_entry (ctxDom : ctx.Dom)
+    {block : BlockPtr} (blockInBounds : block.InBounds ctx.raw)
+    (hsucc : succ ∈ block.getSuccessors! ctx.raw) :
+    value.dominatesIp (InsertPoint.atStart! succ ctx.raw) ctx →
+    value.dominatesIp (InsertPoint.atEnd block) ctx ∨
+      value ∈ succ.getArguments! ctx.raw
+
+/-- An operation dominating the entry of a successor already dominates the predecessor's end. -/
+axiom WfIRContext.Dom.op_dominatesIp_successor_entry (ctxDom : ctx.Dom)
+    {block : BlockPtr} (blockInBounds : block.InBounds ctx.raw)
+    (hsucc : succ ∈ block.getSuccessors! ctx.raw) :
+    op.dominatesIp (InsertPoint.atStart! succ ctx.raw) ctx →
+    op.dominatesIp (InsertPoint.atEnd block) ctx
+
+/-- An argument of a block dominates the block's start. -/
+axiom WfIRContext.Dom.blockArgument_dominatesIp_entry (ctxDom : ctx.Dom)
+    {block : BlockPtr} (blockInBounds : block.InBounds ctx.raw)
+    (hMem : value ∈ block.getArguments! ctx.raw) :
+    value.dominatesIp (InsertPoint.atStart! block ctx.raw) ctx
+
+/-- An argument of a block cannot dominate a program point that dominates the block start. -/
+axiom WfIRContext.Dom.blockArgument_not_dominatesIp_before_of_dominatesIp_firstOp
+    (ctxDom : ctx.Dom) {op : OperationPtr} (opInBounds : op.InBounds ctx.raw)
+    (opDom : op.dominatesIp (InsertPoint.atStart! block ctx.raw) ctx)
+    (hMem : value ∈ block.getArguments! ctx.raw) :
+    ¬ value.dominatesIp (InsertPoint.before op) ctx

Veir/IR/Basic.lean

@@ -2650,6 +2650,14 @@ theorem cons_iff {ctx : IRContext OpInfo} {parent : BlockPtr} {head : OperationP
 
 end BlockPtr.OpChainSlice
 
+/-- Get the successors of a block. This is defined as the successors of the terminator operation.
+In case the block has no terminator, this function returns an empty array. -/
+def BlockPtr.getSuccessors! (block : BlockPtr) (ctx : IRContext OpInfo) : Array BlockPtr :=
+  let term := (block.get! ctx).lastOp
+  match term with
+  | none => #[]
+  | some term => term.getSuccessors! ctx
+
 def IRContext.empty (OpInfo : Type) [HasOpInfo OpInfo] : IRContext OpInfo := {
     nextID := 0,
     operations := Std.HashMap.emptyWithCapacity,

Veir/Interpreter/EquationLemma.lean

@@ -204,6 +204,51 @@ theorem interpretOp_DefinesDominating {ctx : WfIRContext OpCode} {opInBounds}
   case inr =>
     grind [OperationPtr.getResults!.mem_iff_exists_index]
 
+/-- Setting a successor's block arguments preserves `DefinesDominating` in a state satisfying it at
+the predecessor's exit. -/
+theorem InterpreterState.DefinesDominating.setArgumentValues?_succ_entry
+    (ctxDom : ctx.Dom) {exitState : InterpreterState ctx}
+    {block : BlockPtr} (blockInBounds : block.InBounds ctx.raw)
+    (hsucc : succ ∈ block.getSuccessors! ctx.raw)
+    (hExit : exitState.DefinesDominating (InsertPoint.atEnd block))
+    (hArgs : exitState.variables.setArgumentValues? succ res succInBounds = some newVars) :
+    InterpreterState.DefinesDominating ⟨newVars, exitState.memory⟩ (InsertPoint.atStart! succ ctx.raw) := by
+  intro value valueInBounds valueDom
+  cases WfIRContext.Dom.value_dominatesIp_successor_entry ctxDom blockInBounds hsucc valueDom
+  · grind [InterpreterState.DefinesDominating]
+  · grind [BlockPtr.getArguments!.mem_iff_exists_index]
+
+/-- `EquationHolds` for an `op` dominating `succ`'s entry is preserved when setting `succ`'s block
+arguments. -/
+theorem InterpreterState.EquationHolds.setArgumentValues?_of_dominatesIp (ctxDom : ctx.Dom)
+    (opDom : op.dominatesIp (InsertPoint.atStart! succ ctx.raw) ctx)
+    {exitState : InterpreterState ctx} (hEq : exitState.EquationHolds op opIn)
+    (hArgs : exitState.variables.setArgumentValues? succ res succInBounds = some newVars) :
+    InterpreterState.EquationHolds ⟨newVars, exitState.memory⟩ op := by
+  simp only [InterpreterState.EquationHolds] at hEq ⊢
+  obtain ⟨cf, hinterp⟩ := hEq
+  simp only [interpretOp_some_iff] at hinterp ⊢
+  grind [VariableState.getOperandValues_eq_of_getVar?_eq,
+      VariableState.getVar?_setArgumentValues?_of_notMem_getArguments!,
+      WfIRContext.Dom.blockArgument_not_dominatesIp_before_of_dominatesIp_firstOp,
+      → VariableState.setResultValues?_setArgumentValues?_comm]
+
+/-- Setting a successor's block arguments preserves `EquationLemmaAt` in a state satisfying it at
+the predecessor's exit. -/
+theorem InterpreterState.EquationLemmaAt.setArgumentValues?_succ_entry (ctxDom : ctx.Dom)
+    {block : BlockPtr} (blockInBounds : block.InBounds ctx.raw)
+    (hsucc : succ ∈ block.getSuccessors! ctx.raw)
+    {exitState : InterpreterState ctx}
+    (hExit : exitState.EquationLemmaAt (InsertPoint.atEnd block))
+    (hArgs : exitState.variables.setArgumentValues? succ res succInBounds = some newVars) :
+    InterpreterState.EquationLemmaAt ⟨newVars, exitState.memory⟩
+      (InsertPoint.atStart! succ ctx.raw) := by
+  intro op opIn hPure hDom
+  have opDomAtEnd : op.dominatesIp (InsertPoint.atEnd block) ctx := by
+    grind [WfIRContext.Dom.op_dominatesIp_successor_entry]
+  have := hExit op opIn hPure opDomAtEnd
+  grind [InterpreterState.EquationHolds.setArgumentValues?_of_dominatesIp]
+
 /-- Interpreting a verified operation never fails on a state satisfying `DefinesDominating` at the
 operation's location. -/
 theorem InterpreterState.DefinesDominating.interpretOp_ne_none

Veir/Interpreter/Lemmas.lean

@@ -180,6 +180,90 @@ theorem VariableState.getVar?_setResultValues?_of_value_inBounds
   simp only [getVar?_setResultValues? h]
   cases value <;> grind
 
+theorem VariableState.getVar?_setArgumentValues?_loop {block : BlockPtr}
+    {values : Array RuntimeValue} {i : Nat} {iInBounds blockInBounds} :
+    VariableState.setArgumentValues?.loop block values blockInBounds varState i iInBounds = some varState' →
+    varState'.getVar? value =
+    match value with
+    | .blockArgument ⟨block', index⟩ =>
+      if block' = block ∧ index < i then
+        some values[index]!
+      else
+        varState.getVar? value
+    | _ =>
+      varState.getVar? value := by
+  fun_induction VariableState.setArgumentValues?.loop
+  next => grind
+  next =>
+    simp only [Option.bind_eq_bind, Nat.succ_eq_add_one, Option.bind]
+    grind [cases BlockArgumentPtr, OperationPtr.getResult_def, cases ValuePtr]
+
+@[grind =>]
+theorem VariableState.getVar?_setArgumentValues? {block : BlockPtr} {values : Array RuntimeValue}
+    {blockInBounds} :
+    varState.setArgumentValues? block values blockInBounds = some varState' →
+    varState'.getVar? value =
+    match value with
+    | .blockArgument ⟨block', index⟩ =>
+      if block' = block ∧ index < block.getNumArguments! ctx.raw then
+        some values[index]!
+      else
+        varState.getVar? value
+    | _ =>
+      varState.getVar? value := by
+  grind [VariableState.setArgumentValues?, VariableState.getVar?_setArgumentValues?_loop]
+
+theorem VariableState.getVar?_setArgumentValues?_of_notMem_getArguments!
+    {block : BlockPtr} {value values blockInBounds} :
+    value ∉ block.getArguments! ctx.raw →
+    varState.setArgumentValues? block values blockInBounds = some varState' →
+    varState'.getVar? value = varState.getVar? value := by
+  intro hNotMem hSet
+  simp only [VariableState.getVar?_setArgumentValues? hSet]
+  rcases value with _ | ⟨block', index⟩
+  · grind
+  · simp only [BlockPtr.getArguments!.mem_iff_exists_index, not_exists, not_and] at hNotMem
+    grind [BlockPtr.getArgument_def]
+
+/-- `block` arguments are exactly the values set by `setArgumentValues? block` in the new state. -/
+theorem VariableState.getVar?_getArgument_of_setArgumentValues?
+    {block : BlockPtr} {values blockInBounds} :
+    i < block.getNumArguments! ctx.raw →
+    varState.setArgumentValues? block values blockInBounds = some varState' →
+    varState'.getVar? (block.getArgument i) = some values[i]! := by
+  intro hi hSet
+  simp only [VariableState.getVar?_setArgumentValues? hSet]
+  grind [BlockPtr.getArgument_def]
+
+/-- `setArgumentValues?.loop` succeeds iff every argument value it binds conforms to its argument
+type. -/
+theorem VariableState.setArgumentValues?_loop_isSome_iff {block : BlockPtr}
+    {values : Array RuntimeValue} {i : Nat} {iInBounds blockInBounds} :
+    (∀ j, j < i → (values[j]!).Conforms ((block.getArgument j : ValuePtr).getType! ctx.raw)) ↔
+    (∃ v, VariableState.setArgumentValues?.loop block values blockInBounds varState i iInBounds = some v) := by
+  fun_induction VariableState.setArgumentValues?.loop
+  next => grind
+  next varState hin k hk arg value ih =>
+    simp only [Option.bind_eq_bind, Option.bind]
+    constructor
+    · intro hconform
+      rw [VariableState.setVar?_eq_some_setVar (hconform k (by grind))]
+      simp only [← ih]
+      grind
+    · rintro ⟨v, hv⟩ j hj
+      rcases hsv : hin.setVar? (ValuePtr.blockArgument arg) value with _ | varState'
+      · grind
+      · grind [ih varState']
+
+/-- `setArgumentValues?` succeeds iff every argument value conforms to its argument type. -/
+theorem VariableState.setArgumentValues?_isSome_iff_conforms (varState : VariableState ctx)
+    {block : BlockPtr} {values : Array RuntimeValue} {blockInBounds} :
+    (∀ j, j < block.getNumArguments! ctx.raw →
+      (values[j]!).Conforms ((block.getArgument j : ValuePtr).getType! ctx.raw)) ↔
+    (∃ v, varState.setArgumentValues? block values blockInBounds = some v) := by
+  simp only [VariableState.setArgumentValues?]
+  exact VariableState.setArgumentValues?_loop_isSome_iff
+
 /--
 Assert equality between two `interpretOp'` calls that have the same operation type and properties.
 This lemma is useful to avoid introducing extra casts on the `interpretOp'` arguments.
@@ -282,6 +366,48 @@ theorem VariableState.setResultValues?_comm
   ext val value
   cases val <;> grind [getVar?_setResultValues?]
 
+/-- Success of `setArgumentValues?` only depends on the values conforming to the argument types,
+not on the contents of the variable state, so it can be transferred to any other state. -/
+theorem VariableState.setArgumentValues?_eq_some_of_varState
+    (varState₂ : VariableState ctx) :
+    varState.setArgumentValues? block values blockInBounds = some varState' →
+    ∃ varState₂', varState₂.setArgumentValues? block values blockInBounds = some varState₂' := by
+  intro h
+  simp only [← VariableState.setArgumentValues?_isSome_iff_conforms]
+  grind [(VariableState.setArgumentValues?_isSome_iff_conforms varState).mpr ⟨_, h⟩]
+
+/-- Setting block arguments and operation results commute: block arguments and operation results
+are distinct `ValuePtr`s, so they never alias. -/
+theorem VariableState.setArgumentValues?_setResultValues?_comm :
+    varState.setArgumentValues? block argValues blockInBounds = some varState' →
+    varState'.setResultValues? op resValues opInBounds = some varState'' →
+    ∃ varState₂, varState.setResultValues? op resValues opInBounds = some varState₂ ∧
+    varState₂.setArgumentValues? block argValues blockInBounds = some varState'' := by
+  intros h₁ h₂
+  have ⟨varState₂, hvs₂⟩ := setResultValues?_eq_some_of_varState varState h₂
+  have ⟨varState₂', hvs₂'⟩ := setArgumentValues?_eq_some_of_varState varState₂ h₁
+  exists varState₂
+  constructor; grind
+  simp only [hvs₂', Option.some.injEq]
+  ext val value
+  grind [getVar?_setResultValues?, getVar?_setArgumentValues?]
+
+/-- Setting operation results and block arguments commute: block arguments and operation results
+are distinct `ValuePtr`s, so they never alias. -/
+theorem VariableState.setResultValues?_setArgumentValues?_comm :
+    varState.setResultValues? op resValues opInBounds = some varState' →
+    varState'.setArgumentValues? block argValues blockInBounds = some varState'' →
+    ∃ varState₂, varState.setArgumentValues? block argValues blockInBounds = some varState₂ ∧
+    varState₂.setResultValues? op resValues opInBounds = some varState'' := by
+  intros h₁ h₂
+  have ⟨varState₂, hvs₂⟩ := setArgumentValues?_eq_some_of_varState varState h₂
+  have ⟨varState₂', hvs₂'⟩ := setResultValues?_eq_some_of_varState varState₂ h₁
+  exists varState₂
+  constructor; grind
+  simp only [hvs₂', Option.some.injEq]
+  ext val value
+  grind [getVar?_setResultValues?, getVar?_setArgumentValues?]
+
 theorem VariableState.getVar?_setResultValues?_operand_of_dominates
     (ctxDom : ctx.Dom) (hdom : op'.dominates op ctx) :
     value ∈ op'.getOperands! ctx.raw →

Veir/Rewriter/InsertPoint.lean

@@ -56,6 +56,16 @@ theorem InsertPoint.atStart!_eq_atStart (block : BlockPtr) (ctx : IRContext OpIn
     InsertPoint.atStart! block ctx = InsertPoint.atStart block ctx hIn := by
   cases (block.get ctx (by grind)).firstOp <;> grind [InsertPoint.atStart!, InsertPoint.atStart]
 
+@[simp, grind =]
+theorem InsertPoint.inBounds_atStart (ctxWf : ctx.WellFormed) :
+    (InsertPoint.atStart block ctx hIn).InBounds ctx ↔ block.InBounds ctx := by
+  grind [InsertPoint.atStart]
+
+@[simp, grind =]
+theorem InsertPoint.inBounds_atStart! (ctxWf : ctx.WellFormed) (blockInBounds : block.InBounds ctx) :
+    (InsertPoint.atStart! block ctx).InBounds ctx ↔ block.InBounds ctx := by
+  grind [InsertPoint.atStart!]
+
 def InsertPoint.after (op : OperationPtr) (ctx : IRContext OpInfo) (block : BlockPtr)
     (_opHasParent : (op.get! ctx).parent = some block := by grind)
     (opInBounds : op.InBounds ctx := by grind) : InsertPoint :=