module Language.Graphs where

open import Language.Base
open import Language.Semantics

open import Data.Fin as Fin using (Fin; suc; zero; _↑ˡ_; _↑ʳ_)
open import Data.Fin.Properties as FinProp using (suc-injective)
open import Data.List as List using (List; []; _∷_)
open import Data.List.Membership.Propositional as ListMem using ()
open import Data.List.Membership.Propositional.Properties as ListMemProp using ()
open import Data.Nat as Nat using (ℕ; suc)
open import Data.Nat.Properties using (+-assoc; +-comm)
open import Data.Product using (_×_; Σ; _,_)
open import Data.Vec using (Vec; []; _∷_; lookup; cast; _++_)
open import Data.Vec.Properties using (cast-is-id; ++-assoc; lookup-++ˡ; cast-sym; ++-identityʳ)
open import Relation.Binary.PropositionalEquality as Eq using (_≡_; sym; refl; subst)

open import Lattice
open import Utils using (x∈xs⇒fx∈fxs; _⊗_; _,_)

record Graph : Set where
    constructor MkGraph
    field
        size : ℕ

    Index : Set
    Index = Fin size

    Edge : Set
    Edge = Index × Index

    field
        nodes : Vec (List BasicStmt) size
        edges : List Edge

empty : Graph
empty = record
    { size = 0
    ; nodes = []
    ; edges = []
    }

↑ˡ-Edge : ∀ {n} → (Fin n × Fin n) → ∀ m → (Fin (n Nat.+ m) × Fin (n Nat.+ m))
↑ˡ-Edge (idx₁ , idx₂) m = (idx₁ ↑ˡ m , idx₂ ↑ˡ m)

_[_] : ∀ (g : Graph) → Graph.Index g → List BasicStmt
_[_] g idx = lookup (Graph.nodes g) idx

record _⊆_ (g₁ g₂ : Graph) : Set where
    constructor Mk-⊆
    field
        n : ℕ
        sg₂≡sg₁+n : Graph.size g₂ ≡ Graph.size g₁ Nat.+ n
        newNodes : Vec (List BasicStmt) n
        nsg₂≡nsg₁++newNodes : cast sg₂≡sg₁+n (Graph.nodes g₂) ≡ Graph.nodes g₁ ++ newNodes
        e∈g₁⇒e∈g₂ : ∀ {e : Graph.Edge g₁} →
                    e ListMem.∈ (Graph.edges g₁) →
                    (↑ˡ-Edge e n) ListMem.∈ (subst (λ m → List (Fin m × Fin m)) sg₂≡sg₁+n (Graph.edges g₂))

castᵉ : ∀ {n m : ℕ} .(p : n ≡ m) → (Fin n × Fin n) → (Fin m × Fin m)
castᵉ p (idx₁ , idx₂) = (Fin.cast p idx₁ , Fin.cast p idx₂)

↑ˡ-assoc : ∀ {s n₁ n₂} (f : Fin s) (p : s Nat.+ (n₁ Nat.+ n₂) ≡ s Nat.+ n₁ Nat.+ n₂) →
           f ↑ˡ n₁ ↑ˡ n₂ ≡ Fin.cast p (f ↑ˡ (n₁ Nat.+ n₂))
↑ˡ-assoc zero p = refl
↑ˡ-assoc {suc s'} {n₁} {n₂} (suc f') p rewrite ↑ˡ-assoc f' (sym (+-assoc s' n₁ n₂)) = refl

↑ˡ-Edge-assoc : ∀ {s n₁ n₂} (e : Fin s × Fin s) (p : s Nat.+ (n₁ Nat.+ n₂) ≡ s Nat.+ n₁ Nat.+ n₂) →
                ↑ˡ-Edge (↑ˡ-Edge e n₁) n₂ ≡ castᵉ p (↑ˡ-Edge e (n₁ Nat.+ n₂))
↑ˡ-Edge-assoc (idx₁ , idx₂) p
    rewrite ↑ˡ-assoc idx₁ p
    rewrite ↑ˡ-assoc idx₂ p = refl

↑ˡ-identityʳ : ∀ {s} (f : Fin s) (p : s Nat.+ 0 ≡ s) →
           f ≡ Fin.cast p (f ↑ˡ 0)
