Add intermediate proof that rhs of hoisted assignments are unchanged

lean/Spa/Transformation/Licm/Correctness.lean

@@ -0,0 +1,80 @@
+import Spa.Transformation.Licm
+import Spa.Analysis.Reaching.Paths
+
+namespace Spa
+namespace LicmTransformation
+open GGraph Forward ReachingAnalysis
+
+private lemma mem_union_fold {α β : Type} [DecidableEq β] (f : α → Finset β)
+    (xs : List α) (acc : Finset β) (d : β) :
+    d ∈ xs.foldl (fun a x => a ∪ f x) acc ↔ d ∈ acc ∨ ∃ x ∈ xs, d ∈ f x := by
+  induction xs generalizing acc with
+  | nil => simp
+  | cons x xs ih => simp [List.foldl, ih, or_assoc, or_left_comm, or_comm]
+
+/-- The executable invariant test excludes each actual reaching definition. -/
+lemma isInvariant_sound_at {prog : Program} {c : Candidate prog} {ρ ρ' : Env}
+    {x : String} {d : prog.State}
+    (hinv : isInvariant prog c = true) (hx : x ∈ c.rhsVars) (hxp : x ∈ prog.vars)
+    (hr : Reaches c.assignState ρ ρ')
+    (hl : LastAssign prog x (runOfPath prog hr.pre) d) : c.encl.covers d = false := by
+  let entry := joinForKey c.assignState (result (DefSet prog) prog)
+  have hk : FiniteMap.MemKey x entry := hxp
+  have hd : d ∈ lookupDef prog entry x := by
+    have hs := (ReachingAnalysis.analyze_correct_at prog hr).1
+    have hm := (FiniteMap.locate hk).2
+    have hd := hs x (FiniteMap.locate hk).1 hm d hl
+    simpa [lookupDef, hk] using hd
+  have hall : ∀ d ∈ c.rhsVars.foldl (fun acc k => acc ∪ lookupDef prog entry k) ∅,
+      c.encl.covers d = false := by simpa [isInvariant, entry] using hinv
+  exact hall d ((mem_union_fold _ _ _ _).mpr (Or.inr ⟨x, hx, hd⟩))
+
+/-- A path inside the loop can execute statements only within its body range. -/
+lemma loop_steps_covered {prog : Program} {cond : Expr} {x : String} {rhs : Expr} {tail : Stmt}
+    (e : Embed (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg prog.cfg)
+    {a b : Configuration (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg}
+    (seg : Path _ a b) {d : prog.State}
+    (hm : d ∈ (seg.embed e).steps) :
+    (headCandidate prog cond x rhs tail e).encl.covers d = true := by
+  rw [Path.steps_embed] at hm
+  obtain ⟨j, hj, rfl⟩ := List.mem_map.mp hm
+  obtain ⟨bs, hcode⟩ := Option.ne_none_iff_exists'.mp (seg.steps_nonempty hj)
+  obtain ⟨i, hi⟩ := GGraph.loop_node_in_body hcode
+  apply decide_eq_true
+  apply (Embed.mem_range_iff ((Embed.loop _).trans e) _).mp
+  exact ⟨i, congrArg e.f hi⟩
+
+/-- The analysis and the actual intervening path together establish RHS
+stability. This is the bridge from static sites to unchanged runtime values. -/
+lemma head_rhs_agrees {prog : Program} {cond : Expr} {x : String} {rhs : Expr} {tail : Stmt}
+    (e : Embed (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg prog.cfg)
+    {i j : (Stmt.whileLoop cond (.andThen (.basic (.assign x rhs)) tail)).cfg.Index}
+    {ρ₀ ρ₁ ρ₂ : Env}
+    (pre : Traceₗ prog.cfg prog.initialState (e.f i) [] ρ₀)
+    (seg : Traceₗ _ i j ρ₀ ρ₁)
+    (hj : e.f j = (headCandidate prog cond x rhs tail e).assignState)
+    (step : EvalBasicStmt ρ₁ (.assign x rhs) ρ₂)
+    (hinv : isInvariant prog (headCandidate prog cond x rhs tail e) = true) :
+    Env.AgreeOn rhs.vars ρ₀ ρ₁ := by
+  have hj' : j = ((Embed.sequenceLeft (Stmt.basic (.assign x rhs)).cfg tail.cfg).trans
+      (Embed.loop (Stmt.andThen (.basic (.assign x rhs)) tail).cfg)).singletonIndex := e.f_inj hj
+  subst j
+  let c := headCandidate prog cond x rhs tail e
+  have hcode : prog.code c.assignState = some (.assign x rhs) :=
+    Embed.nodes_singletonIndex
+      ((Embed.sequenceLeft (Stmt.basic (.assign x rhs)).cfg tail.cfg).trans
+        ((Embed.loop (Stmt.andThen (.basic (.assign x rhs)) tail).cfg).trans e))
+  let reach : Reaches c.assignState ρ₁ ρ₂ :=
+    ⟨pre.append (seg.embed e), hcode ▸ .some step⟩
+  intro y hy
+  apply ReachingAnalysis.Path.preserves_of_lastAssign_outside pre (seg.embed e)
+    {d | c.encl.covers d = true}
+  · intro d hm; exact loop_steps_covered e seg hm
+  · intro d hl hd
+    have hl' : LastAssign prog y (runOfPath prog reach.pre) d := hl
+    have hf := isInvariant_sound_at hinv (by simpa [c, headCandidate] using hy)
+      (Program.code_vars hcode y (Finset.mem_union_right _ hy)) reach hl'
+    exact Bool.noConfusion (hd.symm.trans hf)
+
+end LicmTransformation
+end Spa