open import Language hiding (_[_])
open import Lattice

module Analysis.Forward.Lattices
    (L : Set) {h}
    {_≈ˡ_ : L → L → Set} {_⊔ˡ_ : L → L → L} {_⊓ˡ_ : L → L → L}
    {{isFiniteHeightLatticeˡ : IsFiniteHeightLattice L h _≈ˡ_ _⊔ˡ_ _⊓ˡ_}}
    {{≈ˡ-Decidable : IsDecidable _≈ˡ_}}
    (prog : Program) where

open import Data.String using (String) renaming (_≟_ to _≟ˢ_)
open import Data.Product using (proj₁; proj₂; _,_)
open import Data.Sum using (inj₁; inj₂)
open import Data.List using (List; _∷_; []; foldr)
open import Data.List.Membership.Propositional using () renaming (_∈_ to _∈ˡ_)
open import Data.List.Relation.Unary.Any as Any using ()
open import Relation.Binary.PropositionalEquality using (_≡_; refl)
open import Utils using (Pairwise; _⇒_; _∨_; it)

open IsFiniteHeightLattice isFiniteHeightLatticeˡ
    using ()
    renaming
        ( isLattice to isLatticeˡ
        ; fixedHeight to fixedHeightˡ
        ; ≈-sym to ≈ˡ-sym
        )
open Program prog

import Lattice.FiniteMap
import Chain

instance
    ≡-Decidable-String = record { R-dec = _≟ˢ_ }
    ≡-Decidable-State = record { R-dec = _≟_ }

-- The variable -> abstract value (e.g. sign) map is a finite value-map
-- with keys strings. Use a bundle to avoid explicitly specifying operators.
-- It's helpful to export these via 'public' since consumers tend to
-- use various variable lattice operations.
module VariableValuesFiniteMap = Lattice.FiniteMap String L vars
open VariableValuesFiniteMap
    using ()
    renaming
        ( FiniteMap to VariableValues
        ; isLattice to isLatticeᵛ
        ; _≈_ to _≈ᵛ_
        ; _⊔_ to _⊔ᵛ_
        ; _≼_ to _≼ᵛ_
        ; ≈-Decidable to ≈ᵛ-Decidable
        ; _∈_ to _∈ᵛ_
        ; _∈k_ to _∈kᵛ_
        ; _updating_via_ to _updatingᵛ_via_
        ; locate to locateᵛ
        ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵛ≼m₂[k]ᵛ
        ; ∈k-dec to ∈k-decᵛ
        ; all-equal-keys to all-equal-keysᵛ
        ; Provenance-union to Provenance-unionᵛ
        )
    public
open IsLattice isLatticeᵛ
    using ()
    renaming
        ( ⊔-Monotonicˡ to ⊔ᵛ-Monotonicˡ
        ; ⊔-Monotonicʳ to ⊔ᵛ-Monotonicʳ
        ; ⊔-idemp to ⊔ᵛ-idemp
        )
    public
open VariableValuesFiniteMap.FixedHeight vars-Unique
    using ()
    renaming
        ( isFiniteHeightLattice to isFiniteHeightLatticeᵛ
        ; fixedHeight to fixedHeightᵛ
        ; ⊥-contains-bottoms to ⊥ᵛ-contains-bottoms
        )
    public

⊥ᵛ = Chain.Height.⊥ fixedHeightᵛ

-- Finally, the map we care about is (state -> (variables -> value)). Bring that in.
module StateVariablesFiniteMap = Lattice.FiniteMap State VariableValues states
open StateVariablesFiniteMap
    using (_[_]; []-∈; m₁≼m₂⇒m₁[ks]≼m₂[ks]; m₁≈m₂⇒k∈m₁⇒k∈km₂⇒v₁≈v₂)
    renaming
        ( FiniteMap to StateVariables
        ; isLattice to isLatticeᵐ
        ; _≈_ to _≈ᵐ_
        ; _∈_ to _∈ᵐ_
        ; _∈k_ to _∈kᵐ_
        ; locate to locateᵐ
        ; _≼_ to _≼ᵐ_
        ; ≈-Decidable to ≈ᵐ-Decidable
        ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ
        ; ≈-sym to ≈ᵐ-sym
        )
    public
