import Spa.Analysis.Forward
import Spa.Lattice.Finset
import Spa.Showable

namespace Spa

open Forward

instance {n : ℕ} : Showable (Finset (Fin n)) :=
  ⟨fun s =>
    "{" ++ (List.finRange n).foldr
        (fun i rest => if i ∈ s then show' i ++ ", " ++ rest else rest) ""
        ++ "}"⟩

abbrev DefSet (prog : Program) : Type := Finset prog.State

namespace ReachingAnalysis

variable (prog : Program)

def eval (s : prog.State) (vs : VariableValues (DefSet prog) prog) : VariableValues (DefSet prog) prog :=
  match prog.code s with
  | none => vs
  | some bs =>
    match bs with
    | .assign k _ => FiniteMap.generalizedUpdate id (fun _ _ => {s}) [k] vs
    | .noop => vs

lemma eval_mono (s : prog.State) :
    Monotone (eval prog s) := by
  intros vs₁ vs₂ hle
  unfold eval; split <;> try simpa
  split <;> try simpa
  apply FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => monotone_const)
  assumption

instance stmtEvaluator : StmtEvaluator (DefSet prog) prog :=
  ⟨eval prog, eval_mono prog⟩

def output : String :=
  show' (result (DefSet prog) prog)

/-- Executed nodes, most recent first. Instructions are read from `prog.code`.
    This is `Path.steps` (chronological) reversed, so facts about concatenating
    traces reduce to mathlib's `List.append`/`List.reverse` lemmas. -/
abbrev Run (prog : Program) : Type := List prog.State

/-- The first node in a newest-first history whose instruction assigns `x`. -/
@[aesop unsafe cases]
inductive LastAssign (prog : Program) (x : String) : Run prog → prog.State → Prop
  | here (s : prog.State) (e : Expr) (rest : Run prog)
      (hc : prog.code s = some (.assign x e)) :
      LastAssign prog x (s :: rest) s
  | there (s : prog.State) (rest : Run prog) {n : prog.State} :
      (∀ e, prog.code s ≠ some (.assign x e)) → LastAssign prog x rest n →
      LastAssign prog x (s :: rest) n

def runOfPath {a b : Configuration prog.cfg} (p : Path prog.cfg a b) : Run prog :=
  p.steps.reverse

abbrev runOfTraceₗ {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
    (tr : Traceₗ prog.cfg s₁ s₂ ρ₁ ρ₂) : Run prog := runOfPath prog tr

abbrev runOfTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) : Run prog := runOfPath prog tr

instance stateInterp : StateInterpretation (DefSet prog) prog where
  Proj := Run prog
  Pre := fun tr => runOfPath prog tr
  Post := fun tr => runOfPath prog tr

  interp vs run := ∀ (x : String) (assigners : DefSet prog), (x, assigners) ∈ vs →
    ∀ (n : prog.State), LastAssign prog x run n → n ∈ assigners
  interp_sup := by
    intro vs₁ vs₂ run h x assigners hmem n hla
    obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
    aesop (add simp Finset.mem_union)
  interp_inf := by
    intro vs₁ vs₂ run h x assigners hmem n hla
    obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_inf hmem
    aesop (add simp Finset.mem_inter)

  post_pre := by
    intro vs s₁ s₂ s₃ ρ₁ ρ₂ tr hedge hvs
    simpa only [runOfPath, Trace.addEdge, Path.steps_append, Path.single,
      Path.steps, Step.steps, List.append_nil] using hvs

private lemma valid_step (s : prog.State)
    {vs : VariableValues (DefSet prog) prog} {run : Run prog}
    (hvs : ⟦vs⟧ run) :
    ⟦eval prog s vs⟧ ((match prog.code s with | none => [] | some _ => [s]) ++ run) := by
  cases hcode : prog.code s with
  | none => simpa [eval, hcode] using hvs
  | some bs =>
    cases bs with
    | noop =>
      simp [eval, hcode]
      intro x assigners hmem n hla; aesop (add simp hcode)
    | assign x e =>
      simp [eval, hcode]; intro k assigners hmem n hla
      by_cases hx : k = x
      · subst hx
        have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem
        rcases hla <;> simp [hd] <;> aesop (add simp hcode)
      · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
          (fun hc => hx (List.mem_singleton.mp hc)) hmem
        aesop (add simp hcode)

instance validStateEvaluator : ValidStateEvaluator (DefSet prog) prog where
  valid := by
    intro s₁ s₂ ρ₁ ρ₂ ρ₃ vs tr hbs hvs
    change ⟦vs⟧ (runOfPath prog tr) at hvs
    change ⟦eval prog s₂ vs⟧ (runOfPath prog (Path.append tr (.single (.execute hbs))))
    cases hcode : prog.code s₂ <;>
      simpa [runOfPath, Path.single, Path.steps, Step.steps, hcode] using valid_step prog s₂ hvs
  botV_init := by intro x assigners _ n hla; cases hla

theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    ⟦ variablesAt prog.finalState (result (DefSet prog) prog) ⟧
      (runOfTrace prog (prog.trace hrun)) :=
  Forward.analyze_correct' (DefSet prog) prog hrun

theorem analyze_correct_at {s : prog.State} {ρin ρout : Env}
    (hr : Reaches s ρin ρout) :
    ⟦ joinForKey s (result (DefSet prog) prog) ⟧ (runOfTraceₗ prog hr.pre)
      ∧ ⟦ variablesAt s (result (DefSet prog) prog) ⟧ (runOfTrace prog hr.post) :=
  Forward.analyze_correct_at (DefSet prog) prog hr

end ReachingAnalysis

end Spa