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

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] [DecidableEq L] {prog : Program} {h : ℕ}
  (fhL : FixedHeight L h) (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)

omit [DecidableEq L] in
/-- Agda: `updateVariablesForState-Monoʳ`. -/
theorem updateVariablesForState_mono (s : prog.State) :
    Monotone (updateVariablesForState E 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 E s sv)
    prog.states sv

omit [DecidableEq L] in
/-- Agda: `updateAll-Mono`. -/
theorem updateAll_mono : Monotone (updateAll E) :=
  FiniteMap.generalizedUpdate_monotone monotone_id (updateVariablesForState_mono E)

omit [DecidableEq L] in
/-- Agda: `updateAll-k∈ks-≡`. -/
theorem updateAll_mem_eq {s : prog.State} {vs : VariableValues L prog}
    {sv : StateVariables L prog} (hmem : (s, vs) ∈ updateAll E sv) :
    vs = updateVariablesForState E s sv :=
  FiniteMap.generalizedUpdate_mem_eq (prog.states_complete s) hmem

omit [DecidableEq L] in
/-- Agda: `variablesAt-updateAll`. -/
theorem variablesAt_updateAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (updateAll E sv) = updateVariablesForState E s sv :=
  updateAll_mem_eq E (variablesAt_mem s (updateAll E sv))

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

omit [DecidableEq L] in
/-- Agda: `analyze-Mono`. -/
theorem analyze_mono : Monotone (analyze fhL E) := fun _ _ hle =>
  updateAll_mono E (joinAll_mono fhL hle)

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

/-- Agda: `result≈analyze-result`. -/
theorem result_eq : result fhL E = analyze fhL E (result fhL E) :=
  Fixedpoint.aFix_eq (statesFixedHeight L prog fhL) (analyze fhL E) (analyze_mono fhL E)

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

variable {I : LatticeInterpretation L} {E}
variable (hE : IsValidStmtEvaluator E I)
include hE

omit [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 I vs ρ₁) :
    interpV I (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 (hE hbs hvs)

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

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

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

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

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

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

/-- 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 I (variablesAt prog.finalState (result fhL E)) ρ :=
  walkTrace fhL hE (interpV_joinForKey_initialState fhL (E := E) (I := I))
    (prog.trace hrun)

end Spa
