{-# OPTIONS --universe-polymorphism #-}

module FreeAnnotatedMonad where

open import Data.Container as Container
open import Data.Container.Combinator as C using (const[_]⟶_)
open import Data.Empty
open import Data.List as List
open import Data.Product
open import Function
open import Level
open import Relation.Binary.PropositionalEquality as P using (_≡_)

------------------------------------------------------------------------
-- Fluff

-- The identity functor and composition are defined as records to make
-- Agda's type inference work a bit better.

record Id {ℓ} (A : Set ℓ) : Set ℓ where
  constructor return
  field ! : id A

record _∘̂_ {ℓ} (F G : Set ℓ → Set ℓ) (A : Set ℓ) : Set ℓ where
  constructor effect
  field ! : (F ∘ G) A

------------------------------------------------------------------------
-- Free (?) annotated monad

-- The free (?) annotated monad over a sequence of functors
-- (represented as containers).

Free : ∀ {c} → List (Container c) → Set c → Set c
Free []       = Id
Free (C ∷ Cs) = ⟦ C ⟧ ∘̂ Free Cs

infixl 5 _>>=_

_>>=_ : ∀ {ℓ Cs₁ Cs₂} {A B : Set ℓ} →
        Free Cs₁ A → (A → Free Cs₂ B) → Free (Cs₁ ++ Cs₂) B
_>>=_ {Cs₁ = []}      (return x) g = g x
_>>=_ {Cs₁ = C ∷ Cs₁} (effect f) g =
  effect (Container.map (λ x → x >>= g) f)

------------------------------------------------------------------------
-- Equality

-- Equality. Heterogeneous to avoid having to cast types.

infix 4 _≈_

data _≈_ {c} {A : Set c} : ∀ {Cs₁ Cs₂} →
         Free Cs₁ A → Free Cs₂ A → Set (suc c) where
  return-cong : ∀ {x₁ x₂} (x₁≡x₂ : x₁ ≡ x₂) → return x₁ ≈ return x₂
  effect-cong : ∀ {C Cs₁ Cs₂}
                  {f₁ : ⟦ C ⟧ (Free Cs₁ A)} {f₂ : ⟦ C ⟧ (Free Cs₂ A)}
                (f₁≈f₂ : Container.Eq _≈_ f₁ f₂) → effect f₁ ≈ effect f₂

refl : ∀  {c Cs} {A : Set c} (f : Free Cs A) → f ≈ f
refl {Cs = []}     (return x) = return-cong P.refl
refl {Cs = C ∷ Cs} (effect f) = effect-cong (P.refl , refl ∘ proj₂ f)

postulate -- Somewhat awkward proofs; not used below.

  sym : ∀ {c Cs₁ Cs₂} {A : Set c} {f₁ : Free Cs₁ A} {f₂ : Free Cs₂ A} →
        f₁ ≈ f₂ → f₂ ≈ f₁

  trans : ∀ {c Cs₁ Cs₂ Cs₃} {A : Set c}
            {f₁ : Free Cs₁ A} {f₂ : Free Cs₂ A} {f₃ : Free Cs₃ A} →
          f₁ ≈ f₂ → f₂ ≈ f₃ → f₁ ≈ f₃

------------------------------------------------------------------------
-- Monad laws

left-identity : ∀ {ℓ Cs} {A B : Set ℓ} (x : A) (f : A → Free Cs B) →
                return x >>= f ≈ f x
left-identity x f = refl (f x)

right-identity : ∀ {ℓ Cs} {A : Set ℓ} (m : Free Cs A) →
                 m >>= return ≈ m
right-identity {Cs = []}     (return x) = refl (return x)
right-identity {Cs = C ∷ Cs} (effect f) =
  effect-cong (P.refl , right-identity ∘ proj₂ f)

assoc : ∀ {ℓ Cs₁ Cs₂ Cs₃} {A B C : Set ℓ}
        (m : Free Cs₁ A) (g : A → Free Cs₂ B) (h : B → Free Cs₃ C) →
        m >>= g >>= h ≈ m >>= λ x → g x >>= h
assoc {Cs₁ = []}      (return x) g h = refl (g x >>= h)
assoc {Cs₁ = C ∷ Cs₁} (effect f) g h =
  effect-cong (P.refl , λ p → assoc (proj₂ f p) g h)

------------------------------------------------------------------------
-- Example

-- The n-ary function type: _⇨_.

infixr 4 _⇨_

_⇨_ : ∀ {ℓ} → List (Set ℓ) → Set ℓ → Set ℓ
_⇨_ As = Free (List.map (λ A → const[ A ]⟶ C.id) As)

-- Abstraction.

Λ : ∀ {ℓ A As} {B : Set ℓ} → (A → As ⇨ B) → A ∷ As ⇨ B
Λ f = effect (_ , f ∘ proj₁)

-- Application.

infixl 10 _·_

_·_ : ∀ {ℓ A As} {B : Set ℓ} → A ∷ As ⇨ B → A → As ⇨ B
effect (_ , f) · x = f (x , _)

-- Extraction of the result.

run : ∀ {ℓ} {B : Set ℓ} → [] ⇨ B → B
run (return x) = x