import Spa.Analysis.Forward
import Spa.Lattice.Finset
import Spa.Language.Tagged.Graphs
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.NodeId

namespace ReachingAnalysis

variable (prog : Program)

def genSet (s : prog.State) : DefSet prog := (prog.nodeIdOf s).elim {} (fun x => {x})

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 _ _ => genSet prog 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)

/-- The statements a trace executed, paired with the state each executed at,
    most recent first (matching `LastAssign`, which scans for the most recent
    assignment). This is `Trace.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 × BasicStmt)

@[aesop unsafe cases]
inductive LastAssign (prog : Program) (x : String) : Run prog → prog.NodeId → Prop
  | here (s : prog.State) (e : Expr) (hc : prog.code s = some (.assign x e))
      (rest : Run prog) :
      LastAssign prog x ((s, .assign x e) :: rest) (prog.nodeIdOfNonempty s hc)
  | there (s : prog.State) (bs : BasicStmt) (hc : prog.code s = some bs)
      (rest : Run prog) {n : prog.NodeId} :
      (∀ e, bs ≠ .assign x e) → LastAssign prog x rest n →
      LastAssign prog x ((s, bs) :: rest) n

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

def runOfTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) : Run prog :=
  tr.steps.reverse

instance stateInterp : StateInterpretation (DefSet prog) prog where
  Proj := Run prog
  Pre := @runOfTraceₗ prog
  Post := @runOfTrace prog

  interp vs run := ∀ (x : String) (assigners : DefSet prog), (x, assigners) ∈ vs →
    ∀ (n : prog.NodeId), 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 [runOfTrace, runOfTraceₗ] using hvs

private lemma valid_step (s : prog.State) {ρ₁ ρ₂ : Env}
    {obs : Option BasicStmt} (hcode : prog.code s = obs)
    (hbs : EvalBasicStmtOpt ρ₁ obs ρ₂)
    {vs : VariableValues (DefSet prog) prog} {run : Run prog}
    (hvs : ⟦vs⟧ run) :
    ⟦eval prog s vs⟧ ((hbs.steps s).reverse ++ run) := by
  cases hbs with
  | none => simpa [eval, hcode, EvalBasicStmtOpt.steps] using hvs
  | some hbs =>
  cases hbs with
  | noop =>
    simp [eval, hcode, EvalBasicStmtOpt.steps]
    intro x assigners hmem n hla; aesop
  | assign x e v hev =>
    simp [eval, hcode, EvalBasicStmtOpt.steps]; 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 [Program.nodeIdOfNonempty, hd, genSet, Option.get] <;> aesop
    · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
        (fun hc => hx (List.mem_singleton.mp hc)) hmem
      aesop

instance validStateEvaluator : ValidStateEvaluator (DefSet prog) prog where
  valid := by
    intro s₁ s₂ ρ₁ ρ₂ ρ₃ vs tr hbs hvs
    show ⟦eval prog s₂ vs⟧ (runOfTrace prog (tr ++ hbs))
    simpa [runOfTrace, runOfTraceₗ] using valid_step prog s₂ rfl hbs 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 {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
    {s : prog.State} {ρin ρout : Env}
    (hr : Reaches (prog.trace hrun) 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 hrun hr

end ReachingAnalysis

end Spa