open StateVariablesFiniteMap.FixedHeight states-Unique
    using ()
    renaming
        ( isFiniteHeightLattice to isFiniteHeightLatticeᵐ
        )
    public

-- We now have our (state -> (variables -> value)) map.
-- Define a couple of helpers to retrieve values from it. Specifically,
-- since the State type is as specific as possible, it's always possible to
-- retrieve the variable values at each state.

states-in-Map : ∀ (s : State) (sv : StateVariables) → s ∈kᵐ sv
states-in-Map s sv@(m , ksv≡states) rewrite ksv≡states = states-complete s

variablesAt : State → StateVariables → VariableValues
variablesAt s sv = proj₁ (locateᵐ {s} {sv} (states-in-Map s sv))

variablesAt-∈ : ∀ (s : State) (sv : StateVariables) → (s , variablesAt s sv) ∈ᵐ sv
variablesAt-∈ s sv = proj₂ (locateᵐ {s} {sv} (states-in-Map s sv))

variablesAt-≈ : ∀ s sv₁ sv₂ → sv₁ ≈ᵐ sv₂ → variablesAt s sv₁ ≈ᵛ variablesAt s sv₂
variablesAt-≈ s sv₁ sv₂ sv₁≈sv₂ =
    m₁≈m₂⇒k∈m₁⇒k∈km₂⇒v₁≈v₂ sv₁ sv₂ sv₁≈sv₂
                           (states-in-Map s sv₁) (states-in-Map s sv₂)

-- build up the 'join' function, which follows from Exercise 4.26's
--
--    L₁ → (A → L₂)
--
-- Construction, with L₁ = (A → L₂), and f = id

joinForKey : State → StateVariables → VariableValues
joinForKey k states = foldr _⊔ᵛ_ ⊥ᵛ (states [ incoming k ])

-- The per-key join is made up of map key accesses (which are monotonic)
-- and folds using the join operation (also monotonic)

joinForKey-Mono : ∀ (k : State) → Monotonic _≼ᵐ_ _≼ᵛ_ (joinForKey k)
joinForKey-Mono k {fm₁} {fm₂} fm₁≼fm₂ =
    foldr-Mono it it (fm₁ [ incoming k ]) (fm₂ [ incoming k ]) _⊔ᵛ_ ⊥ᵛ ⊥ᵛ
               (m₁≼m₂⇒m₁[ks]≼m₂[ks] fm₁ fm₂ (incoming k) fm₁≼fm₂)
               (⊔ᵛ-idemp ⊥ᵛ) ⊔ᵛ-Monotonicʳ ⊔ᵛ-Monotonicˡ

-- The name f' comes from the formulation of Exercise 4.26.
open StateVariablesFiniteMap.GeneralizedUpdate {{isLatticeᵐ}} (λ x → x) (λ a₁≼a₂ → a₁≼a₂) joinForKey joinForKey-Mono states
    using ()
    renaming
        ( f' to joinAll
        ; f'-Monotonic to joinAll-Mono
        ; f'-k∈ks-≡ to joinAll-k∈ks-≡
        )
    public

variablesAt-joinAll : ∀ (s : State) (sv : StateVariables) →
                      variablesAt s (joinAll sv) ≡ joinForKey s sv
variablesAt-joinAll s sv
    with (vs , s,vs∈usv) ← locateᵐ {s} {joinAll sv} (states-in-Map s (joinAll sv)) =
    joinAll-k∈ks-≡ {l = sv} (states-complete s) s,vs∈usv

