import Spa.Analysis.Forward.Lattices
import Spa.Analysis.Forward.Evaluation
import Spa.Analysis.Forward.Adapters
import Spa.Fixedpoint

namespace Spa

namespace Forward

variable {L : Type} [FiniteHeightLattice L] {prog : Program} [E : StmtEvaluator L prog]

def updateVariablesForState (s : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog := E.eval s (variablesAt s sv)

lemma updateVariablesForState_mono (s : prog.State) :
    Monotone (updateVariablesForState (L := L) s) := fun _ _ hle =>
  E.eval_mono s (variablesAt_le hle s)

def updateAll (sv : StateVariables L prog) : StateVariables L prog :=
  FiniteMap.generalizedUpdate id updateVariablesForState
    prog.states sv

lemma updateAll_mono : Monotone (updateAll (L := L) (prog := prog)) :=
  FiniteMap.generalizedUpdate_monotone monotone_id updateVariablesForState_mono

lemma updateAll_mem_eq {s : prog.State} {vs : VariableValues L prog}
    {sv : StateVariables L prog} (hmem : (s, vs) ∈ updateAll sv) :
    vs = updateVariablesForState s sv :=
  FiniteMap.generalizedUpdate_mem_eq (prog.states_complete s) hmem

lemma variablesAt_updateAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (updateAll sv) = updateVariablesForState s sv :=
  updateAll_mem_eq (variablesAt_mem s (updateAll sv))

def analyze (sv : StateVariables L prog) : StateVariables L prog :=
  updateAll (joinAll sv)

lemma analyze_mono : Monotone (analyze (L := L) (prog := prog)) := fun _ _ hle =>
  updateAll_mono (joinAll_mono hle)

variable [DecidableEq L]

variable (L prog) in
def result : StateVariables L prog :=
  Fixedpoint.aFix analyze analyze_mono

variable (L prog) in
lemma result_eq : result L prog = analyze (result L prog) :=
  Fixedpoint.aFix_eq analyze analyze_mono

lemma joinForKey_initialState :
    joinForKey prog.initialState (result L prog) = botV L prog := by
  rw [joinForKey, prog.incoming_initialState_eq_nil]
  rfl

class ValidStateEvaluator (L : Type) [FiniteHeightLattice L] (prog : Program)
    [E : StmtEvaluator L prog] [S : StateInterpretation L prog] where
  step : (s : prog.State) → {ρ₁ ρ₂ : Env} → EvalBasicStmtOpt ρ₁ (prog.code s) ρ₂ → S.St ρ₁ → S.St ρ₂
  valid : ∀ (s : prog.State) {ρ₁ ρ₂ : Env}
    {vs : VariableValues L prog} {st : S.St ρ₁},
    (hbs : EvalBasicStmtOpt ρ₁ (prog.code s) ρ₂) → ⟦ vs ⟧ ρ₁ st →
    ⟦ E.eval s vs ⟧ ρ₂ (step s hbs st)
  botV_init : ⟦ botV L prog ⟧ [] S.init

instance [LatticeInterpretation L] [ValidStmtEvaluator L prog] :
    ValidStateEvaluator L prog where
  step := by intro _ _ _ _ _; exact PUnit.unit
  valid := by intro _ _ _ _ _ hbs hvs; exact ValidStmtEvaluator.valid hbs hvs
  botV_init := by intro k l _ v hmem; cases hmem

section
variable [S : StateInterpretation L prog] [V : ValidStateEvaluator L prog]

noncomputable def stepNode (s : prog.State) {ρ₁ ρ₂ : Env}
    (h : EvalBasicStmtOpt ρ₁ (prog.code s) ρ₂) (st : S.St ρ₁) : S.St ρ₂ :=
  V.step s h st

noncomputable def stepTraceState :
    {s₁ s₂ : prog.State} → {ρ₁ ρ₂ : Env} →
    Trace prog.cfg s₁ s₂ ρ₁ ρ₂ → S.St ρ₁ → S.St ρ₂
  | s₁, _, _, _, .single hnode, st => stepNode s₁ hnode st
  | s₁, _, _, _, .edge hnode _ subtr, st =>
      stepTraceState subtr (stepNode s₁ hnode st)

/-- `Reaches tr st₁ s ρin ρout stin stout` witnesses that, when the trace `tr`
    (starting at state `st₁`) is executed, node `s` is visited at some point: `ρin`
    and `ρout` are the concrete environments just before and after `s`'s basic block
    runs, and `stin`/`stout` are the corresponding abstract execution states. A node
    inside a loop is reached once per iteration, each with its own environments. -/
inductive Reaches : {s₁ s₂ : prog.State} → {ρ₁ ρ₂ : Env} →
    Trace prog.cfg s₁ s₂ ρ₁ ρ₂ → S.St ρ₁ →
    (s : prog.State) → (ρin ρout : Env) → S.St ρin → S.St ρout → Prop
  | single_here {s₁ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂) :
      Reaches (.single hnode) st₁ s₁ ρ₁ ρ₂ st₁ (stepNode s₁ hnode st₁)
  | edge_here {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env} {st₁ : S.St ρ₁}
      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃) :
      Reaches (.edge hnode hedge rest) st₁ s₁ ρ₁ ρ₂ st₁ (stepNode s₁ hnode st₁)
  | edge_there {s₁ s₂ s₃ : prog.State} {ρ₁ ρ₂ ρ₃ : Env} {st₁ : S.St ρ₁}
      (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂)
      (hedge : (s₁, s₂) ∈ prog.cfg.edges) (rest : Trace prog.cfg s₂ s₃ ρ₂ ρ₃)
      {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout} :
      Reaches rest (stepNode s₁ hnode st₁) s ρin ρout stin stout →
      Reaches (.edge hnode hedge rest) st₁ s ρin ρout stin stout

