    ------------------------------------------------------------------------
    -- 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 using (Container; _◃_/_; ⟦_⟧;
      Command; Response; ◇; next)
    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 using (Level; _⊔_)
    open import Relation.Unary using (Pred; _⊆_; _∪_; _∩_; ⋃; ⋂)
      renaming (_⟨×⟩_ to _⟪×⟫_; _⟨⊙⟩_ to _⟪⊙⟫_; _⟨⊎⟩_ to _⟪⊎⟫_)
    open import Relation.Binary.PropositionalEquality.Core
      using (_≗_; refl; cong)

    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 = 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 , _) =
          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₁ _ , _) = cong (inj₁ _ ,_) (ext λ{ (All.inj₁ r) → refl})
        from∘to (inj₂ _ , _) = 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}