Data.Container.Indexed.Combinator

------------------------------------------------------------------------
-- The Agda standard library
--
-- Indexed container combinators
------------------------------------------------------------------------

{-# OPTIONS --cubical-compatible --safe #-}

module Data.Container.Indexed.Combinator where

open import Axiom.Extensionality.Propositional using (Extensionality)
open import Data.Container.Indexed
open import Data.Empty.Polymorphic using (⊥; ⊥-elim)
open import Data.Unit.Polymorphic.Base using (⊤)
open import Data.Product.Base as Prod hiding (Σ) renaming (_×_ to _⟨×⟩_)
open import Data.Sum.Base renaming (_⊎_ to _⟨⊎⟩_)
open import Data.Sum.Relation.Unary.All as All using (All)
open import Function.Base as F hiding (id; const) renaming (_∘_ to _⟨∘⟩_)
open import Function.Bundles using (mk↔ₛ′)
open import Function.Indexed.Bundles using (_↔ᵢ_)
open import Level
open import Relation.Unary using (Pred; _⊆_; _∪_; _∩_; ⋃; ⋂)
  renaming (_⟨×⟩_ to _⟪×⟫_; _⟨⊙⟩_ to _⟪⊙⟫_; _⟨⊎⟩_ to _⟪⊎⟫_)
open import Relation.Binary.PropositionalEquality as P
  using (_≗_; refl)

private
  variable
    ℓ ℓ₁ ℓ₂ i j k o c c₁ c₂ r r₁ r₂ x z : Level
    I J K O X Z : Set _

------------------------------------------------------------------------
-- Combinators


-- Identity.

id : Container O O c r
id = F.const ⊤ ◃ F.const ⊤ / (λ {o} _ _ → o)

-- Constant.

const : Pred O c → Container I O c r
const X = X ◃ F.const ⊥ / F.const ⊥-elim

-- Composition.

infixr 9 _∘_

_∘_ : Container J K c₁ r₁ → Container I J c₂ r₂ → Container I K _ _
C₁ ∘ C₂ = C ◃ R / n
  where
  C : ∀ k → Set _
  C = ⟦ C₁ ⟧ (Command C₂)

  R : ∀ {k} → ⟦ C₁ ⟧ (Command C₂) k → Set _
  R (c , k) = ◇ C₁ {X = Command C₂} (Response C₂ ⟨∘⟩ proj₂) (_ , c , k)

  n : ∀ {k} (c : ⟦ C₁ ⟧ (Command C₂) k) → R c → _
  n (_ , f) (r₁ , r₂) = next C₂ (f r₁) r₂

-- Duality.

_^⊥ : Container I O c r → Container I O (c ⊔ r) c
(C ^⊥) .Command  o     = (c : C .Command o) → C .Response c
(C ^⊥) .Response {o} _ = C .Command o
(C ^⊥) .next     f c   = C .next c (f c)

-- Strength.

infixl 3 _⋊_

_⋊_ : Container I O c r → (Z : Set z) → Container (I ⟨×⟩ Z) (O ⟨×⟩ Z) c r
(C ⋊ Z) .Command  (o , z)     = C .Command o
(C ⋊ Z) .Response             = C .Response
(C ⋊ Z) .next     {o , z} c r = C .next c r , z

infixr 3 _⋉_

_⋉_ : (Z : Set z) → Container I O c r → Container (Z ⟨×⟩ I) (Z ⟨×⟩ O) c r
(Z ⋉ C) .Command  (z , o)     = C .Command o
(Z ⋉ C) .Response             = C .Response
(Z ⋉ C) .next     {z , o} c r = z , C .next c r



-- Product. (Note that, up to isomorphism, and ignoring universe level
-- issues, this is a special case of indexed product.)

infixr 2 _×_

_×_ : Container I O c₁ r₁ → Container I O c₂ r₂ → Container I O _ _
(C₁ ◃ R₁ / n₁) × (C₂ ◃ R₂ / n₂) = record
  { Command  = C₁ ∩ C₂
  ; Response = R₁ ⟪⊙⟫ R₂
  ; next     = λ { (c₁ , c₂) → [ n₁ c₁ , n₂ c₂ ] }
  }

-- Indexed product.

Π : (X → Container I O c r) → Container I O _ _
Π {X = X} C = record
  { Command  = ⋂ X (Command ⟨∘⟩ C)
  ; Response = ⋃[ x ∶ X ] λ c → Response (C x) (c x)
  ; next     = λ { c (x , r) → next (C x) (c x) r }
  }

-- Sum. (Note that, up to isomorphism, and ignoring universe level
-- issues, this is a special case of indexed sum.)

infixr 1  _⊎_ _⊎′_

_⊎_ : Container I O c₁ r₁ → Container I O c₂ r₂ → Container I O _ _
(C₁ ⊎ C₂) .Command  = C₁ .Command ∪ C₂ .Command
(C₁ ⊎ C₂) .Response = All (C₁ .Response) (C₂ .Response)
(C₁ ⊎ C₂) .next     = All.[ C₁ .next , C₂ .next ]

-- A simplified version for responses at the same level r:

_⊎′_ : Container I O c₁ r → Container I O c₂ r → Container I O _ r
(C₁ ◃ R₁ / n₁) ⊎′ (C₂ ◃ R₂ / n₂) = record
  { Command  = C₁ ∪ C₂
  ; Response = [ R₁ , R₂ ]
  ; next     = [ n₁ , n₂ ]
  }

-- Indexed sum.

Σ : (X → Container I O c r) → Container I O _ r
Σ {X = X} C = record
  { Command  = ⋃ X (Command ⟨∘⟩ C)
  ; Response = λ { (x , c) → Response (C x) c }
  ; next     = λ { (x , c) r → next (C x) c r }
  }

-- Constant exponentiation. (Note that this is a special case of
-- indexed product.)

infix 0 const[_]⟶_

const[_]⟶_ : (X : Set ℓ) → Container I O c r → Container I O _ _
const[ X ]⟶ C = Π {X = X} (F.const C)

------------------------------------------------------------------------
-- Correctness proofs

module Identity where

  correct : {X : Pred O ℓ} → ⟦ id {c = c}{r} ⟧ X ↔ᵢ F.id X
  correct {X = X} = mk↔ₛ′ to from (λ _ → refl) (λ _ → refl)
    where
    to : ∀ {x} → ⟦ id ⟧ X x → F.id X x
    to xs = proj₂ xs _

    from : ∀ {x} → F.id X x → ⟦ id ⟧ X x
    from x = (_ , λ _ → x)

module Constant (ext : ∀ {ℓ} → Extensionality ℓ ℓ) where

  correct : (X : Pred O ℓ₁) {Y : Pred O ℓ₂} → ⟦ const X ⟧ Y ↔ᵢ F.const X Y
  correct X {Y} = mk↔ₛ′ to from (λ _ → refl) to∘from
    where
    to : ⟦ const X ⟧ Y ⊆ X
    to = proj₁

    from : X ⊆ ⟦ const X ⟧ Y
    from = < F.id , F.const ⊥-elim >

    to∘from : _
    to∘from xs = P.cong (proj₁ xs ,_) (ext ⊥-elim)

module Duality where

  correct : (C : Container I O c r) (X : Pred I ℓ) →
            ⟦ C ^⊥ ⟧ X ↔ᵢ (λ o → (c : Command C o) → ∃ λ r → X (next C c r))
  correct C X = mk↔ₛ′ (λ { (f , g) → < f , g > }) (λ f → proj₁ ⟨∘⟩ f , proj₂ ⟨∘⟩ f)
                        (λ _ → refl) (λ _ → refl)

module Composition where

  correct : (C₁ : Container J K c r) (C₂ : Container I J c r) →
            {X : Pred I ℓ} → ⟦ C₁ ∘ C₂ ⟧ X ↔ᵢ (⟦ C₁ ⟧ ⟨∘⟩ ⟦ C₂ ⟧) X
  correct C₁ C₂ {X} = mk↔ₛ′ to from (λ _ → refl) (λ _ → refl)
    where
    to : ⟦ C₁ ∘ C₂ ⟧ X ⊆ ⟦ C₁ ⟧ (⟦ C₂ ⟧ X)
    to ((c , f) , g) = (c , < f , curry g >)

    from : ⟦ C₁ ⟧ (⟦ C₂ ⟧ X) ⊆ ⟦ C₁ ∘ C₂ ⟧ X
    from (c , f) = ((c , proj₁ ⟨∘⟩ f) , uncurry (proj₂ ⟨∘⟩ f))

module Product (ext : ∀ {ℓ} → Extensionality ℓ ℓ) where

  correct : (C₁ C₂ : Container I O c r) {X : Pred I _} →
            ⟦ C₁ × C₂ ⟧ X ↔ᵢ (⟦ C₁ ⟧ X ∩ ⟦ C₂ ⟧ X)
  correct C₁ C₂ {X} = mk↔ₛ′ to from (λ _ → refl) from∘to
    where
    to : ⟦ C₁ × C₂ ⟧ X ⊆ ⟦ C₁ ⟧ X ∩ ⟦ C₂ ⟧ X
    to ((c₁ , c₂) , k) = ((c₁ , k ⟨∘⟩ inj₁) , (c₂ , k ⟨∘⟩ inj₂))

    from : ⟦ C₁ ⟧ X ∩ ⟦ C₂ ⟧ X ⊆ ⟦ C₁ × C₂ ⟧ X
    from ((c₁ , k₁) , (c₂ , k₂)) = ((c₁ , c₂) , [ k₁ , k₂ ])

    from∘to : from ⟨∘⟩ to ≗ F.id
    from∘to (c , _) =
      P.cong (c ,_) (ext [ (λ _ → refl) , (λ _ → refl) ])

module IndexedProduct where

  correct : (C : X → Container I O c r) {Y : Pred I ℓ} →
            ⟦ Π C ⟧ Y ↔ᵢ ⋂[ x ∶ X ] ⟦ C x ⟧ Y
  correct {X = X} C {Y} = mk↔ₛ′ to from (λ _ → refl) (λ _ → refl)
    where
    to : ⟦ Π C ⟧ Y ⊆ ⋂[ x ∶ X ] ⟦ C x ⟧ Y
    to (c , k) = λ x → (c x , λ r → k (x , r))

    from : ⋂[ x ∶ X ] ⟦ C x ⟧ Y ⊆ ⟦ Π C ⟧ Y
    from f = (proj₁ ⟨∘⟩ f , uncurry (proj₂ ⟨∘⟩ f))

module Sum (ext : ∀ {ℓ₁ ℓ₂} → Extensionality ℓ₁ ℓ₂) where

  correct : (C₁ C₂ : Container I O c r) {X : Pred I ℓ} →
            ⟦ C₁ ⊎ C₂ ⟧ X ↔ᵢ (⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X)
  correct C₁ C₂ {X} = mk↔ₛ′ to from to∘from from∘to
    where
    to : ⟦ C₁ ⊎ C₂ ⟧ X ⊆ ⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X
    to (inj₁ c₁ , k) = inj₁ (c₁ , λ r → k (All.inj₁ r))
    to (inj₂ c₂ , k) = inj₂ (c₂ , λ r → k (All.inj₂ r))

    from : ⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X ⊆ ⟦ C₁ ⊎ C₂ ⟧ X
    from (inj₁ (c , f)) = inj₁ c , λ{ (All.inj₁ r) → f r}
    from (inj₂ (c , f)) = inj₂ c , λ{ (All.inj₂ r) → f r}

    from∘to : from ⟨∘⟩ to ≗ F.id
    from∘to (inj₁ _ , _) = P.cong (inj₁ _ ,_) (ext λ{ (All.inj₁ r) → refl})
    from∘to (inj₂ _ , _) = P.cong (inj₂ _ ,_) (ext λ{ (All.inj₂ r) → refl})

    to∘from : to ⟨∘⟩ from ≗ F.id
    to∘from =  [ (λ _ → refl) , (λ _ → refl) ]

module Sum′ where

  correct : (C₁ C₂ : Container I O c r) {X : Pred I ℓ} →
            ⟦ C₁ ⊎′ C₂ ⟧ X ↔ᵢ (⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X)
  correct C₁ C₂ {X} = mk↔ₛ′ to from to∘from from∘to
    where
    to : ⟦ C₁ ⊎′ C₂ ⟧ X ⊆ ⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X
    to (inj₁ c₁ , k) = inj₁ (c₁ , k)
    to (inj₂ c₂ , k) = inj₂ (c₂ , k)

    from : ⟦ C₁ ⟧ X ∪ ⟦ C₂ ⟧ X ⊆ ⟦ C₁ ⊎′ C₂ ⟧ X
    from = [ Prod.map inj₁ F.id , Prod.map inj₂ F.id ]

    from∘to : from ⟨∘⟩ to ≗ F.id
    from∘to (inj₁ _ , _) = refl
    from∘to (inj₂ _ , _) = refl

    to∘from : to ⟨∘⟩ from ≗ F.id
    to∘from = [ (λ _ → refl) , (λ _ → refl) ]

module IndexedSum where

  correct : (C : X → Container I O c r) {Y : Pred I ℓ} →
            ⟦ Σ C ⟧ Y ↔ᵢ ⋃[ x ∶ X ] ⟦ C x ⟧ Y
  correct {X = X} C {Y} = mk↔ₛ′ to from (λ _ → refl) (λ _ → refl)
    where
    to : ⟦ Σ C ⟧ Y ⊆ ⋃[ x ∶ X ] ⟦ C x ⟧ Y
    to ((x , c) , k) = (x , (c , k))

    from : ⋃[ x ∶ X ] ⟦ C x ⟧ Y ⊆ ⟦ Σ C ⟧ Y
    from (x , (c , k)) = ((x , c) , k)

module ConstantExponentiation where

  correct : (C : Container I O c r) {Y : Pred I ℓ} →
            ⟦ const[ X ]⟶ C ⟧ Y ↔ᵢ (⋂ X (F.const (⟦ C ⟧ Y)))
  correct C {Y} = IndexedProduct.correct (F.const C) {Y}