/-
Port of `Analysis/Forward.agda` (`WithProg`, `WithStmtEvaluator`,
`WithValidInterpretation`).

As in Agda, the statement evaluator, the lattice interpretation and the
evaluator's validity proof are instance arguments (`{{evaluator}}`,
`{{latticeInterpretationˡ}}`, `{{validEvaluator}}`); `result` and
`analyze_correct` take `L` and `prog` explicitly, mirroring the Agda call
shape `WithProg.result L prog`.

Correspondence:
  updateVariablesForState, -Monoʳ   ↦ updateVariablesForState, _mono
  updateAll, updateAll-Mono,
  updateAll-k∈ks-≡                  ↦ updateAll, updateAll_mono, updateAll_mem_eq
  analyze, analyze-Mono             ↦ analyze, analyze_mono
  result, result≈analyze-result     ↦ result, result_eq
  variablesAt-updateAll             ↦ variablesAt_updateAll
  eval-fold-valid                   ↦ eval_fold_valid
  updateVariablesForState-matches   ↦ updateVariablesForState_matches
  updateAll-matches                 ↦ updateAll_matches
  stepTrace                         ↦ stepTrace   (the `subst`/`⟦⟧ᵛ-respects-≈ᵛ`
                                       plumbing becomes plain rewriting with `=`)
  walkTrace                         ↦ walkTrace
  joinForKey-initialState-⊥ᵛ        ↦ joinForKey_initialState
  ⟦joinAll-initialState⟧ᵛ∅          ↦ interpV_joinForKey_initialState
  analyze-correct                   ↦ analyze_correct
-/
import Spa.Analysis.Forward.Lattices
import Spa.Analysis.Forward.Evaluation
import Spa.Analysis.Forward.Adapters
import Spa.Fixedpoint

namespace Spa

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

