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