import Spa.Language.Properties
import Spa.Language.Equivalence

namespace Spa
open GGraph

/-- Recorded nodes contain instructions; empty CFG nodes are omitted from the history. -/
lemma Path.steps_nonempty {g : Graph} {a b : Configuration g} (p : Path g a b)
    {d : g.Index} (hm : d ∈ p.steps) : g.nodes d ≠ none := by
  induction p with
  | nil => simp [Path.steps] at hm
  | cons st p ih =>
      rcases List.mem_append.mp hm with hs | hp
      · cases st with
        | edge => simp [Step.steps] at hs
        | @execute i ρ σ h =>
            cases hc : g.nodes i <;> aesop (add simp [Step.steps, hc])
      · exact ih hp

private lemma optional_preserves_unwritten {ρ σ : Env} {obs : Option BasicStmt}
    (h : EvalBasicStmtOpt ρ obs σ) (x : String)
    (hn : ∀ rhs, obs ≠ some (.assign x rhs)) :
    ∀ v, Env.Mem (x, v) ρ ↔ Env.Mem (x, v) σ := by
  cases h with
  | none => exact fun _ => Iff.rfl
  | some h =>
      cases h with
      | noop => exact fun _ => Iff.rfl
      | assign y rhs w hv =>
          have hxy : x ≠ y := by
            rintro rfl
            exact hn rhs rfl
          intro v; simp [Env.mem_cons, hxy]

/-- A path whose executed nodes do not assign `x` preserves its binding. -/
lemma Path.preserves_unwritten {g : Graph} {a b : Configuration g} (p : Path g a b)
    {x : String} (hn : ∀ d ∈ p.steps, ∀ rhs, g.nodes d ≠ some (.assign x rhs)) :
    ∀ v, Env.Mem (x, v) a.2 ↔ Env.Mem (x, v) b.2 := by
  induction p with
  | nil => exact fun _ => Iff.rfl
  | cons st p ih =>
      have ht := ih (fun d hm => hn d (List.mem_append_right _ hm))
      suffices hs : ∀ v, Env.Mem (x, v) _ ↔ Env.Mem (x, v) _ from
        fun v => (hs v).trans (ht v)
      cases st with
      | edge => exact fun _ => Iff.rfl
      | execute h =>
          apply optional_preserves_unwritten h x
          intro rhs hc
          exact hn _ (List.mem_append_left _ (by simp [Step.steps, hc])) rhs hc

lemma Step.steps_embed {g h : Graph} (e : Embed g h) {a b : Configuration g}
    (s : Step g a b) :
    (s.embed e).steps = s.steps.map e.f := by
  cases s with
  | edge => rfl
  | @execute i ρ σ h =>
      simp only [Step.embed, Step.steps, e.nodes_eq]
      cases g.nodes i <;> rfl

lemma Path.steps_embed {g h : Graph} (e : Embed g h) {a b : Configuration g}
    (p : Path g a b) :
    (p.embed e).steps = p.steps.map e.f := by
  induction p <;> aesop (add simp [Path.embed, Path.steps, Step.steps_embed])

/-- Every nonempty node in a loop belongs to its body. -/
lemma GGraph.loop_node_in_body {g : Graph} {i : (Graph.loop g).Index} {bs : BasicStmt}
    (hc : (Graph.loop g).nodes i = some bs) : ∃ j, (Embed.loop g).f j = i := by
  refine Fin.addCases ?_ ?_ i hc
  · intro j hj
    simp [Graph.loop, Fin.append_left] at hj
  · intro j _; exact ⟨j, rfl⟩

/-- Variables at any CFG statement occur in its source statement. -/
lemma Stmt.cfg_node_vars {s : Stmt} {i : s.cfg.Index} {bs : BasicStmt}
    (hc : s.cfg.nodes i = some bs) : bs.vars ⊆ s.vars := by
  induction s with
  | basic b =>
      have : b = bs := Option.some.inj hc
      subst bs; exact Finset.Subset.refl _
  | andThen a b iha ihb =>
      refine Fin.addCases ?_ ?_ i hc
      · intro j hj; have hv := iha (by simpa [Stmt.cfg, Graph.sequence] using hj)
        exact fun x hx => Finset.mem_union_left _ (hv hx)
      · intro j hj; have hv := ihb (by simpa [Stmt.cfg, Graph.sequence] using hj)
        exact fun x hx => Finset.mem_union_right _ (hv hx)
  | ifElse cond a b iha ihb =>
      refine Fin.addCases ?_ ?_ i hc
      · intro j hj; have hv := iha (by simpa [Stmt.cfg, Graph.overlay] using hj)
        exact fun x hx => Finset.mem_union_left _ (Finset.mem_union_right _ (hv hx))
      · intro j hj; have hv := ihb (by simpa [Stmt.cfg, Graph.overlay] using hj)
        exact fun x hx => Finset.mem_union_right _ (hv hx)
  | whileLoop cond body ih =>
      obtain ⟨j, rfl⟩ := GGraph.loop_node_in_body hc
      have hv := ih (((Embed.loop body.cfg).nodes_eq j).symm.trans hc)
      exact fun x hx => Finset.mem_union_right _ (hv hx)

lemma Program.code_vars {prog : Program} {i : prog.State} {bs : BasicStmt}
    (hc : prog.code i = some bs) : ∀ x ∈ bs.vars, x ∈ prog.vars := by
  have hroot : ∃ j, prog.rootStmt.cfg.nodes j = some bs := by
    unfold Program.code Program.cfg Graph.wrap at hc
    revert hc
    refine Fin.addCases ?_ ?_ i
    · intro j hj; simp [Graph.sequence, Graph.singleton] at hj
    · intro j
      refine Fin.addCases ?_ ?_ j
      · intro k hk
        exact ⟨k, by simpa [Graph.sequence] using hk⟩
      · intro k hk; simp [Graph.sequence, Graph.singleton] at hk
  obtain ⟨j, hj⟩ := hroot
  intro x hx
  simpa [Program.vars] using Stmt.cfg_node_vars hj hx

end Spa