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)

inductive Run (prog : Program) where
  | nil : Run prog
  | cons (s : prog.State) (bs : BasicStmt)
      (rest : Run prog) : Run prog

@[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 (Run.cons 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 (Run.cons s bs rest) n

instance stateInterp : StateInterpretation (DefSet prog) prog where
  St := fun _ => Run prog
  init := Run.nil
  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)

private def stepAt (s : prog.State) (obs : Option BasicStmt) { ρ₁ ρ₂ : Env} : EvalBasicStmtOpt ρ₁ obs ρ₂ → Run prog → Run prog
  | .none, rest => rest
  | .some (bs := bs) _, rest => Run.cons s bs rest

instance validStateEvaluator : ValidStateEvaluator (DefSet prog) prog where
  step := fun s ρ₁ ρ₂ => stepAt prog s (prog.code s)
  valid := by
    simp [StmtEvaluator.eval, eval];
    intro s ρ₁ ρ₂ vs; generalize prog.code s = obs; intro hst hbs hvs
    rcases hbs with _ | @⟨_, bs, hbs⟩; try (simpa [stepAt])
    cases hbs with
    | noop => intro x assigners hmem n hla; aesop
    | assign x e v hev =>
      simp; 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
  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) ⟧ ρ
      (stepTraceState (prog.trace hrun) (stateInterp prog).init) :=
  Forward.analyze_correct_state (DefSet prog) prog hrun

theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
    {s : prog.State} {ρin ρout : Env} {stin : Run prog} {stout : Run prog}
    (hr : Reaches (prog.trace hrun) (stateInterp prog).init s ρin ρout stin stout) :
    ⟦ joinForKey s (result (DefSet prog) prog) ⟧ ρin stin
      ∧ ⟦ variablesAt s (result (DefSet prog) prog) ⟧ ρout stout :=
  Forward.analyze_correct_at (DefSet prog) prog hrun hr

end ReachingAnalysis

end Spa