/-
Port of `Analysis/Forward/Lattices.agda`.

The Agda module instantiates `Lattice.FiniteMap` twice (variables ↦ abstract
values, states ↦ variable maps) and re-exports everything with ᵛ/ᵐ suffixes.
In Lean the two instantiations are `abbrev`s and the FiniteMap API is used
directly; the module parameters (the finite-height lattice `L`, the program)
become section variables.

Correspondence:
  VariableValues, StateVariables  ↦ VariableValues, StateVariables
  isLatticeᵛ/isLatticeᵐ, ⊔ᵛ, ≼ᵛ … ↦ (the FiniteMap Lattice instances)
  fixedHeightᵛ                    ↦ varsFixedHeight
  ⊥ᵛ, ⊥ᵛ-contains-bottoms         ↦ botV, FiniteMap.bot_contains_bots
  states-in-Map                   ↦ states_memKey
  variablesAt                     ↦ variablesAt
  variablesAt-∈                   ↦ variablesAt_mem
  variablesAt-≈                   ↦ (congruence, trivial with `=`)
  joinForKey, joinForKey-Mono     ↦ joinForKey, joinForKey_mono
  joinAll, joinAll-Mono,
  joinAll-k∈ks-≡                  ↦ joinAll, joinAll_mono, joinAll_mem_eq
  variablesAt-joinAll             ↦ variablesAt_joinAll
  ⟦_⟧ᵛ                            ↦ interpV
  ⟦⊥ᵛ⟧ᵛ∅                          ↦ interpV_botV_nil
  ⟦⟧ᵛ-respects-≈ᵛ                 ↦ (trivial with `=`)
  ⟦⟧ᵛ-⊔ᵛ-∨                        ↦ interpV_sup
  ⟦⟧ᵛ-foldr                       ↦ interpV_foldr
-/
import Spa.Language
import Spa.Lattice.FiniteMap

namespace Spa

variable (L : Type) [Lattice L] (prog : Program)

/-- Agda: `VariableValues`. -/
abbrev VariableValues : Type := FiniteMap String L prog.vars

/-- Agda: `StateVariables`. -/
abbrev StateVariables : Type := FiniteMap prog.State (VariableValues L prog) prog.states

variable {h : ℕ}

/-- Agda: `fixedHeightᵛ`. -/
def varsFixedHeight (fhL : FixedHeight L h) :
    FixedHeight (VariableValues L prog) (prog.vars.length * h) :=
  FiniteMap.fixedHeight fhL prog.vars

/-- Agda: `⊥ᵛ`. -/
def botV (fhL : FixedHeight L h) : VariableValues L prog :=
  (varsFixedHeight L prog fhL).bot

/-- Agda: `fixedHeight` on `StateVariables` (assembled in `Forward.agda`'s
fixpoint call; named here for reuse). -/
def statesFixedHeight (fhL : FixedHeight L h) :
    FixedHeight (StateVariables L prog) (prog.states.length * (prog.vars.length * h)) :=
  FiniteMap.fixedHeight (varsFixedHeight L prog fhL) prog.states

variable {L prog}

omit [Lattice L] in
/-- Agda: `states-in-Map`. -/
theorem states_memKey (s : prog.State) (sv : StateVariables L prog) :
    FiniteMap.MemKey s sv :=
  FiniteMap.memKey_iff.mpr (prog.states_complete s)

