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

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

open import Data.Empty using (⊥-elim)
open import Data.String using (String) renaming (_≟_ to _≟ˢ_)
open import Data.Nat using (suc)
open import Data.Product using (_×_; proj₁; _,_)
open import Data.List using (List; _∷_; []; foldr; cartesianProduct; cartesianProductWith)
open import Data.List.Membership.Propositional as MemProp using () renaming (_∈_ to _∈ˡ_)
open import Data.List.Relation.Unary.Any as Any using ()
open import Relation.Binary.PropositionalEquality using (_≡_; refl; sym; trans; subst)
open import Relation.Nullary using (¬_; Dec; yes; no)
open import Data.Unit using (⊤)
open import Function using (_∘_)

open import Utils using (Pairwise)
import Lattice.FiniteValueMap

open IsFiniteHeightLattice isFiniteHeightLatticeˡ
    using ()
    renaming
        ( isLattice to isLatticeˡ
        ; fixedHeight to fixedHeightˡ
        ; _≼_ to _≼ˡ_
        )

module WithProg (prog : Program) where
    open Program prog

    -- The variable -> abstract value (e.g. sign) map is a finite value-map
    -- with keys strings. Use a bundle to avoid explicitly specifying operators.
    module VariableValuesFiniteMap = Lattice.FiniteValueMap.WithKeys _≟ˢ_ isLatticeˡ vars
    open VariableValuesFiniteMap
        using ()
        renaming
            ( FiniteMap to VariableValues
            ; isLattice to isLatticeᵛ
            ; _≈_ to _≈ᵛ_
            ; _⊔_ to _⊔ᵛ_
            ; _≼_ to _≼ᵛ_
            ; ≈₂-dec⇒≈-dec to ≈ˡ-dec⇒≈ᵛ-dec
            ; _∈_ 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ᵛ
            ; forget to forgetᵛ
            )
        public
    open IsLattice isLatticeᵛ
        using ()
        renaming
            ( ⊔-Monotonicˡ to ⊔ᵛ-Monotonicˡ
            ; ⊔-Monotonicʳ to ⊔ᵛ-Monotonicʳ
            ; ⊔-idemp to ⊔ᵛ-idemp
            )
    open Lattice.FiniteValueMap.IterProdIsomorphism.WithUniqueKeysAndFixedHeight _≟ˢ_ isLatticeˡ vars-Unique ≈ˡ-dec _ fixedHeightˡ
        using ()
        renaming
            ( isFiniteHeightLattice to isFiniteHeightLatticeᵛ
            )

    ≈ᵛ-dec = ≈ˡ-dec⇒≈ᵛ-dec ≈ˡ-dec
    joinSemilatticeᵛ = IsFiniteHeightLattice.joinSemilattice isFiniteHeightLatticeᵛ
    fixedHeightᵛ = IsFiniteHeightLattice.fixedHeight isFiniteHeightLatticeᵛ
    ⊥ᵛ = proj₁ (proj₁ (proj₁ fixedHeightᵛ))

    -- Finally, the map we care about is (state -> (variables -> value)). Bring that in.
    module StateVariablesFiniteMap = Lattice.FiniteValueMap.WithKeys _≟_ isLatticeᵛ states
    open StateVariablesFiniteMap
        using (_[_]; m₁≼m₂⇒m₁[ks]≼m₂[ks])
        renaming
            ( FiniteMap to StateVariables
            ; isLattice to isLatticeᵐ
            ; _∈k_ to _∈kᵐ_
            ; locate to locateᵐ
            ; _≼_ to _≼ᵐ_
            ; ≈₂-dec⇒≈-dec to ≈ᵛ-dec⇒≈ᵐ-dec
            ; m₁≼m₂⇒m₁[k]≼m₂[k] to m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ
            )
        public
    open Lattice.FiniteValueMap.IterProdIsomorphism.WithUniqueKeysAndFixedHeight _≟_ isLatticeᵛ states-Unique ≈ᵛ-dec _ fixedHeightᵛ
        using ()
        renaming
            ( isFiniteHeightLattice to isFiniteHeightLatticeᵐ
            )

    ≈ᵐ-dec = ≈ᵛ-dec⇒≈ᵐ-dec ≈ᵛ-dec
    fixedHeightᵐ = IsFiniteHeightLattice.fixedHeight isFiniteHeightLatticeᵐ

    -- 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 joinSemilatticeᵛ joinSemilatticeᵛ (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 states isLatticeᵐ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) joinForKey joinForKey-Mono states
        renaming
            ( f' to joinAll
            ; f'-Monotonic to joinAll-Mono
            )

    -- With 'join' in hand, we need to perform abstract evaluation.
    module WithEvaluator (eval : Expr → VariableValues → L)
                         (eval-Mono : ∀ (e : Expr) → Monotonic _≼ᵛ_ _≼ˡ_ (eval e)) where

        -- For a particular evaluation function, we need to perform an evaluation
        -- for an assignment, and update the corresponding key. Use Exercise 4.26's
        -- generalized update to set the single key's value.

        private module _ (k : String) (e : Expr) where
            open VariableValuesFiniteMap.GeneralizedUpdate vars isLatticeᵛ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) (λ _ → eval e) (λ _ {vs₁} {vs₂} vs₁≼vs₂ → eval-Mono e {vs₁} {vs₂} vs₁≼vs₂) (k ∷ [])
                renaming
                    ( f' to updateVariablesFromExpression
                    ; f'-Monotonic to updateVariablesFromExpression-Mono
                    ; f'-k∈ks-≡ to updateVariablesFromExpression-k∈ks-≡
                    ; f'-k∉ks-backward to updateVariablesFromExpression-k∉ks-backward
                    )
                public

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

        -- The per-state update function makes use of the single-key setter,
        -- updateVariablesFromExpression, for the case where the statement
        -- is an assignment.
        --
        -- This per-state function adjusts the variables in that state,
        -- also monotonically; we derive the for-each-state update from
        -- the Exercise 4.26 again.

        updateVariablesFromStmt : BasicStmt → VariableValues → VariableValues
        updateVariablesFromStmt (k ← e) vs = updateVariablesFromExpression k e vs
        updateVariablesFromStmt noop vs = vs

        updateVariablesFromStmt-Monoʳ : ∀ (bs : BasicStmt) → Monotonic _≼ᵛ_ _≼ᵛ_ (updateVariablesFromStmt bs)
        updateVariablesFromStmt-Monoʳ (k ← e) {vs₁} {vs₂} vs₁≼vs₂ = updateVariablesFromExpression-Mono k e {vs₁} {vs₂} vs₁≼vs₂
        updateVariablesFromStmt-Monoʳ noop vs₁≼vs₂ = vs₁≼vs₂

        updateVariablesForState : State → StateVariables → VariableValues
        updateVariablesForState s sv =
            let
                bss = code s
                (vs , s,vs∈sv) = locateᵐ {s} {sv} (states-in-Map s sv)
            in
                foldr updateVariablesFromStmt vs bss

        updateVariablesForState-Monoʳ : ∀ (s : State) → Monotonic _≼ᵐ_ _≼ᵛ_ (updateVariablesForState s)
        updateVariablesForState-Monoʳ s {sv₁} {sv₂} sv₁≼sv₂ =
            let
                bss = code s
                (vs₁ , s,vs₁∈sv₁) = locateᵐ {s} {sv₁} (states-in-Map s sv₁)
                (vs₂ , s,vs₂∈sv₂) = locateᵐ {s} {sv₂} (states-in-Map s sv₂)
                vs₁≼vs₂ = m₁≼m₂⇒m₁[k]ᵐ≼m₂[k]ᵐ sv₁ sv₂ sv₁≼sv₂ s,vs₁∈sv₁ s,vs₂∈sv₂
            in
                foldr-Mono' (IsLattice.joinSemilattice isLatticeᵛ) bss
                            updateVariablesFromStmt updateVariablesFromStmt-Monoʳ
                            vs₁≼vs₂

        open StateVariablesFiniteMap.GeneralizedUpdate states isLatticeᵐ (λ x → x) (λ a₁≼a₂ → a₁≼a₂) updateVariablesForState updateVariablesForState-Monoʳ states
            renaming
                ( f' to updateAll
                ; f'-Monotonic to updateAll-Mono
                )

        -- Finally, the whole analysis consists of getting the 'join'
        -- of all incoming states, then applying the per-state evaluation
        -- function. This is just a composition, and is trivially monotonic.

        analyze : StateVariables → StateVariables
        analyze = updateAll ∘ joinAll

        analyze-Mono : Monotonic _≼ᵐ_ _≼ᵐ_ analyze
        analyze-Mono {sv₁} {sv₂} sv₁≼sv₂ =
            updateAll-Mono {joinAll sv₁} {joinAll sv₂}
                           (joinAll-Mono {sv₁} {sv₂} sv₁≼sv₂)

        -- The fixed point of the 'analyze' function is our final goal.
        open import Fixedpoint ≈ᵐ-dec isFiniteHeightLatticeᵐ analyze (λ {m₁} {m₂} m₁≼m₂ → analyze-Mono {m₁} {m₂} m₁≼m₂)
            using ()
            renaming (aᶠ to result)
            public

        module WithInterpretation (latticeInterpretationˡ : LatticeInterpretation isLatticeˡ) where
            open LatticeInterpretation latticeInterpretationˡ
                using ()
                renaming (⟦_⟧ to ⟦_⟧ˡ)

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

            InterpretationValid : Set
            InterpretationValid = ∀ {vs ρ e v} → ρ , e ⇒ᵉ v → ⟦ vs ⟧ᵛ ρ → ⟦ eval e vs ⟧ˡ v

            module WithValidity (interpretationValidˡ : InterpretationValid) where

                updateVariablesFromStmt-matches : ∀ {bs vs ρ₁ ρ₂} → ρ₁ , bs ⇒ᵇ ρ₂ → ⟦ vs ⟧ᵛ ρ₁ → ⟦ updateVariablesFromStmt bs vs ⟧ᵛ ρ₂
                updateVariablesFromStmt-matches {_} {vs} {ρ₁} {ρ₁} (⇒ᵇ-noop ρ₁) ⟦vs⟧ρ₁ = ⟦vs⟧ρ₁
                updateVariablesFromStmt-matches {_} {vs} {ρ₁} {_} (⇒ᵇ-← ρ₁ k e v ρ,e⇒v) ⟦vs⟧ρ₁ {k'} {l} k',l∈vs' {v'} k',v'∈ρ₂
                    with k ≟ˢ k' | k',v'∈ρ₂
                ...   | yes refl | here _ v _
                        rewrite updateVariablesFromExpression-k∈ks-≡ k e {l = vs} (Any.here refl) k',l∈vs' =
                        interpretationValidˡ ρ,e⇒v ⟦vs⟧ρ₁
                ...   | yes k≡k' | there _ _ _ _ _ k'≢k _ = ⊥-elim (k'≢k (sym k≡k'))
                ...   | no k≢k' | here _ _ _ = ⊥-elim (k≢k' refl)
                ...   | no k≢k'  | there _ _ _ _ _ _ k',v'∈ρ₁ =
                        let
                            k'∉[k] = (λ { (Any.here refl) → k≢k' refl })
                            k',l∈vs = updateVariablesFromExpression-k∉ks-backward k e {l = vs} k'∉[k] k',l∈vs'
                        in
                            ⟦vs⟧ρ₁ k',l∈vs k',v'∈ρ₁