↑ˡ-identityʳ zero p = refl
↑ˡ-identityʳ {suc s'} (suc f') p rewrite sym (↑ˡ-identityʳ f' (+-comm s' 0)) = refl

↑ˡ-Edge-identityʳ : ∀ {s} (e : Fin s × Fin s) (p : s Nat.+ 0 ≡ s) →
                    e ≡ castᵉ p (↑ˡ-Edge e 0)
↑ˡ-Edge-identityʳ (idx₁ , idx₂) p
    rewrite sym (↑ˡ-identityʳ idx₁ p)
    rewrite sym (↑ˡ-identityʳ idx₂ p) = refl

cast∈⇒∈subst : ∀ {n m : ℕ} (p : n ≡ m) (q : m ≡ n)
               (e : Fin n × Fin n) (es : List (Fin m × Fin m)) →
               castᵉ p e ListMem.∈ es →
               e ListMem.∈ subst (λ m → List (Fin m × Fin m)) q es
cast∈⇒∈subst refl refl (idx₁ , idx₂) es e∈es
    rewrite FinProp.cast-is-id refl idx₁
    rewrite FinProp.cast-is-id refl idx₂ = e∈es

⊆-trans : ∀ {g₁ g₂ g₃ : Graph} → g₁ ⊆ g₂ → g₂ ⊆ g₃ → g₁ ⊆ g₃
⊆-trans {MkGraph s₁ ns₁ es₁} {MkGraph s₂ ns₂ es₂} {MkGraph s₃ ns₃ es₃}
        (Mk-⊆ n₁ p₁@refl newNodes₁ nsg₂≡nsg₁++newNodes₁ e∈g₁⇒e∈g₂)
        (Mk-⊆ n₂ p₂@refl newNodes₂ nsg₃≡nsg₂++newNodes₂ e∈g₂⇒e∈g₃)
    rewrite cast-is-id refl ns₂
    rewrite cast-is-id refl ns₃
    with refl ← nsg₂≡nsg₁++newNodes₁
    with refl ← nsg₃≡nsg₂++newNodes₂ =
    record
        { n = n₁ Nat.+ n₂
        ; sg₂≡sg₁+n = +-assoc s₁ n₁ n₂
        ; newNodes = newNodes₁ ++ newNodes₂
        ; nsg₂≡nsg₁++newNodes = ++-assoc (+-assoc s₁ n₁ n₂) ns₁ newNodes₁ newNodes₂
        ; e∈g₁⇒e∈g₂ = λ {e} e∈g₁ →
            cast∈⇒∈subst (sym (+-assoc s₁ n₁ n₂)) (+-assoc s₁ n₁ n₂) _ _
                         (subst (λ e' → e' ListMem.∈ es₃)
                                (↑ˡ-Edge-assoc e (sym (+-assoc s₁ n₁ n₂)))
                                (e∈g₂⇒e∈g₃ (e∈g₁⇒e∈g₂ e∈g₁)))
        }

record Relaxable (T : Graph → Set) : Set where
    field relax : ∀ {g₁ g₂ : Graph} → g₁ ⊆ g₂ → T g₁ → T g₂

instance
    IndexRelaxable : Relaxable Graph.Index
    IndexRelaxable = record
        { relax = λ { (Mk-⊆ n refl _ _ _) idx → idx ↑ˡ n }
        }

    EdgeRelaxable : Relaxable Graph.Edge
    EdgeRelaxable = record
        { relax = λ g₁⊆g₂ (idx₁ , idx₂) →
            ( Relaxable.relax IndexRelaxable g₁⊆g₂ idx₁
            , Relaxable.relax IndexRelaxable g₁⊆g₂ idx₂
            )
        }

    ProdRelaxable : ∀ {P : Graph → Set} {Q : Graph → Set} →
                    {{ PRelaxable : Relaxable P }} → {{ QRelaxable : Relaxable Q }} →
                    Relaxable (P ⊗ Q)
    ProdRelaxable {{pr}} {{qr}} = record
        { relax = (λ { g₁⊆g₂ (p , q) →
            ( Relaxable.relax pr g₁⊆g₂ p
            , Relaxable.relax qr g₁⊆g₂ q) }
            )
        }

open Relaxable {{...}} public

relax-preserves-[]≡ : ∀ (g₁ g₂ : Graph) (g₁⊆g₂ : g₁ ⊆ g₂) (idx : Graph.Index g₁) →
                      g₁ [ idx ] ≡ g₂ [ relax g₁⊆g₂ idx ]