/-- Agda: `updateVariablesForState`. -/
def updateVariablesForState (s : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (prog.code s).foldl (fun vs bs => E.eval s bs vs) (variablesAt s sv)

/-- Agda: `updateVariablesForState-Monoʳ`. -/
theorem updateVariablesForState_mono (s : prog.State) :
    Monotone (updateVariablesForState (L := L) s) := fun _ _ hle =>
  foldl_mono' (prog.code s) _ (fun bs => E.eval_mono s bs) (variablesAt_le hle s)

/-- Agda: `updateAll`. -/
def updateAll (sv : StateVariables L prog) : StateVariables L prog :=
  FiniteMap.generalizedUpdate id (fun s sv => updateVariablesForState s sv)
    prog.states sv

/-- Agda: `updateAll-Mono`. -/
theorem updateAll_mono : Monotone (updateAll (L := L) (prog := prog)) :=
  FiniteMap.generalizedUpdate_monotone monotone_id updateVariablesForState_mono

/-- Agda: `updateAll-k∈ks-≡`. -/
theorem 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

/-- Agda: `variablesAt-updateAll`. -/
theorem variablesAt_updateAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (updateAll sv) = updateVariablesForState s sv :=
  updateAll_mem_eq (variablesAt_mem s (updateAll sv))

variable [FiniteHeightLattice L]

/-- Agda: `analyze`. -/
def analyze (sv : StateVariables L prog) : StateVariables L prog :=
  updateAll (joinAll sv)

/-- Agda: `analyze-Mono`. -/
theorem analyze_mono : Monotone (analyze (L := L) (prog := prog)) := fun _ _ hle =>
  updateAll_mono (joinAll_mono hle)

variable [DecidableEq L]

variable (L prog) in
/-- Agda: `result` (the least fixpoint of `analyze`). -/
def result : StateVariables L prog :=
  Fixedpoint.aFix analyze analyze_mono

variable (L prog) in
/-- Agda: `result≈analyze-result`. -/
theorem result_eq : result L prog = analyze (result L prog) :=
  Fixedpoint.aFix_eq analyze analyze_mono

/-- Agda: `joinForKey-initialState-⊥ᵛ`. -/
theorem joinForKey_initialState :
    joinForKey prog.initialState (result L prog) = botV L prog := by
  rw [joinForKey, prog.incoming_initialState_eq_nil]
  rfl

/-! ### Semantic correctness (Agda: `WithValidInterpretation`) -/

variable [I : LatticeInterpretation L] [V : ValidStmtEvaluator L prog]

omit [FiniteHeightLattice L] [DecidableEq L] in
/-- Agda: `eval-fold-valid`. -/
theorem eval_fold_valid {s : prog.State} {bss : List BasicStmt}
    {vs : VariableValues L prog} {ρ₁ ρ₂ : Env}
    (hbss : EvalBasicStmts ρ₁ bss ρ₂) (hvs : interpV vs ρ₁) :
    interpV (bss.foldl (fun vs bs => E.eval s bs vs) vs) ρ₂ := by
  induction hbss generalizing vs with
  | nil => exact hvs
  | cons hbs _ ih => exact ih (ValidStmtEvaluator.valid hbs hvs)

omit [FiniteHeightLattice L] [DecidableEq L] in
/-- Agda: `updateVariablesForState-matches`. -/
theorem updateVariablesForState_matches {s : prog.State}
    {sv : StateVariables L prog} {ρ₁ ρ₂ : Env}
    (hbss : EvalBasicStmts ρ₁ (prog.code s) ρ₂)
    (hvs : interpV (variablesAt s sv) ρ₁) :
    interpV (updateVariablesForState s sv) ρ₂ :=
  eval_fold_valid hbss hvs

omit [FiniteHeightLattice L] [DecidableEq L] in
/-- Agda: `updateAll-matches`. -/
theorem updateAll_matches {s : prog.State} {sv : StateVariables L prog}
    {ρ₁ ρ₂ : Env} (hbss : EvalBasicStmts ρ₁ (prog.code s) ρ₂)
    (hvs : interpV (variablesAt s sv) ρ₁) :
    interpV (variablesAt s (updateAll sv)) ρ₂ := by
  rw [variablesAt_updateAll]
  exact updateVariablesForState_matches hbss hvs

/-- Agda: `stepTrace`. -/
theorem stepTrace {s₁ : prog.State} {ρ₁ ρ₂ : Env}
    (hjoin : interpV (joinForKey s₁ (result L prog)) ρ₁)
    (hbss : EvalBasicStmts ρ₁ (prog.code s₁) ρ₂) :
    interpV (variablesAt s₁ (result L prog)) ρ₂ := by
  rw [result_eq L prog]
  refine updateAll_matches hbss ?_
  rw [variablesAt_joinAll]
  exact hjoin

/-- Agda: `walkTrace`. -/
theorem walkTrace {s₁ s₂ : prog.State} {ρ₁ ρ₂ : Env}
    (hjoin : interpV (joinForKey s₁ (result L prog)) ρ₁)
    (tr : Trace prog.graph s₁ s₂ ρ₁ ρ₂) :
    interpV (variablesAt s₂ (result L prog)) ρ₂ := by
  induction tr with
  | single hbss => exact stepTrace hjoin hbss
  | @edge _ ρ' _ i₁ i₂ _ hbss hedge _ ih =>
    have hstep : interpV (variablesAt i₁ (result L prog)) ρ' :=
      stepTrace hjoin hbss
    have hmem : variablesAt i₁ (result L prog)
        ∈ (result L prog).valuesAt (prog.incoming i₂) :=
      FiniteMap.mem_valuesAt prog.states_nodup
        (prog.mem_incoming_of_edge hedge) (variablesAt_mem i₁ (result L prog))
    exact ih (interpV_foldr hstep hmem)

omit V in
/-- Agda: `⟦joinAll-initialState⟧ᵛ∅`. -/
theorem interpV_joinForKey_initialState :
    interpV (joinForKey prog.initialState (result L prog)) [] := by
  rw [joinForKey_initialState]
  exact interpV_botV_nil

variable (L prog) in
/-- Agda: `analyze-correct` — the analysis result at the final state soundly
describes every terminating execution of the program. -/
theorem analyze_correct {ρ : Env} (hrun : EvalStmt [] prog.rootStmt ρ) :
    interpV (variablesAt prog.finalState (result L prog)) ρ :=
  walkTrace interpV_joinForKey_initialState (prog.trace hrun)

end Spa