-- Elements of the lattice type L describe individual variables. What
-- exactly each lattice element says about the variable is defined
-- by a LatticeInterpretation element. We've now constructed the
-- (Variable → L) lattice, which describes states, and we need to lift
-- the "meaning" of the element lattice to a descriptions of states.
module _ {{latticeInterpretationˡ : LatticeInterpretation isLatticeˡ}} where
    open LatticeInterpretation latticeInterpretationˡ
        using ()
        renaming
            ( ⟦_⟧ to ⟦_⟧ˡ
            ; ⟦⟧-respects-≈ to ⟦⟧ˡ-respects-≈ˡ
            ; ⟦⟧-⊔-∨ to ⟦⟧ˡ-⊔ˡ-∨
            )
        public

    ⟦_⟧ᵛ : VariableValues → Env → Set
    ⟦_⟧ᵛ vs ρ = ∀ {k l} → (k , l) ∈ᵛ vs → ∀ {v} → (k , v) Language.∈ ρ → ⟦ l ⟧ˡ v

    ⟦⊥ᵛ⟧ᵛ∅ : ⟦ ⊥ᵛ ⟧ᵛ []
    ⟦⊥ᵛ⟧ᵛ∅ _ ()

    ⟦⟧ᵛ-respects-≈ᵛ : ∀ {vs₁ vs₂ : VariableValues} → vs₁ ≈ᵛ vs₂ → ⟦ vs₁ ⟧ᵛ ⇒ ⟦ vs₂ ⟧ᵛ
    ⟦⟧ᵛ-respects-≈ᵛ {m₁ , _} {m₂ , _}
                    (m₁⊆m₂ , m₂⊆m₁) ρ ⟦vs₁⟧ρ {k} {l} k,l∈m₂ {v} k,v∈ρ =
        let
            (l' , (l≈l' , k,l'∈m₁)) = m₂⊆m₁ _ _ k,l∈m₂
            ⟦l'⟧v = ⟦vs₁⟧ρ k,l'∈m₁ k,v∈ρ
        in
            ⟦⟧ˡ-respects-≈ˡ (≈ˡ-sym l≈l') v ⟦l'⟧v

    ⟦⟧ᵛ-⊔ᵛ-∨ : ∀ {vs₁ vs₂ : VariableValues} → (⟦ vs₁ ⟧ᵛ ∨ ⟦ vs₂ ⟧ᵛ) ⇒ ⟦ vs₁ ⊔ᵛ vs₂ ⟧ᵛ
    ⟦⟧ᵛ-⊔ᵛ-∨ {vs₁} {vs₂} ρ ⟦vs₁⟧ρ∨⟦vs₂⟧ρ {k} {l} k,l∈vs₁₂ {v} k,v∈ρ
        with ((l₁ , l₂) , (refl , (k,l₁∈vs₁ , k,l₂∈vs₂)))
             ← Provenance-unionᵛ vs₁ vs₂ k,l∈vs₁₂
        with ⟦vs₁⟧ρ∨⟦vs₂⟧ρ
    ...   | inj₁ ⟦vs₁⟧ρ = ⟦⟧ˡ-⊔ˡ-∨ {l₁} {l₂} v (inj₁ (⟦vs₁⟧ρ k,l₁∈vs₁ k,v∈ρ))
    ...   | inj₂ ⟦vs₂⟧ρ = ⟦⟧ˡ-⊔ˡ-∨ {l₁} {l₂} v (inj₂ (⟦vs₂⟧ρ k,l₂∈vs₂ k,v∈ρ))

    ⟦⟧ᵛ-foldr : ∀ {vs : VariableValues} {vss : List VariableValues} {ρ : Env} →
                ⟦ vs ⟧ᵛ ρ → vs ∈ˡ vss → ⟦ foldr _⊔ᵛ_ ⊥ᵛ vss ⟧ᵛ ρ
    ⟦⟧ᵛ-foldr {vs} {vs ∷ vss'} {ρ = ρ} ⟦vs⟧ρ (Any.here refl) =
        ⟦⟧ᵛ-⊔ᵛ-∨ {vs₁ = vs} {vs₂ = foldr _⊔ᵛ_ ⊥ᵛ vss'} ρ (inj₁ ⟦vs⟧ρ)
    ⟦⟧ᵛ-foldr {vs} {vs' ∷ vss'} {ρ = ρ} ⟦vs⟧ρ (Any.there vs∈vss') =
        ⟦⟧ᵛ-⊔ᵛ-∨ {vs₁ = vs'} {vs₂ = foldr _⊔ᵛ_ ⊥ᵛ vss'} ρ
                 (inj₂ (⟦⟧ᵛ-foldr ⟦vs⟧ρ vs∈vss'))