import Spa.Analysis.Forward
import Spa.Lattice.Bool
import Spa.Lattice.Tuple
import Spa.Language.Tagged.Graphs
import Spa.Showable

namespace Spa

open Forward

instance : Showable Bool := ⟨fun b => if b then "true" else "false"⟩

instance {n : ℕ} {β : Type*} [Showable β] : Showable (Fin n → β) :=
  ⟨fun f =>
    "{" ++ (List.finRange n).foldr
        (fun i rest => show' i ++ " ↦ " ++ show' (f i) ++ ", " ++ rest) ""
        ++ "}"⟩

abbrev DefSet (prog : Program) : Type := prog.NodeId → Bool

namespace ReachingAnalysis

variable (prog : Program)

def genSet (s : prog.State) {bs : BasicStmt} (h : prog.code s = some bs) :
    DefSet prog :=
  Function.update (⊥ : DefSet prog) (prog.nodeIdOfNonempty s h) true

def eval (s : prog.State) :
    (bs : BasicStmt) → prog.code s = some bs →
    VariableValues (DefSet prog) prog → VariableValues (DefSet prog) prog
  | .assign k _, h, vs =>
    FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s h) [k] vs
  | .noop, _, vs => vs

lemma eval_mono (s : prog.State) (bs : BasicStmt) (h : prog.code s = some bs) :
    Monotone (eval prog s bs h) := by
  cases bs with
  | assign k e =>
    exact FiniteMap.generalizedUpdate_monotone monotone_id (fun _ => monotone_const)
  | noop => exact monotone_id

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) (hc : prog.code s = some bs)
      (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) hc 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 hc rest) n

instance stateInterp : StateInterp (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 → assigners n = true
  interp_sup := by
    intro vs₁ vs₂ ρ run h x assigners hmem n hla
    obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
    aesop
  interp_inf := by
    intro vs₁ vs₂ ρ run h x assigners hmem n hla
    obtain ⟨a₁, a₂, rfl, h₁, h₂⟩ := FiniteMap.mem_inf hmem
    aesop

instance validStateEvaluator : ValidStateEvaluator (DefSet prog) prog where
  step := by intro s _ _ bs hcode _ rest; exact Run.cons s bs hcode rest
  valid := by
    intro s ρ₁ ρ₂ bs vs st hcode hbs hvs
    cases hbs with
    | noop => intro x assigners hmem n hla; aesop
    | assign x e v hev =>
      intro k assigners hmem n hla
      have hmem2 : (k, assigners) ∈
          FiniteMap.generalizedUpdate id (fun _ _ => genSet prog s hcode) [x] vs := hmem
      by_cases hx : k = x
      · subst hx
        have hd := FiniteMap.generalizedUpdate_mem_eq (List.mem_singleton.mpr rfl) hmem2
        aesop (add simp genSet)
      · have hmem' := FiniteMap.generalizedUpdate_not_mem_backward
          (fun hc => hx (List.mem_singleton.mp hc)) hmem2
        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

end ReachingAnalysis

end Spa