omit [DecidableEq L] in
lemma updateAll_matches {s : prog.State} {sv : StateVariables L prog}
    {ρ₁ ρ₂ : Env} {st : S.St ρ₁}
    (hnode : EvalBasicStmtOpt ρ₁ (prog.code s) ρ₂)
    (hvs : ⟦ variablesAt s sv ⟧ ρ₁ st) :
    ⟦ variablesAt s (updateAll sv) ⟧ ρ₂ (stepNode s hnode st) := by
  rw [variablesAt_updateAll]
  exact V.valid s hnode hvs

lemma stepTrace {s₁ : prog.State} {ρ₁ ρ₂ : Env} {st : S.St ρ₁}
    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁ st)
    (hnode : EvalBasicStmtOpt ρ₁ (prog.code s₁) ρ₂) :
    ⟦ variablesAt s₁ (result L prog) ⟧ ρ₂ (stepNode s₁ hnode st) := by
  rw [result_eq L prog]
  refine updateAll_matches hnode ?_
  rw [variablesAt_joinAll]
  exact hjoin

/-- Soundness at *every* visited node: if the analysis result over-approximates the
    incoming environment at the start of the trace, then at each node reached along the
    way it over-approximates both the environment entering that node (via `joinForKey`)
    and the environment leaving it (via `variablesAt`). The intermediate `variablesAt`
    evidence used to be computed and discarded inside `walkTrace`; here it is returned. -/
lemma walkTrace_reaches {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
    {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout}
    {tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂}
    (hr : Reaches tr st₁ s ρin ρout stin stout)
    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁ st₁) :
    ⟦ joinForKey s (result L prog) ⟧ ρin stin
      ∧ ⟦ variablesAt s (result L prog) ⟧ ρout stout := by
  induction hr with
  | single_here hnode => exact ⟨hjoin, stepTrace hjoin hnode⟩
  | edge_here hnode hedge rest => exact ⟨hjoin, stepTrace hjoin hnode⟩
  | edge_there hnode hedge rest hr' ih =>
    have hstep := stepTrace hjoin hnode
    have hmem := FiniteMap.mem_valuesAt prog.states_nodup
      (prog.mem_incoming_of_edge hedge) (variablesAt_mem _ (result L prog))
    exact ih (interp_foldr hstep hmem)

omit [DecidableEq L] in
/-- The final node of a trace is always reached, with the environment/state the trace
    ends in. Used to recover the final-state soundness theorem from `walkTrace_reaches`. -/
lemma reaches_final {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} (st₁ : S.St ρ₁)
    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
    ∃ ρin, ∃ stin : S.St ρin,
      Reaches tr st₁ s₂ ρin ρ₂ stin (stepTraceState tr st₁) := by
  induction tr with
  | single hnode => exact ⟨_, _, .single_here hnode⟩
  | edge hnode hedge rest ih =>
    obtain ⟨ρin, stin, hr⟩ := ih (stepNode _ hnode st₁)
    exact ⟨ρin, stin, .edge_there hnode hedge rest hr⟩

lemma walkTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env} {st₁ : S.St ρ₁}
    (hjoin : ⟦ joinForKey s₁ (result L prog) ⟧ ρ₁ st₁)
    (tr : Trace prog.cfg s₁ s₂ ρ₁ ρ₂) :
    ⟦ variablesAt s₂ (result L prog) ⟧ ρ₂ (stepTraceState tr st₁) := by
  obtain ⟨_, _, hr⟩ := reaches_final st₁ tr
  exact (walkTrace_reaches hr hjoin).2

variable (L prog) in
theorem analyze_correct_state {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    ⟦ variablesAt prog.finalState (result L prog) ⟧ ρ
      (stepTraceState (prog.trace hrun) S.init) := by
  refine walkTrace ?_ (prog.trace hrun)
  rw [joinForKey_initialState]
  exact ValidStateEvaluator.botV_init

variable (L prog) in
/-- Soundness at every program point reached during execution: for any node `s` visited
    by the run `hrun` (witnessed by `hr`), the analysis result over-approximates both the
    environment entering `s` and the one leaving it. The final-state theorem
    `analyze_correct_state` is the special case where `s` is `prog.finalState`. -/
theorem analyze_correct_at {ρf : Env} (hrun : EvalStmt [] prog.rootStmt ρf)
    {s : prog.State} {ρin ρout : Env} {stin : S.St ρin} {stout : S.St ρout}
    (hr : Reaches (prog.trace hrun) S.init s ρin ρout stin stout) :
    ⟦ joinForKey s (result L prog) ⟧ ρin stin
      ∧ ⟦ variablesAt s (result L prog) ⟧ ρout stout := by
  refine walkTrace_reaches hr ?_
  rw [joinForKey_initialState]
  exact ValidStateEvaluator.botV_init

end

variable (L prog) in
theorem analyze_correct [LatticeInterpretation L] [ValidStmtEvaluator L prog]
    {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    ⟦ variablesAt prog.finalState (result L prog) ⟧ ρ () :=
  analyze_correct_state L prog hrun

end Forward

end Spa