relax-preserves-[]≡ g₁ g₂ (Mk-⊆ n refl newNodes nsg₂≡nsg₁++newNodes _) idx
    rewrite cast-is-id refl (Graph.nodes g₂)
    with refl ← nsg₂≡nsg₁++newNodes = sym (lookup-++ˡ (Graph.nodes g₁) _ _)

-- Tools for graph construction. The most important is a 'monotonic function':
-- one that takes a graph, and produces another graph, such that the
-- new graph includes all the information from the old one.

MonotonicGraphFunction : (Graph → Set) → Set
MonotonicGraphFunction T = (g₁ : Graph) → Σ Graph (λ g₂ → T g₂ × g₁ ⊆ g₂)

-- Now, define some operations on monotonic functions; these are useful
-- to save the work of threading intermediate graphs in and out of operations.

infixr 2 _⟨⊗⟩_
_⟨⊗⟩_ : ∀ {T₁ T₂ : Graph → Set} {{ T₁Relaxable : Relaxable T₁ }} →
      MonotonicGraphFunction T₁ → MonotonicGraphFunction T₂ →
      MonotonicGraphFunction (T₁ ⊗ T₂)
_⟨⊗⟩_ {{r}} f₁ f₂ g
    with (g' , (t₁ , g⊆g')) ← f₁ g
    with (g'' , (t₂ , g'⊆g'')) ← f₂ g' =
    (g'' , ((Relaxable.relax r g'⊆g'' t₁ , t₂) , ⊆-trans g⊆g' g'⊆g''))

infixl 2 _update_
_update_ : ∀ {T : Graph → Set} {{ TRelaxable : Relaxable T }} →
         MonotonicGraphFunction T → (∀ (g : Graph) → T g → Σ Graph (λ g' → g ⊆ g')) →
         MonotonicGraphFunction T
_update_ {{r}} f mod g
    with (g' , (t , g⊆g')) ← f g
    with (g'' , g'⊆g'') ← mod g' t =
    (g'' , ((Relaxable.relax r g'⊆g'' t , ⊆-trans g⊆g' g'⊆g'')))

infixl 2 _map_
_map_ : ∀ {T₁ T₂ : Graph → Set} →
      MonotonicGraphFunction T₁ → (∀ (g : Graph) → T₁ g → T₂ g) →
      MonotonicGraphFunction T₂
_map_ f fn g = let (g' , (t₁ , g⊆g')) = f g in (g' , (fn g' t₁ , g⊆g'))


-- To reason about monotonic functions and what we do, we need a way
-- to describe values they produce. A 'graph-value predicate' is
-- just a predicate for some (dependent) value.

GraphValuePredicate : (Graph → Set) → Set₁
GraphValuePredicate T = ∀ (g : Graph) → T g → Set

Both : {T₁ T₂ : Graph → Set} → GraphValuePredicate T₁ → GraphValuePredicate T₂ →
       GraphValuePredicate (T₁ ⊗ T₂)
Both P Q = (λ { g (t₁ , t₂) → (P g t₁ × Q g t₂) })

-- Since monotnic functions keep adding on to a function, proofs of
-- graph-value predicates go stale fast (they describe old values of
-- the graph). To keep propagating them through, we need them to still
-- on 'bigger graphs'. We call such predicates monotonic as well, since
-- they respect the ordering of graphs.

MonotonicPredicate : ∀ {T : Graph → Set} {{ TRelaxable : Relaxable T }} →
                     GraphValuePredicate T → Set
MonotonicPredicate {T} P = ∀ (g₁ g₂ : Graph) (t₁ : T g₁) (g₁⊆g₂ : g₁ ⊆ g₂) →
                           P g₁ t₁ → P g₂ (relax g₁⊆g₂ t₁)


-- A 'map' has a certain property if its ouputs satisfy that property
-- for all inputs.

always : ∀ {T : Graph → Set} → GraphValuePredicate T → MonotonicGraphFunction T → Set
always P m = ∀ g₁ → let (g₂ , t , _) = m g₁ in P g₂ t

⟨⊗⟩-reason : ∀ {T₁ T₂ : Graph → Set} {{ T₁Relaxable : Relaxable T₁ }}
             {P : GraphValuePredicate T₁} {Q : GraphValuePredicate T₂}
             {P-Mono : MonotonicPredicate P}
             {m₁ : MonotonicGraphFunction T₁} {m₂ : MonotonicGraphFunction T₂} →
             always P m₁ → always Q m₂ → always (Both P Q) (m₁ ⟨⊗⟩ m₂)
⟨⊗⟩-reason {P-Mono = P-Mono} {m₁ = m₁} {m₂ = m₂} aP aQ g
    with p ← aP g
    with (g' , (t₁ , g⊆g')) ← m₁ g
    with q ← aQ g'
    with (g'' , (t₂ , g'⊆g'')) ← m₂ g' = (P-Mono _ _ _ g'⊆g'' p , q)

module Construction where
    pushBasicBlock : List BasicStmt → MonotonicGraphFunction Graph.Index
    pushBasicBlock bss g =
        ( record
            { size = Graph.size g Nat.+ 1
            ; nodes = Graph.nodes g ++ (bss ∷ [])
            ; edges = List.map (λ e → ↑ˡ-Edge e 1) (Graph.edges g)
            }
        , ( Graph.size g ↑ʳ zero
          , record
              { n = 1
              ; sg₂≡sg₁+n = refl
              ; newNodes = (bss ∷ [])
              ; nsg₂≡nsg₁++newNodes = cast-is-id refl _
              ; e∈g₁⇒e∈g₂ = λ e∈g₁ → x∈xs⇒fx∈fxs (λ e → ↑ˡ-Edge e 1) e∈g₁
              }
          )
        )

    addEdges : ∀ (g : Graph) → List (Graph.Edge g) → Σ Graph (λ g' → g ⊆ g')
    addEdges (MkGraph s ns es) es' =
        ( record
            { size = s
            ; nodes = ns
            ; edges = es' List.++ es
            }
        , record
            { n = 0
            ; sg₂≡sg₁+n = +-comm 0 s
            ; newNodes = []
            ; nsg₂≡nsg₁++newNodes = cast-sym _ (++-identityʳ (+-comm s 0) ns)
            ; e∈g₁⇒e∈g₂ = λ {e} e∈es →
                cast∈⇒∈subst (+-comm s 0) (+-comm 0 s) _ _
                             (subst (λ e' → e' ListMem.∈ _)
                                    (↑ˡ-Edge-identityʳ e (+-comm s 0))
                                    (ListMemProp.∈-++⁺ʳ es' e∈es))
            }
        )

    pushEmptyBlock : MonotonicGraphFunction Graph.Index
    pushEmptyBlock = pushBasicBlock []

    buildCfg : Stmt → MonotonicGraphFunction (Graph.Index ⊗ Graph.Index)
    buildCfg ⟨ bs₁ ⟩ = pushBasicBlock (bs₁ ∷ []) map (λ g idx → (idx , idx))
    buildCfg (s₁ then s₂) =
        (buildCfg s₁ ⟨⊗⟩ buildCfg s₂)
        update (λ { g ((idx₁ , idx₂) , (idx₃ , idx₄)) → addEdges g ((idx₂ , idx₃) ∷ []) })
        map (λ { g ((idx₁ , idx₂) , (idx₃ , idx₄)) → (idx₁ , idx₄) })
    buildCfg (if _ then s₁ else s₂) =
        (buildCfg s₁ ⟨⊗⟩ buildCfg s₂ ⟨⊗⟩ pushEmptyBlock ⟨⊗⟩ pushEmptyBlock)
        update (λ { g ((idx₁ , idx₂) , (idx₃ , idx₄) , idx , idx') →
                  addEdges g ((idx , idx₁) ∷ (idx , idx₃) ∷ (idx₂ , idx') ∷ (idx₄ , idx') ∷ []) })
        map (λ { g ((idx₁ , idx₂) , (idx₃ , idx₄) , idx , idx') → (idx , idx') })
    buildCfg (while _ repeat s) =
        (buildCfg s ⟨⊗⟩ pushEmptyBlock ⟨⊗⟩ pushEmptyBlock)
        update (λ { g ((idx₁ , idx₂) , idx , idx') →
                  addEdges g ((idx , idx') ∷ (idx , idx₁) ∷ (idx₂ , idx) ∷ []) })
        map (λ { g ((idx₁ , idx₂) , idx , idx') → (idx , idx') })