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₁)))
        }

open import MonotonicState _⊆_ ⊆-trans renaming (MonotonicState to MonotonicGraphFunction)

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₂
            )
        }

open Relaxable {{...}}

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₁) _ _)

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') })