/-- Agda: `variablesAt`. -/
def variablesAt (s : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (FiniteMap.locate (states_memKey s sv)).1

omit [Lattice L] in
/-- Agda: `variablesAt-∈`. -/
theorem variablesAt_mem (s : prog.State) (sv : StateVariables L prog) :
    (s, variablesAt s sv) ∈ sv :=
  (FiniteMap.locate (states_memKey s sv)).2

/-- Agda: `m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ`, specialized the way `Forward.agda` uses it. -/
theorem variablesAt_le {sv₁ sv₂ : StateVariables L prog} (hle : sv₁ ≤ sv₂)
    (s : prog.State) : variablesAt s sv₁ ≤ variablesAt s sv₂ :=
  FiniteMap.le_of_mem_mem prog.states_nodup hle
    (variablesAt_mem s sv₁) (variablesAt_mem s sv₂)

variable (fhL : FixedHeight L h)

/-- Agda: `joinForKey`. -/
def joinForKey (k : prog.State) (sv : StateVariables L prog) :
    VariableValues L prog :=
  (sv.valuesAt (prog.incoming k)).foldr (· ⊔ ·) (botV L prog fhL)

/-- Agda: `joinForKey-Mono`. -/
theorem joinForKey_mono (k : prog.State) :
    Monotone (joinForKey fhL k) := by
  intro sv₁ sv₂ hle
  exact foldr_mono _ (FiniteMap.valuesAt_le hle (prog.incoming k)) (le_refl _)
    (fun b _ _ hab => sup_le_sup_right hab b)
    (fun a _ _ hab => sup_le_sup_left hab a)

/-- Agda: `joinAll` (the "Exercise 4.26" generalized update with `f = id`). -/
def joinAll (sv : StateVariables L prog) : StateVariables L prog :=
  FiniteMap.generalizedUpdate id (joinForKey fhL) prog.states sv

/-- Agda: `joinAll-Mono`. -/
theorem joinAll_mono : Monotone (joinAll (prog := prog) fhL) :=
  FiniteMap.generalizedUpdate_monotone monotone_id (joinForKey_mono fhL)

/-- Agda: `joinAll-k∈ks-≡`. -/
theorem joinAll_mem_eq {s : prog.State} {vs : VariableValues L prog}
    {sv : StateVariables L prog} (h : (s, vs) ∈ joinAll fhL sv) :
    vs = joinForKey fhL s sv :=
  FiniteMap.generalizedUpdate_mem_eq (prog.states_complete s) h

/-- Agda: `variablesAt-joinAll`. -/
theorem variablesAt_joinAll (s : prog.State) (sv : StateVariables L prog) :
    variablesAt s (joinAll fhL sv) = joinForKey fhL s sv :=
  joinAll_mem_eq fhL (variablesAt_mem s (joinAll fhL sv))

/-! ### Lifting an interpretation to variable maps -/

variable (I : LatticeInterpretation L)

/-- Agda: `⟦_⟧ᵛ`. -/
def interpV (vs : VariableValues L prog) (ρ : Env) : Prop :=
  ∀ (k : String) (l : L), (k, l) ∈ vs →
    ∀ (v : Value), Env.Mem (k, v) ρ → I.interp l v

/-- Agda: `⟦⊥ᵛ⟧ᵛ∅`. -/
theorem interpV_botV_nil : interpV I (botV L prog fhL) [] := by
  intro k l _ v hmem
  cases hmem

/-- Agda: `⟦⟧ᵛ-⊔ᵛ-∨`. -/
theorem interpV_sup {vs₁ vs₂ : VariableValues L prog} {ρ : Env}
    (h : interpV I vs₁ ρ ∨ interpV I vs₂ ρ) : interpV I (vs₁ ⊔ vs₂) ρ := by
  intro k l hmem v hv
  obtain ⟨l₁, l₂, rfl, h₁, h₂⟩ := FiniteMap.mem_sup hmem
  rcases h with h | h
  · exact I.interp_sup v (Or.inl (h _ _ h₁ _ hv))
  · exact I.interp_sup v (Or.inr (h _ _ h₂ _ hv))

/-- Agda: `⟦⟧ᵛ-foldr`. -/
theorem interpV_foldr {vs : VariableValues L prog}
    {vss : List (VariableValues L prog)} {ρ : Env}
    (hvs : interpV I vs ρ) (hmem : vs ∈ vss) :
    interpV I (vss.foldr (· ⊔ ·) (botV L prog fhL)) ρ := by
  induction vss with
  | nil => cases hmem
  | cons vs' vss' ih =>
    rcases List.mem_cons.mp hmem with rfl | hmem'
    · exact interpV_sup I (Or.inl hvs)
    · exact interpV_sup I (Or.inr (ih hmem'))

end Spa