------------------------------------------------------------------------
-- Generalised reverse composition
------------------------------------------------------------------------

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

module ReverseComposition where

open import Category.Monad
open import Data.Empty
open import Data.List as List using (List; []; _∷_; _++_)
open import Data.Maybe as Maybe hiding (map)
open import Data.Nat
open import Data.Product using (_×_; _,_; ∃; ,_; proj₁; proj₂)
open import Data.Unit using (⊤; tt)
open import Data.Vec using (Vec; []; _∷_)
open import Function
open import Level using (Lift; lift)
open import Relation.Binary.PropositionalEquality as P using (_≡_; _≗_)
open import Relation.Binary.HeterogeneousEquality as H using (_≅_)

------------------------------------------------------------------------
-- Functions from a given number of arguments, with given types

-- Injection of values into the monad.

record Done {ℓ} (A : Set ℓ) : Set ℓ where
  constructor return

  -- Extraction of the result.

  field run : A

open Done public

-- Abstraction.

infixl  _→̂_

record _→̂_ {ℓ} (A B : Set ℓ) : Set ℓ where
  constructor Λ

  -- Application.

  infixl  _·_

  field _·_ : A → B

open _→̂_ public

-- Functions from a given number of arguments, with given types.

infixr  _⇨_

_⇨_ : ∀ {ℓ} → List (Set ℓ) → Set ℓ → Set ℓ
[]       ⇨ B = Done B
(A ∷ As) ⇨ B = A →̂ (As ⇨ B)

------------------------------------------------------------------------
-- Heterogeneous equality

infix  _≈_

data _≈_ {ℓ} : ∀ {As₁ As₂} {B₁ B₂ : Set ℓ} →
               As₁ ⇨ B₁ → As₂ ⇨ B₂ → Set (Level.suc ℓ) where
  return-cong : ∀ {B} {x₁ x₂ : B} →
                (x₁≡x₂ : x₁ ≡ x₂) → return x₁ ≈ return x₂
  Λ-cong      : ∀ {A As₁ B₁ As₂ B₂}
                  {f₁ : A → As₁ ⇨ B₁} {f₂ : A → As₂ ⇨ B₂}
                (f₁≈f₂ : ∀ x → f₁ x ≈ f₂ x) → Λ f₁ ≈ Λ f₂

refl : ∀  {ℓ As} {B : Set ℓ} (f : As ⇨ B) → f ≈ f
refl {As = []}     (return x) = return-cong P.refl
refl {As = A ∷ As} (Λ f)      = Λ-cong (λ x → refl (f x))

sym : ∀ {ℓ As₁ As₂} {B₁ B₂ : Set ℓ} {f : As₁ ⇨ B₁} {g : As₂ ⇨ B₂} →
      f ≈ g → g ≈ f
sym (return-cong x≡y) = return-cong (P.sym x≡y)
sym (Λ-cong f≈g)      = Λ-cong (sym ∘ f≈g)

trans : ∀ {ℓ As₁ As₂ As₃} {B₁ B₂ B₃ : Set ℓ}
          {f : As₁ ⇨ B₁} {g : As₂ ⇨ B₂} {h : As₃ ⇨ B₃} →
        f ≈ g → g ≈ h → f ≈ h
trans (return-cong x≡y) (return-cong y≡z) =
  return-cong (P.trans x≡y y≡z)
trans (Λ-cong f≈g) (Λ-cong g≈h) = Λ-cong λ x → trans (f≈g x) (g≈h x)

-- The equality is complete with respect to propositional equality.
-- For (a limited form of) soundness, see run-cong.

complete : ∀ {ℓ As} {B : Set ℓ} {f g : As ⇨ B} → f ≡ g → f ≈ g
complete P.refl = refl _

-- Equational reasoning.

module ≈-Reasoning where

  infix   _∎
  infixr  _≈⟨_⟩_ _≡⟨_⟩_

  _≈⟨_⟩_ : ∀ {ℓ As₁ As₂ As₃} {B₁ B₂ B₃ : Set ℓ}
           (f : As₁ ⇨ B₁) {g : As₂ ⇨ B₂} {h : As₃ ⇨ B₃} →
           f ≈ g → g ≈ h → f ≈ h
  _ ≈⟨ f≈g ⟩ g≈h = trans f≈g g≈h

  _≡⟨_⟩_ : ∀ {ℓ As₁ As₂} {B₁ B₂ : Set ℓ}
           f {g : As₁ ⇨ B₁} {h : As₂ ⇨ B₂} →
           f ≡ g → g ≈ h → f ≈ h
  _ ≡⟨ P.refl ⟩ g≈h = g≈h

  _∎ : ∀ {ℓ As} {B : Set ℓ} (f : As ⇨ B) → f ≈ f
  f ∎ = refl f

-- A module which can be used when deriving code from a specification
-- of the form "∃ λ g → f ≈ g".

module Derivation where

  open ≈-Reasoning public

  infix  ▷_

  EqualTo : ∀ {ℓ As} {B : Set ℓ} → As ⇨ B → Set (Level.suc ℓ)
  EqualTo {As = As} {B} f = ∃ λ (g : As ⇨ B) → f ≈ g

  ▷_ : ∀ {ℓ As} {B : Set ℓ} {f g : As ⇨ B} → f ≈ g → EqualTo f
  ▷_ = ,_

  witness : ∀ {ℓ As} {B : Set ℓ} {f : As ⇨ B} → EqualTo f → As ⇨ B
  witness = proj₁

  proof : ∀ {ℓ As} {B : Set ℓ} {f : As ⇨ B}
          (f≈ : EqualTo f) → f ≈ witness f≈
  proof = proj₂

-- Some congruences.

infixl  _·-cong_

_·-cong_ : ∀ {ℓ As₁ As₂} {A₁ B₁ A₂ B₂ : Set ℓ}
             {x₁ x₂} {f₁ : A₁ ∷ As₁ ⇨ B₁} {f₂ : A₂ ∷ As₂ ⇨ B₂} →
           f₁ ≈ f₂ → x₁ ≅ x₂ → f₁ · x₁ ≈ f₂ · x₂
Λ-cong f₁≈f₂ ·-cong H.refl = f₁≈f₂ _

run-cong : ∀ {ℓ} {B : Set ℓ} {x₁ x₂ : [] ⇨ B} →
           x₁ ≈ x₂ → run x₁ ≡ run x₂
run-cong (return-cong x₁≡x₂) = x₁≡x₂

------------------------------------------------------------------------
-- Generalised reverse composition

-- This construction seems to be closely related to free monads, see
-- FreeAnnotatedMonad.

infixl  _>>=_

_>>=_ : ∀ {ℓ As Bs} {B C : Set ℓ} →
        As ⇨ B → (B → Bs ⇨ C) → As ++ Bs ⇨ C
_>>=_ {As = []}     x g = g (run x)
_>>=_ {As = A ∷ As} f g = Λ λ x → f · x >>= g

-- Bind and return form an annotated monad.

module MonadLaws {ℓ} where

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

  right-identity : ∀ {As} {B : Set ℓ}
                   (f : As ⇨ B) → f >>= return ≈ f
  right-identity {As = []}     x = refl x
  right-identity {As = A ∷ As} f = Λ-cong λ x → right-identity (f · x)

  assoc : ∀ {As Bs Cs} {B C D : Set ℓ}
          (f : As ⇨ B) (g : B → Bs ⇨ C) (h : C → Cs ⇨ D) →
          (f >>= g) >>= h ≈ f >>= (λ x → g x >>= h)
  assoc {As = []}     x g h = refl (g (run x) >>= h)
  assoc {As = A ∷ As} f g h = Λ-cong λ x → assoc (f · x) g h

  infixl  _>>=-cong_ _>>=-cong′_

  _>>=-cong_ : ∀ {As₁ Bs₁ As₂ Bs₂} {B₁ C₁ B₂ C₂ : Set ℓ}
                 {f₁ : As₁ ⇨ B₁} {g₁ : B₁ → Bs₁ ⇨ C₁}
                 {f₂ : As₂ ⇨ B₂} {g₂ : B₂ → Bs₂ ⇨ C₂} →
               f₁ ≈ f₂ → (∀ {x₁ x₂} → x₁ ≅ x₂ → g₁ x₁ ≈ g₂ x₂) →
               f₁ >>= g₁ ≈ f₂ >>= g₂
  return-cong x₁≡x₂ >>=-cong g₁≈g₂ = g₁≈g₂ (H.≡-to-≅ x₁≡x₂)
  Λ-cong f₁≈f₂      >>=-cong g₁≈g₂ = Λ-cong λ x → f₁≈f₂ x >>=-cong g₁≈g₂

  _>>=-cong′_ : ∀ {As₁ Bs₁ As₂ Bs₂} {B C₁ C₂ : Set ℓ}
                  {f₁ : As₁ ⇨ B} {g₁ : B → Bs₁ ⇨ C₁}
                  {f₂ : As₂ ⇨ B} {g₂ : B → Bs₂ ⇨ C₂} →
                f₁ ≈ f₂ → (∀ x → g₁ x ≈ g₂ x) →
                f₁ >>= g₁ ≈ f₂ >>= g₂
  _>>=-cong′_ {g₁ = g₁} {g₂ = g₂} f₁≈f₂ g₁≈g₂ =
    f₁≈f₂ >>=-cong λ {x₁ x₂} x₁≅x₂ →
    H.subst (λ x → g₁ x₁ ≈ g₂ x) x₁≅x₂ (g₁≈g₂ x₁)

-- We can turn this into some kind of restricted monoid.

infixr  _⊙_

_⊙_ : ∀ {ℓ As Bs} {B C : Set ℓ} →
      As ⇨ B → (B ∷ Bs) ⇨ C → As ++ Bs ⇨ C
f ⊙ g = f >>= _·_ g

identity : ∀ {ℓ} {A : Set ℓ} → A ∷ [] ⇨ A
identity = Λ λ x → return x

module MonoidLaws {ℓ} where

  open MonadLaws using (_>>=-cong_)
  open ≈-Reasoning

  right-identity : ∀ {As} {B : Set ℓ}
                   (f : As ⇨ B) → f ⊙ identity ≈ f
  right-identity f =
    f ⊙ identity  ≡⟨ P.refl ⟩
    f >>= return  ≈⟨ MonadLaws.right-identity f ⟩
    f             ∎

  -- Note that the type of f is not completely general.

  left-identity : ∀ {As} {A B : Set ℓ}
                  (f : A ∷ As ⇨ B) → identity ⊙ f ≈ f
  left-identity f =
    identity ⊙ f                  ≡⟨ P.refl ⟩
    (Λ λ x → return x) >>= _·_ f  ≡⟨ P.refl ⟩
    (Λ λ x → return x >>= _·_ f)  ≡⟨ P.refl ⟩
    (Λ λ x → f · x)               ≡⟨ P.refl ⟩
    f                             ∎

  -- Note that the types of g and h are not completely general.

  assoc : ∀ {As Bs Cs} {B C D : Set ℓ}
          (f : As ⇨ B) (g : B ∷ Bs ⇨ C) (h : C ∷ Cs ⇨ D) →
          (f ⊙ g) ⊙ h ≈ f ⊙ (g ⊙ h)
  assoc f g h =
    (f ⊙ g) ⊙ h                    ≡⟨ P.refl ⟩
    f >>= _·_ g >>= _·_ h          ≈⟨ MonadLaws.assoc f (_·_ g) (_·_ h) ⟩
    f >>= (λ x → g · x >>= _·_ h)  ≡⟨ P.refl ⟩
    f ⊙ (g ⊙ h)                    ∎

  infixr  _⊙-cong_

  _⊙-cong_ : ∀ {As₁ Bs₁ As₂ Bs₂} {B₁ C₁ B₂ C₂ : Set ℓ}
               {f₁ : As₁ ⇨ B₁} {g₁ : B₁ ∷ Bs₁ ⇨ C₁}
               {f₂ : As₂ ⇨ B₂} {g₂ : B₂ ∷ Bs₂ ⇨ C₂} →
             f₁ ≈ f₂ → g₁ ≈ g₂ → f₁ ⊙ g₁ ≈ f₂ ⊙ g₂
  f₁≈f₂ ⊙-cong g₁≈g₂ = f₁≈f₂ >>=-cong _·-cong_ g₁≈g₂

  identity-cong : {A : Set ℓ} → identity {A = A} ≈ identity {A = A}
  identity-cong = refl identity

------------------------------------------------------------------------
-- Join and split

join : ∀ {ℓ As Bs} {C : Set ℓ} → As ⇨ Bs ⇨ C → As ++ Bs ⇨ C
join f = f >>= id

split : ∀ {ℓ} As {Bs} {C : Set ℓ} → As ++ Bs ⇨ C → As ⇨ Bs ⇨ C
split []       f = return f
split (A ∷ As) f = Λ λ a → split As (f · a)

module JoinSplit {ℓ} where

  join∘split : ∀ As {Bs} {C : Set ℓ} (f : As ++ Bs ⇨ C) →
               join (split As f) ≈ f
  join∘split []       f = refl f
  join∘split (A ∷ As) f = Λ-cong λ a → join∘split As (f · a)

  split∘join : ∀ {As Bs} {C : Set ℓ} (f : As ⇨ Bs ⇨ C) →
               split As (join f) ≈ f
  split∘join {[]}     f = refl f
  split∘join {A ∷ As} f = Λ-cong λ a → split∘join (f · a)

------------------------------------------------------------------------
-- Generalised currying

Args : ∀ {ℓ} → List (Set ℓ) → Set ℓ
Args = List.foldr _×_ (Lift ⊤)

uncurry : ∀ {ℓ As} {B : Set ℓ} → As ⇨ B → (Args As → B)
uncurry {As = []}     x _        = run x
uncurry {As = A ∷ As} f (a , as) = uncurry (f · a) as

curry : ∀ {ℓ As} {B : Set ℓ} → (Args As → B) → As ⇨ B
curry {As = []}     f = return (f _)
curry {As = A ∷ As} f = Λ λ a → curry (λ as → f (a , as))

module CurryLaws {ℓ} where

  curry∘uncurry : ∀ {As} {B : Set ℓ} (f : As ⇨ B) →
                  curry {As = As} (uncurry f) ≈ f
  curry∘uncurry {As = []}     x = return-cong P.refl
  curry∘uncurry {As = A ∷ As} f = Λ-cong λ a → curry∘uncurry (f · a)

  uncurry∘curry : ∀ As {B : Set ℓ} (f : Args As → B) →
                  uncurry (curry {As = As} f) ≗ f
  uncurry∘curry []       f as = P.refl
  uncurry∘curry (A ∷ As) f as = uncurry∘curry As _ _

  curry-cong : ∀ {As} {B : Set ℓ} {f₁ f₂ : Args As → B} →
               f₁ ≗ f₂ → curry {As = As} f₁ ≈ curry {As = As} f₂
  curry-cong {[]}     f₁≗f₂ = return-cong (f₁≗f₂ _)
  curry-cong {A ∷ As} f₁≗f₂ = Λ-cong λ a →
    curry-cong (λ as → f₁≗f₂ (a , as))

  uncurry-cong : ∀ {As} {B : Set ℓ} {f₁ f₂ : As ⇨ B} →
                 f₁ ≈ f₂ → uncurry f₁ ≗ uncurry f₂
  uncurry-cong (return-cong x₁≡x₂) _        = x₁≡x₂
  uncurry-cong (Λ-cong f₁≈f₂)      (a , as) = uncurry-cong (f₁≈f₂ a) as

------------------------------------------------------------------------
-- A specialised variant of _⇨_

infixr  _↑_⇨_

_↑_⇨_ : ∀ {ℓ} → Set ℓ → ℕ → Set ℓ → Set ℓ
A ↑ n ⇨ B = List.replicate n A ⇨ B

-- Applies the function to a vector of arguments.
--
-- Note that app is morally a special case of uncurry.

app : ∀ {ℓ} {A B : Set ℓ} {n} → A ↑ n ⇨ B → Vec A n → B
app {n = zero}  x []       = run x
app {n = suc n} f (x ∷ xs) = app (f · x) xs

------------------------------------------------------------------------
-- Casts

cast : ∀ {ℓ As₁ As₂} {B : Set ℓ} → As₁ ≡ As₂ → As₁ ⇨ B → As₂ ⇨ B
cast {B = B} = P.subst (λ As → As ⇨ B)

module Cast where

  drop : ∀ {ℓ As₁ As₂} {B : Set ℓ} (eq : As₁ ≡ As₂) (f : As₁ ⇨ B) →
         cast eq f ≈ f
  drop P.refl = refl

------------------------------------------------------------------------
-- _⇨_ can be turned into a monad in a second way

-- Above it was shown that _⇨_ is a parametrised monad. Here it is
-- shown that _⇨_ As is an ordinary monad.

return∥ : ∀ {ℓ As} {B : Set ℓ} → B → As ⇨ B
return∥ {As = As} x = curry {As = As} (const x)

infixl  _>>=∥_

_>>=∥_ : ∀ {ℓ As} {B C : Set ℓ} → As ⇨ B → (B → As ⇨ C) → As ⇨ C
f >>=∥ g = curry λ as → uncurry (g (uncurry f as)) as

-- The monad laws.

module MonadLaws∥ {ℓ} where

  open ≈-Reasoning

  left-identity : ∀ {As} {A B : Set ℓ} (f : A → As ⇨ B) x →
                  return∥ x >>=∥ f ≈ f x
  left-identity {As = []}     f x = refl (f x)
  left-identity {As = A ∷ As} f x =
    Λ (λ y → return∥ x >>=∥ λ z → f z · y)  ≈⟨ Λ-cong (λ y → left-identity (λ z → f z · y) x) ⟩
    Λ (λ y → f x · y)                       ≡⟨ P.refl ⟩
    f x                                     ∎

  right-identity : ∀ {As} {B : Set ℓ}
                   (f : As ⇨ B) → f >>=∥ return∥ ≈ f
  right-identity {As = []}    x = refl x
  right-identity {As = _ ∷ _} f = Λ-cong λ x → right-identity (f · x)

  return∥-cong : ∀ {As} {B : Set ℓ} {x₁ x₂ : B} →
                 x₁ ≡ x₂ → return∥ {As = As} x₁ ≈ return∥ {As = As} x₂
  return∥-cong {As = []}     x₁≡x₂ = return-cong x₁≡x₂
  return∥-cong {As = A ∷ As} x₁≡x₂ = Λ-cong λ _ → return∥-cong x₁≡x₂

  infixl  _>>=∥-cong_ _>>=∥-cong′_

  _>>=∥-cong_ : ∀ {As₁ As₂} {B₁ C₁ B₂ C₂ : Set ℓ}
                  {f₁ : As₁ ⇨ B₁} {g₁ : B₁ → As₁ ⇨ C₁}
                  {f₂ : As₂ ⇨ B₂} {g₂ : B₂ → As₂ ⇨ C₂} →
                f₁ ≈ f₂ → (∀ {x₁ x₂} → x₁ ≅ x₂ → g₁ x₁ ≈ g₂ x₂) →
                f₁ >>=∥ g₁ ≈ f₂ >>=∥ g₂
  return-cong x₁≡x₂ >>=∥-cong g₁≈g₂ = g₁≈g₂ (H.≡-to-≅ x₁≡x₂)
  Λ-cong f₁≈f₂      >>=∥-cong g₁≈g₂ = Λ-cong λ x →
    f₁≈f₂ x >>=∥-cong λ x₁≅x₂ → g₁≈g₂ x₁≅x₂ ·-cong H.refl

  _>>=∥-cong′_ : ∀ {As₁ As₂} {B C₁ C₂ : Set ℓ}
                   {f₁ : As₁ ⇨ B} {g₁ : B → As₁ ⇨ C₁}
                   {f₂ : As₂ ⇨ B} {g₂ : B → As₂ ⇨ C₂} →
                 f₁ ≈ f₂ → (∀ x → g₁ x ≈ g₂ x) →
                 f₁ >>=∥ g₁ ≈ f₂ >>=∥ g₂
  _>>=∥-cong′_ {g₁ = g₁} {g₂ = g₂} f₁≈f₂ g₁≈g₂ =
    f₁≈f₂ >>=∥-cong λ {x₁ x₂} x₁≅x₂ →
    H.subst (λ x → g₁ x₁ ≈ g₂ x) x₁≅x₂ (g₁≈g₂ x₁)

  assoc : ∀ {As} {B C D : Set ℓ}
          (f : As ⇨ B) (g : B → As ⇨ C) (h : C → As ⇨ D) →
          (f >>=∥ g) >>=∥ h ≈ f >>=∥ (λ x → g x >>=∥ h)
  assoc {As = []}    x g h = refl (g (run x) >>=∥ h)
  assoc {As = _ ∷ _} f g h = Λ-cong λ x →
    (f · x >>=∥ λ y → g y · x) >>=∥ (λ y → h y · x)  ≈⟨ assoc (f · x) (λ y → g y · x) (λ y → h y · x) ⟩
    f · x >>=∥ (λ y → g y · x >>=∥ λ z → h z · x)    ≡⟨ P.refl ⟩
    f · x >>=∥ (λ y → (g y >>=∥ h) · x)              ∎

-- Some derived combinators.

infixr  _⊙∥_

_⊙∥_ : ∀ {ℓ As} {B C : Set ℓ} → As ⇨ B → B ∷ As ⇨ C → As ⇨ C
f ⊙∥ g = f >>=∥ _·_ g

identity∥ : ∀ {ℓ As} {A : Set ℓ} → A ∷ As ⇨ A
identity∥ = Λ λ x → return∥ x

-- Note that map corresponds to Wand's combinator B_k (see "Deriving
-- Target Code as a Representation of Continuation Semantics").

map : ∀ {ℓ As} {B C : Set ℓ} → (B → C) → As ⇨ B → As ⇨ C
map f g = g >>=∥ return∥ ∘ f

infixl  _⊛∥_

_⊛∥_ : ∀ {ℓ As} {B C : Set ℓ} → As ⇨ (B → C) → As ⇨ B → As ⇨ C
f ⊛∥ g = f >>=∥ λ h → map h g

zipWith∥ : ∀ {ℓ As} {B C D : Set ℓ} →
           (B → C → D) → As ⇨ B → As ⇨ C → As ⇨ D
zipWith∥ f g h = return∥ f ⊛∥ g ⊛∥ h

const⋆ : ∀ {ℓ} As {Bs} {C : Set ℓ} → Bs ⇨ C → As ++ Bs ⇨ C
const⋆ As f = return∥ {As = As} (lift tt) >>= const f

-- More laws.

module MoreLaws∥ {ℓ} where

  open ≈-Reasoning

  -- _⊙_ distributes from the left over _⊙∥_.

  ⊙-⊙∥ : ∀ {As Bs} {C D E : Set ℓ}
         (f : As ⇨ C) (g : C ∷ Bs ⇨ D) (h : D ∷ C ∷ Bs ⇨ E) →
         f ⊙ (g ⊙∥ h) ≈ (f ⊙ g) ⊙∥ Λ (λ x → f ⊙ h · x)
  ⊙-⊙∥ {[]}     x g h = refl _
  ⊙-⊙∥ {A ∷ As} f g h = Λ-cong λ x → ⊙-⊙∥ (f · x) g h

  swap->>=∥ : ∀ {As} {B C D : Set ℓ}
              (f : As ⇨ B) (g : As ⇨ C) (h : B → C → As ⇨ D) →
              (f >>=∥ λ x → g >>=∥ h x) ≈ (g >>=∥ λ x → f >>=∥ flip h x)
  swap->>=∥ {[]}     f g h = refl _
  swap->>=∥ {A ∷ As} f g h = Λ-cong λ x →
    swap->>=∥ (f · x) (g · x) (λ y z → h y z · x)

  return∥->>= : ∀ As {Bs} {B C : Set ℓ} (x : B) (f : B → Bs ⇨ C) →
                return∥ {As = As} x >>= f ≈ const⋆ As (f x)
  return∥->>= []       x f = refl (f x)
  return∥->>= (A ∷ As) x f = Λ-cong λ _ → return∥->>= As x f

  >>=-return∥ : ∀ {As Bs} {B C : Set ℓ} (f : As ⇨ B) (x : C) →
                f >>= const (return∥ {As = Bs} x) ≈
                return∥ {As = As ++ Bs} x
  >>=-return∥ {As = []}    x y = refl (return∥ y)
  >>=-return∥ {As = _ ∷ _} f y = Λ-cong λ x → >>=-return∥ (f · x) y

  const⋆-return∥ : ∀ As {Bs} {C : Set ℓ} (x : C) →
                   const⋆ As (return∥ {As = Bs} x) ≈
                   return∥ {As = As ++ Bs} x
  const⋆-return∥ As x = >>=-return∥ (return∥ {As = As} (lift tt)) x

  const⋆-const⋆ : ∀ As Bs {Cs} {D : Set ℓ} (f : Cs ⇨ D) →
                  const⋆ As (const⋆ Bs f) ≈ const⋆ (As ++ Bs) f
  const⋆-const⋆ []       Bs f = refl (const⋆ Bs f)
  const⋆-const⋆ (A ∷ As) Bs f = Λ-cong (λ _ → const⋆-const⋆ As Bs f)

  map-id : ∀ {As} {B : Set ℓ} (f : As ⇨ B) → map id f ≈ f
  map-id f = MonadLaws∥.right-identity f

  map-∘ : ∀ {As} {B C D : Set ℓ} (f : C → D) (g : B → C) (h : As ⇨ B) →
          map f (map g h) ≈ map (f ∘ g) h
  map-∘ {As = As} f g h =
    (h >>=∥ return∥ ∘ g) >>=∥ return∥ ∘ f          ≈⟨ assoc h (return∥ ∘ g) (return∥ ∘ f) ⟩
    h >>=∥ (λ x → return∥ (g x) >>=∥ return∥ ∘ f)  ≈⟨ (refl h >>=∥-cong′ λ x → left-identity (return∥ {As = As} ∘ f) (g x)) ⟩
    h >>=∥ return∥ ∘ f ∘ g                         ∎
    where open MonadLaws∥

  -- map f commutes with _⊙_ g.

  map-⊙ : ∀ {As Bs} {B C D : Set ℓ}
          (f : C → D) (g : As ⇨ B) (h : B ∷ Bs ⇨ C) →
          map f (g ⊙ h) ≈ g ⊙ map f h
  map-⊙ {As = []} f x h =
    map f (x ⊙ h)      ≡⟨ P.refl ⟩
    map f (h · run x)  ≡⟨ P.refl ⟩
    x ⊙ map f h        ∎
  map-⊙ {As = _ ∷ _} f g h =
    map f (g ⊙ h)                                     ≡⟨ P.refl ⟩
    (Λ λ x → g · x ⊙ h >>=∥ λ y → return∥ (f y) · x)  ≈⟨ (Λ-cong λ x → refl (g · x ⊙ h) >>=∥-cong′ λ y → refl (return∥ (f y))) ⟩
    (Λ λ x → g · x ⊙ h >>=∥ return∥ ∘ f)              ≡⟨ P.refl ⟩
    (Λ λ x → map f (g · x ⊙ h))                       ≈⟨ (Λ-cong λ x → map-⊙ f (g · x) h) ⟩
    (Λ λ x → g · x ⊙ map f h)                         ≡⟨ P.refl ⟩
    g ⊙ map f h                                       ∎
    where open MonadLaws∥

  map->>=∥ : ∀ {As} {B C D : Set ℓ}
             (f : B → C) (g : As ⇨ B) (h : C → As ⇨ D) →
             map f g >>=∥ h ≈ g >>=∥ h ∘ f
  map->>=∥ f g h =
    map f g >>=∥ h                       ≡⟨ P.refl ⟩
    g >>=∥ return∥ ∘ f >>=∥ h            ≈⟨ assoc g (return∥ ∘ f) h ⟩
    (g >>=∥ λ x → return∥ (f x) >>=∥ h)  ≈⟨ (refl g >>=∥-cong′ λ x → left-identity h (f x)) ⟩
    (g >>=∥ λ x → h (f x))               ∎
    where open MonadLaws∥

  map->>=∥′ : ∀ {As} {B C D : Set ℓ}
              (f : C → D) (g : As ⇨ B) (h : B → As ⇨ C) →
              map f (g >>=∥ h) ≈ g >>=∥ map f ∘ h
  map->>=∥′ f g h =
    map f (g >>=∥ h)                     ≡⟨ P.refl ⟩
    (g >>=∥ h) >>=∥ return∥ ∘ f          ≈⟨ MonadLaws∥.assoc g h (return∥ ∘ f) ⟩
    g >>=∥ (λ x → h x >>=∥ return∥ ∘ f)  ≡⟨ P.refl ⟩
    g >>=∥ map f ∘ h                     ∎

  return∥-⊛∥ : ∀ {As} {B C : Set ℓ} (f : B → C) (g : As ⇨ B) →
               return∥ f ⊛∥ g ≈ map f g
  return∥-⊛∥ f g =
    return∥ f ⊛∥ g                  ≡⟨ P.refl ⟩
    return∥ f >>=∥ (λ h → map h g)  ≈⟨ MonadLaws∥.left-identity (λ h → map h g) f ⟩
    map f g                         ∎

  identity-⊛∥ : ∀ {As} {B : Set ℓ} (f : As ⇨ B) →
                return∥ id ⊛∥ f ≈ f
  identity-⊛∥ f =
    return∥ id ⊛∥ f  ≈⟨ return∥-⊛∥ id f ⟩
    map id f         ≈⟨ map-id f ⟩
    f                ∎

  identity-⊙∥ : ∀ {As} {B : Set ℓ}
                (f : As ⇨ B) → f ⊙∥ identity∥ ≈ f
  identity-⊙∥ f =
    f ⊙∥ identity∥  ≡⟨ P.refl ⟩
    f >>=∥ return∥  ≈⟨ MonadLaws∥.right-identity f ⟩
    f               ∎

  ⊛∥-return∥ : ∀ {As} {B C : Set ℓ} (f : As ⇨ (B → C)) (x : B) →
               f ⊛∥ return∥ x ≈ map (λ g → g x) f
  ⊛∥-return∥ {As} f x =
    f ⊛∥ return∥ x                             ≡⟨ P.refl ⟩
    f >>=∥ (λ h → map h (return∥ x))           ≡⟨ P.refl ⟩
    f >>=∥ (λ h → return∥ x >>=∥ return∥ ∘ h)  ≈⟨ (refl f >>=∥-cong′ λ h →
                                                   left-identity (return∥ {As = As} ∘ h) x) ⟩
    f >>=∥ (λ h → return∥ (h x))               ≡⟨ P.refl ⟩
    map (λ g → g x) f                          ∎
    where open MonadLaws∥

  composition :
    ∀ {As} {B C D : Set ℓ} (f : As ⇨ (C → D)) (g : As ⇨ (B → C)) h →
    map {As = As} _∘′_ f ⊛∥ g ⊛∥ h ≈ f ⊛∥ (g ⊛∥ h)
  composition {As} f g h =
    map {As = As} _∘′_ f ⊛∥ g ⊛∥ h                        ≡⟨ P.refl ⟩
    (map {As = As} _∘′_ f >>=∥ pam g >>=∥ pam h)          ≈⟨ (map->>=∥ _∘′_ f (pam g) >>=∥-cong′ λ f →
                                                              refl (map f h)) ⟩
    ((f >>=∥ λ f → map {As = As} (_∘′_ f) g) >>=∥ pam h)  ≈⟨ assoc f (λ f → map {As = As} (_∘′_ f) g) (pam h) ⟩
    (f >>=∥ λ f → map {As = As} (_∘′_ f) g >>=∥ pam h)    ≈⟨ (refl f >>=∥-cong′ λ f →
                                                              map->>=∥ (_∘′_ f) g (pam h)) ⟩
    (f >>=∥ λ f → g >>=∥ λ g → map (f ∘′ g) h)            ≈⟨ (refl f >>=∥-cong′ λ f →
                                                              refl g >>=∥-cong′ λ g →
                                                              sym $ map-∘ f g h) ⟩
    (f >>=∥ λ f → g >>=∥ λ g → map f (map g h))           ≈⟨ (refl f >>=∥-cong′ λ f →
                                                              sym $ assoc g (pam h) (return∥ ∘ f)) ⟩
    (f >>=∥ pam (g >>=∥ pam h))                           ≡⟨ P.refl ⟩
    f ⊛∥ (g ⊛∥ h)                                         ∎
    where
    open MonadLaws∥

    pam : ∀ {As} {B C : Set ℓ} → As ⇨ B → (B → C) → As ⇨ C
    pam = flip map

  left-identity : ∀ {As} {B C D : Set ℓ}
                  (f : B → C → D) (x : B) (ys : As ⇨ C) →
                  zipWith∥ f (return∥ x) ys ≈ map (f x) ys
  left-identity {As = As} f x ys =
    (return∥ f >>=∥ λ h → return∥ x >>=∥ return∥ ∘ h) >>=∥ (λ h → map h ys)  ≈⟨ (MonadLaws∥.left-identity
                                                                                   (λ h → return∥ x >>=∥ return∥ ∘ h) f >>=∥-cong′ λ h →
                                                                                 refl (map h ys)) ⟩
    (return∥ x >>=∥ return∥ ∘ f) >>=∥ (λ h → map h ys)                       ≈⟨ (MonadLaws∥.left-identity
                                                                                   (return∥ {As = As} ∘ f) x >>=∥-cong′ λ h →
                                                                                 refl (map h ys)) ⟩
    return∥ (f x) >>=∥ (λ h → map h ys)                                      ≈⟨ MonadLaws∥.left-identity (λ h → map h ys) (f x) ⟩
    map (f x) ys                                                             ∎
    where open MonadLaws∥

  right-identity : ∀ {As} {B C D : Set ℓ}
                   (f : B → C → D) (xs : As ⇨ B) (y : C) →
                   zipWith∥ f xs (return∥ y) ≈ map (flip f y) xs
  right-identity {As = As} f xs y =
    (return∥ f >>=∥ λ h → map h xs) >>=∥ (λ h → return∥ y >>=∥ return∥ ∘ h)  ≈⟨ (MonadLaws∥.left-identity (λ h → map h xs) f
                                                                                   >>=∥-cong′ λ h →
                                                                                 MonadLaws∥.left-identity (return∥ {As = As} ∘ h) y) ⟩
    map f xs >>=∥ (λ h → return∥ (h y))                                      ≡⟨ P.refl ⟩
    (xs >>=∥ return∥ ∘ f) >>=∥ (λ h → return∥ (h y))                         ≈⟨ assoc xs (return∥ ∘ f) (λ h → return∥ (h y)) ⟩
    xs >>=∥ (λ x → return∥ (f x) >>=∥ λ h → return∥ (h y))                   ≈⟨ (refl xs >>=∥-cong′ λ x →
                                                                                 MonadLaws∥.left-identity (λ h → return∥ (h y)) (f x)) ⟩
    map (flip f y) xs                                                        ∎
    where open MonadLaws∥

  assoc : ∀ {As} {B C D E F G : Set ℓ}
            {f₁ : B → C → F} {g₁ : F → D → E}
            {f₂ : B → G → E} {g₂ : C → D → G}
          (xs : As ⇨ B) (ys : As ⇨ C) (zs : As ⇨ D) →
          (∀ x y z → g₁ (f₁ x y) z ≡ f₂ x (g₂ y z)) →
          zipWith∥ g₁ (zipWith∥ f₁ xs ys) zs ≈
          zipWith∥ f₂ xs (zipWith∥ g₂ ys zs)
  assoc {As = []}    x y z hyp =
    return-cong (hyp (run x) (run y) (run z))
  assoc {As = _ ∷ _} f g h hyp =
    Λ-cong λ x → assoc (f · x) (g · x) (h · x) hyp

  infixr  _⊙∥-cong_

  _⊙∥-cong_ : ∀ {As₁ As₂} {B C : Set ℓ}
                {f₁ : As₁ ⇨ B} {g₁ : B ∷ As₁ ⇨ C}
                {f₂ : As₂ ⇨ B} {g₂ : B ∷ As₂ ⇨ C} →
              f₁ ≈ f₂ → g₁ ≈ g₂ → f₁ ⊙∥ g₁ ≈ f₂ ⊙∥ g₂
  f₁≈f₂ ⊙∥-cong Λ-cong g₁≈g₂ = f₁≈f₂ >>=∥-cong′ g₁≈g₂
    where open MonadLaws∥

  identity∥-cong :
    ∀ {As} {A : Set ℓ} →
    identity∥ {As = As} {A = A} ≈ identity∥ {As = As} {A = A}
  identity∥-cong = refl identity∥

  map-cong : ∀ {As} {B C : Set ℓ} {f₁ f₂ : B → C} {g₁ g₂ : As ⇨ B} →
             (∀ x → f₁ x ≡ f₂ x) → g₁ ≈ g₂ → map f₁ g₁ ≈ map f₂ g₂
  map-cong f₁≡f₂ g₁≈g₂ = g₁≈g₂ >>=∥-cong′ λ x → return∥-cong (f₁≡f₂ x)
    where open MonadLaws∥

  infixl  _⊛∥-cong_

  _⊛∥-cong_ : ∀ {As} {B C : Set ℓ}
                {f₁ f₂ : As ⇨ (B → C)} {g₁ g₂ : As ⇨ B} →
              f₁ ≈ f₂ → g₁ ≈ g₂ → f₁ ⊛∥ g₁ ≈ f₂ ⊛∥ g₂
  f₁≈f₂ ⊛∥-cong g₁≈g₂ =
    f₁≈f₂ >>=∥-cong λ h₁≅h₂ →
    map-cong (λ x → P.cong (λ f → f x) (H.≅-to-≡ h₁≅h₂)) g₁≈g₂
    where open MonadLaws∥

  zipWith∥-cong : ∀ {As} {B C D : Set ℓ} {f₁ f₂ : B → C → D}
                    {g₁ g₂ : As ⇨ B} {h₁ h₂ : As ⇨ C} →
                  (∀ x y → f₁ x y ≡ f₂ x y) → g₁ ≈ g₂ → h₁ ≈ h₂ →
                  zipWith∥ f₁ g₁ h₁ ≈ zipWith∥ f₂ g₂ h₂
  zipWith∥-cong f₁≡f₂ (return-cong P.refl) (return-cong P.refl) =
    return-cong (f₁≡f₂ _ _)
  zipWith∥-cong f₁≡f₂ (Λ-cong g₁≈g₂)       (Λ-cong h₁≈h₂)       =
    Λ-cong λ x → zipWith∥-cong f₁≡f₂ (g₁≈g₂ x) (h₁≈h₂ x)

  const⋆-cong : ∀ As {Bs} {C : Set ℓ} {f₁ f₂ : Bs ⇨ C} →
                f₁ ≈ f₂ → const⋆ As f₁ ≈ const⋆ As f₂
  const⋆-cong As f₁≈f₂ =
    refl (return∥ {As = As} (lift tt)) >>=-cong′ λ _ → f₁≈f₂
    where open MonadLaws

  -- map can be defined using curry and uncurry.

  map-curry-uncurry :
    ∀ {As} {B C : Set ℓ} (f : B → C) (g : As ⇨ B) →
    map f g ≈ curry {As = As} (λ as → f (uncurry g as))
  map-curry-uncurry {As} f g =
    g >>=∥ return∥ ∘ f                                                ≡⟨ P.refl ⟩
    curry (λ as → uncurry (return∥ {As = As} (f (uncurry g as))) as)  ≈⟨ curry-cong (λ as →
                                                                           uncurry∘curry As (const (f (uncurry g as))) as) ⟩
    curry (λ as → f (uncurry g as))                                   ∎
    where open CurryLaws

  -- _⊛∥_ can be defined using curry and uncurry.

  ⊛∥-curry-uncurry :
    ∀ {As} {B C : Set ℓ} (f : As ⇨ (B → C)) (g : As ⇨ B) →
    f ⊛∥ g ≈ curry {As = As} (λ as → (uncurry f as) (uncurry g as))
  ⊛∥-curry-uncurry {As} f g =
    (f >>=∥ λ h → map h g)                            ≡⟨ P.refl ⟩
    curry (λ as → uncurry (map (uncurry f as) g) as)  ≈⟨ curry-cong (λ as → uncurry-cong (map-curry-uncurry (uncurry f as) g) as) ⟩
    curry (λ as →
      uncurry (curry {As = As} (λ as′ →
        uncurry f as (uncurry g as′))) as)            ≈⟨ curry-cong (λ as → uncurry∘curry As (uncurry f as ∘ uncurry g) as) ⟩
    curry (λ as → uncurry f as (uncurry g as))        ∎
    where open CurryLaws

  -- const⋆ distributes over various operations.

  const⋆-map : ∀ As {Bs} {C D : Set ℓ} (f : C → D) (g : Bs ⇨ C) →
               const⋆ As (map f g) ≈ map f (const⋆ As g)
  const⋆-map As f g =
    const⋆ As (map f g)                              ≡⟨ P.refl ⟩
    return∥ {As = As} (lift tt) ⊙ map f (Λ λ _ → g)  ≈⟨ sym $ map-⊙ f (return∥ {As = As} (lift tt)) (Λ λ _ → g) ⟩
    map f (return∥ {As = As} (lift tt) ⊙ Λ λ _ → g)  ≡⟨ P.refl ⟩
    map f (const⋆ As g)                              ∎

  const⋆->>=∥ : ∀ As {Bs} {C D : Set ℓ} (f : Bs ⇨ C) (g : C → Bs ⇨ D) →
               const⋆ As (f >>=∥ g) ≈ const⋆ As f >>=∥ const⋆ As ∘ g
  const⋆->>=∥ As f g =
    const⋆ As (f >>=∥ g)                                 ≡⟨ P.refl ⟩
    return∥ {As = As} (lift tt) >>= const (f >>=∥ g)     ≈⟨ ⊙-⊙∥ (return∥ {As = As} (lift tt)) (Λ λ _ → f) (Λ λ x → Λ λ _ → g x) ⟩
    (return∥ {As = As} (lift tt) >>= const f >>=∥ λ x →
     return∥ {As = As} (lift tt) >>= const (g x))        ≡⟨ P.refl ⟩
    const⋆ As f >>=∥ const⋆ As ∘ g                       ∎

  const⋆-⊛∥ : ∀ As {Bs} {C D : Set ℓ} (f : Bs ⇨ (C → D)) (g : Bs ⇨ C) →
              const⋆ As (f ⊛∥ g) ≈ const⋆ As f ⊛∥ const⋆ As g
  const⋆-⊛∥ As f g =
    const⋆ As (f ⊛∥ g)                            ≡⟨ P.refl ⟩
    const⋆ As (f >>=∥ λ f → map f g)              ≈⟨ const⋆->>=∥ As f (λ f → map f g) ⟩
    const⋆ As f >>=∥ (λ f → const⋆ As (map f g))  ≈⟨ (refl (const⋆ As f) >>=∥-cong′ λ f → const⋆-map As f g) ⟩
    const⋆ As f >>=∥ (λ f → map f (const⋆ As g))  ≡⟨ P.refl ⟩
    const⋆ As f ⊛∥ const⋆ As g                    ∎
    where open MonadLaws∥

------------------------------------------------------------------------
-- Some combinators which manipulate the "stack" of arguments

-- Discards the top-most stack element.

wk : ∀ {ℓ As} {A B : Set ℓ} → As ⇨ B → A ∷ As ⇨ B
wk f = Λ λ _ → f

-- Application.

ap : ∀ {ℓ As} {B C D : Set ℓ} →
     C ∷ As ⇨ D → (B → C) ∷ B ∷ As ⇨ D
ap g = Λ λ f → Λ λ x → g · f x

-- Pushes an argument onto the stack.

push : ∀ {ℓ Bs} {A C : Set ℓ} → A ∷ Bs ⇨ C → Bs ⇨ (A → C)
push f = curry (λ bs a → uncurry f (a , bs))

-- Pops the stack.

pop : ∀ {ℓ Bs} {A C : Set ℓ} → Bs ⇨ (A → C) → A ∷ Bs ⇨ C
pop f = Λ λ x → map (λ g → g x) f

module Stack {ℓ} where

  open ≈-Reasoning

  wk-ap :
    ∀ {As} {B C D : Set ℓ}
    (f : As ⇨ (B → C)) (g : As ⇨ B) (h : C ∷ As ⇨ D) →
    (f ⊛∥ g) ⊙∥ h ≈ g ⊙∥ wk f ⊙∥ ap h
  wk-ap f g h =
    (f ⊛∥ g) ⊙∥ h                        ≡⟨ P.refl ⟩
    (f >>=∥ λ f → map f g) >>=∥ _·_ h    ≈⟨ MonadLaws∥.assoc f (λ f → map f g) (_·_ h) ⟩
    (f >>=∥ λ f → map f g >>=∥ _·_ h)    ≈⟨ (refl f >>=∥-cong′ λ f → map->>=∥ f g (_·_ h)) ⟩
    (f >>=∥ λ f → g >>=∥ _·_ h ∘ f)      ≈⟨ swap->>=∥ f g (_∘_ (_·_ h)) ⟩
    (g >>=∥ λ x → f >>=∥ λ f → h · f x)  ≡⟨ P.refl ⟩
    g ⊙∥ wk f ⊙∥ ap h                    ∎
    where
    open MonadLaws∥
    open MoreLaws∥

  map-push : ∀ {Bs} {A C : Set ℓ} (f : A ∷ Bs ⇨ C) x →
             map (λ f → f x) (push f) ≈ f · x
  map-push {Bs} f x =
    map (λ f → f x) (push f)                             ≈⟨ MoreLaws∥.map-curry-uncurry (λ f → f x) (push f) ⟩
    curry (λ bs → uncurry (push f) bs x)                 ≡⟨ P.refl ⟩
    curry (λ bs →
      uncurry {B = _ → _}
        (curry {As = Bs} (λ bs x → uncurry f (x , bs)))
        bs x)                                            ≈⟨ curry-cong (λ bs → P.cong (λ (f : _ → _) → f x) $
                                                              uncurry∘curry Bs (λ bs x → uncurry f (x , bs)) bs) ⟩
    curry (λ bs → uncurry f (x , bs))                    ≈⟨ curry∘uncurry (f · x) ⟩
    f · x                                                ∎
    where open CurryLaws

  push-cong : P.Extensionality ℓ ℓ →
              ∀ {Bs} {A C : Set ℓ} {f₁ f₂ : A ∷ Bs ⇨ C} →
              f₁ ≈ f₂ → push f₁ ≈ push f₂
  push-cong ext f₁≈f₂ =
    curry-cong λ bs →
    ext λ a →
    uncurry-cong f₁≈f₂ (a , bs)
    where open CurryLaws

  pop-cong : ∀ {Bs} {A C : Set ℓ} {f₁ f₂ : Bs ⇨ (A → C)} →
             f₁ ≈ f₂ → pop f₁ ≈ pop f₂
  pop-cong f₁≈f₂ = Λ-cong λ x →
    MoreLaws∥.map-cong (λ g → P.refl {x = g x}) f₁≈f₂

  push-pop : P.Extensionality ℓ ℓ →
             ∀ {Bs} {A C : Set ℓ} (f : Bs ⇨ (A → C)) →
             push (pop f) ≈ f
  push-pop ext {Bs} f =
    push (pop f)                                             ≡⟨ P.refl ⟩
    push (Λ λ a → map (λ g → g a) f)                         ≈⟨ push-cong ext (Λ-cong λ a → MoreLaws∥.map-curry-uncurry (λ g → g a) f) ⟩
    push (Λ λ a → curry (λ bs → uncurry f bs a))             ≡⟨ P.refl ⟩
    curry (λ bs a →
      uncurry (curry {As = Bs} (λ bs → uncurry f bs a)) bs)  ≈⟨ curry-cong (λ bs → ext λ a →
                                                                  uncurry∘curry Bs (λ bs → uncurry f bs a) bs) ⟩
    curry (λ bs a → uncurry f bs a)                          ≡⟨ P.refl ⟩
    curry (uncurry f)                                        ≈⟨ curry∘uncurry f ⟩
    f                                                        ∎
    where open CurryLaws

  pop-push : ∀ {Bs} {A C : Set ℓ} (f : A ∷ Bs ⇨ C) →
             pop (push f) ≈ f
  pop-push f = Λ-cong λ a →
    map (λ g → g a) (push f)  ≈⟨ map-push f a ⟩
    f · a                     ∎

------------------------------------------------------------------------
-- A special case

-- _⇨_ where the target type is Maybe something.

infixl  _⇨M_

_⇨M_ : ∀ {ℓ} → List (Set ℓ) → Set ℓ → Set ℓ
As ⇨M B = As ⇨ Maybe B

-- Variants of the monadic combinators.

returnM : ∀ {ℓ} {B : Set ℓ} → B → [] ⇨M B
returnM = return ∘′ just

infixr  _then_else_

_then_else_ : ∀ {ℓ As Bs} {B C : Set ℓ} →
              As ⇨M B → (B → Bs ⇨M C) → Bs ⇨M C → As ++ Bs ⇨M C
g then s else f = g >>= maybe s f

-- Failure and left-biased choice.

zeroM : ∀ {ℓ As} {B : Set ℓ} → As ⇨M B
zeroM = return∥ nothing

infixr  _⊕_

_⊕_ : ∀ {ℓ As} {B : Set ℓ} → As ⇨M B → As ⇨M B → As ⇨M B
_⊕_ = zipWith∥ (RawMonadPlus._∣_ Maybe.monadPlus)

-- Various laws.

module LawsM {ℓ} where

  open ≈-Reasoning

  >>=M-cong : ∀ {As₁ Bs₁ As₂ Bs₂} {B C₁ C₂ : Set ℓ}
                {f₁ : As₁ ⇨M B} {g₁ : Maybe B → Bs₁ ⇨M C₁}
                {f₂ : As₂ ⇨M B} {g₂ : Maybe B → Bs₂ ⇨M C₂} →
              f₁ ≈ f₂ →
              (∀ x → g₁ (just x) ≈ g₂ (just x)) →
              g₁ nothing ≈ g₂ nothing →
              f₁ >>= g₁ ≈ f₂ >>= g₂
  >>=M-cong f₁≈f₂ g₁j≈g₂j g₁n≈g₂n =
    f₁≈f₂ >>=-cong′ maybe g₁j≈g₂j g₁n≈g₂n
    where open MonadLaws

  maybe′-cong : ∀ {Bs₁ Bs₂} {B C₁ C₂ : Set ℓ}
                 {s₁ : B → Bs₁ ⇨M C₁} {f₁}
                 {s₂ : B → Bs₂ ⇨M C₂} {f₂} →
    (∀ x → s₁ x ≈ s₂ x) → f₁ ≈ f₂ →
    ∀ x → maybe′ s₁ f₁ x ≈ maybe′ s₂ f₂ x
  maybe′-cong s₁≈s₂ f₁≈f₂ (just x) = s₁≈s₂ x
  maybe′-cong s₁≈s₂ f₁≈f₂ nothing  = f₁≈f₂

  infixr  _then_else-cong_

  _then_else-cong_ : ∀ {As₁ Bs₁ As₂ Bs₂} {B C₁ C₂ : Set ℓ}
                       {g₁ : As₁ ⇨M B} {s₁ : B → Bs₁ ⇨M C₁} {f₁}
                       {g₂ : As₂ ⇨M B} {s₂ : B → Bs₂ ⇨M C₂} {f₂} →
    g₁ ≈ g₂ → (∀ x → s₁ x ≈ s₂ x) → f₁ ≈ f₂ →
    g₁ then s₁ else f₁ ≈ g₂ then s₂ else f₂
  g₁≈g₂ then s₁≈s₂ else-cong f₁≈f₂ =
    g₁≈g₂ >>=-cong′ maybe′-cong s₁≈s₂ f₁≈f₂
    where open MonadLaws

  zeroM-cong : ∀ {As} {B : Set ℓ} →
               zeroM {As = As} {B = B} ≈ zeroM {As = As} {B = B}
  zeroM-cong = MonadLaws∥.return∥-cong P.refl

  infixr  _⊕-cong_

  _⊕-cong_ : ∀ {As} {B : Set ℓ} {f₁ f₂ : As ⇨M B} {g₁ g₂ : As ⇨M B} →
             f₁ ≈ f₂ → g₁ ≈ g₂ → f₁ ⊕ g₁ ≈ f₂ ⊕ g₂
  f₁≈f₂ ⊕-cong g₁≈g₂ =
    MoreLaws∥.zipWith∥-cong (λ _ _ → P.refl) f₁≈f₂ g₁≈g₂

  left-zero : ∀ As {Bs} {B C : Set ℓ}
              (s : B → Bs ⇨M C) (f : Bs ⇨M C) →
              zeroM {As = As} then s else f ≈ const⋆ As f
  left-zero As s f = MoreLaws∥.return∥->>= As nothing (maybe s f)

  right-zeros : ∀ {As Bs} {B C : Set ℓ} (g : As ⇨M B) →
                g then const zeroM else zeroM {As = Bs} {B = C} ≈
                zeroM {As = As ++ Bs} {B = C}
  right-zeros {B = B} g =
    g >>= maybe (const zeroM) zeroM  ≈⟨ >>=M-cong (refl g) (λ _ → refl zeroM) (refl zeroM) ⟩
    g >>= const zeroM                ≈⟨ MoreLaws∥.>>=-return∥ g nothing ⟩
    zeroM                            ∎

  right-identities : ∀ {As} {B : Set ℓ} (g : As ⇨M B) →
                     g then returnM else zeroM ≈ g
  right-identities {B = B} g =
    g >>= maybe returnM zeroM  ≈⟨ >>=M-cong (refl g) (refl ∘ returnM) (refl zeroM) ⟩
    g >>= return               ≈⟨ MonadLaws.right-identity g ⟩
    g                          ∎

  distrib : ∀ {As Bs Cs} {B C D : Set ℓ}
            (g : As ⇨M B) (s : B → Bs ⇨M C) {f : Bs ⇨M C}
            (h : Maybe C → Cs ⇨M D) →
            (g then s else f) >>= h ≈
            g then (λ x → s x >>= h) else (f >>= h)
  distrib g s {f} h =
    (g >>= maybe s f) >>= h                  ≈⟨ MonadLaws.assoc g (maybe s f) h ⟩
    g >>= (λ x → maybe′ s f x >>= h)         ≈⟨ >>=M-cong (refl g) (λ x → refl (s x >>= h)) (refl (f >>= h)) ⟩
    g >>= maybe (λ x → s x >>= h) (f >>= h)  ∎

  private
    _∣_ : {A : Set ℓ} → Maybe A → Maybe A → Maybe A
    _∣_ = RawMonadPlus._∣_ Maybe.monadPlus

  left-identity : ∀ {As} {B : Set ℓ} (f : As ⇨M B) → zeroM ⊕ f ≈ f
  left-identity f =
    zeroM ⊕ f            ≈⟨ MoreLaws∥.left-identity _ _ f ⟩
    map (_∣_ nothing) f  ≈⟨ MoreLaws∥.map-id _ ⟩
    f                    ∎

  right-identity : ∀ {As} {B : Set ℓ} (f : As ⇨M B) → f ⊕ zeroM ≈ f
  right-identity f =
    f ⊕ zeroM                  ≈⟨ MoreLaws∥.right-identity _ f _ ⟩
    map (λ x → x ∣ nothing) f  ≈⟨ refl f >>=∥-cong′ lemma ⟩
    map id                  f  ≈⟨ MoreLaws∥.map-id _ ⟩
    f                          ∎
    where
    open MonadLaws∥

    lemma : ∀ x → return∥ (x ∣ nothing) ≈ return∥ x
    lemma (just x) = refl (return∥ (just x))
    lemma nothing  = refl zeroM

  assoc : ∀ {As} {B : Set ℓ} (f g h : As ⇨M B) →
          f ⊕ (g ⊕ h) ≈ (f ⊕ g) ⊕ h
  assoc f g h = sym $ MoreLaws∥.assoc f g h lemma
    where
    lemma : ∀ x y z → (x ∣ y) ∣ z ≡ x ∣ (y ∣ z)
    lemma nothing  y z = P.refl
    lemma (just x) y z = P.refl

  -- In the nullary case _⊕_ can be expressed concisely in terms of
  -- _then_else_.

  ⊕₀ : ∀ {B : Set ℓ} (f g : [] ⇨M B) →
       f ⊕ g ≈ f then returnM else g
  ⊕₀ (return (just x)) (return _) = refl _
  ⊕₀ (return nothing)  (return y) = refl _

  -- Another lemma which applies in the nullary case.

  then-else-·₀ :
    ∀ {Bs} {A B C : Set ℓ}
    (g : [] ⇨M A) (s : A → B ∷ Bs ⇨M C) {f : B ∷ Bs ⇨M C} {x} →
    (g then s else f) · x ≡ g then (λ y → s y · x) else f · x
  then-else-·₀ (return (just x)) _ = P.refl
  then-else-·₀ (return nothing)  _ = P.refl

-- Monoid-like operators.

identityM : ∀ {ℓ} {A : Set ℓ} → A ∷ [] ⇨M A
identityM = Λ λ x → returnM x

infixr  _⊙M_

_⊙M_ : ∀ {ℓ As Bs} {B C : Set ℓ} →
       As ⇨M B → (B ∷ Bs) ⇨M C → As ++ Bs ⇨M C
f ⊙M g = f then _·_ g else zeroM

-- Laws for the monoid-like operators.

module MonoidLawsM {ℓ} where

  open ≈-Reasoning

  right-identity : ∀ {As} {B : Set ℓ}
                   (f : As ⇨M B) → f ⊙M identityM ≈ f
  right-identity f =
    f ⊙M identityM            ≡⟨ P.refl ⟩
    f then returnM else zeroM ≈⟨ LawsM.right-identities f ⟩
    f                         ∎

  -- Note that the type of f is not completely general.

  left-identity : ∀ {As} {A B : Set ℓ}
                  (f : A ∷ As ⇨M B) → identityM ⊙M f ≈ f
  left-identity f =
    identityM ⊙M f                             ≡⟨ P.refl ⟩
    (Λ λ x → returnM x) then _·_ f else zeroM  ≡⟨ P.refl ⟩
    (Λ λ x → returnM x then _·_ f else zeroM)  ≡⟨ P.refl ⟩
    (Λ λ x → f · x)                            ≡⟨ P.refl ⟩
    f                                          ∎

  -- Note that the types of g and h are not completely general.

  assoc : ∀ {As Bs Cs} {B C D : Set ℓ}
          (f : As ⇨M B) (g : B ∷ Bs ⇨M C) (h : C ∷ Cs ⇨M D) →
          (f ⊙M g) ⊙M h ≈ f ⊙M (g ⊙M h)
  assoc {Bs = Bs} f (Λ g) h =
    (f ⊙M Λ g) ⊙M h                                                  ≡⟨ P.refl ⟩
    (f then g else zeroM) then _·_ h else zeroM                      ≈⟨ MonadLaws.assoc f _ _ ⟩
    f >>= (λ x → maybe′ g zeroM x then _·_ h else zeroM)             ≈⟨ LawsM.>>=M-cong (refl f) (λ _ → refl _)
                                                                                        (LawsM.left-zero Bs (_·_ h) zeroM) ⟩
    f >>= maybe (λ x → g x then _·_ h else zeroM) (const⋆ Bs zeroM)  ≈⟨ LawsM.>>=M-cong (refl f) (λ _ → refl _)
                                                                                        (MoreLaws∥.const⋆-return∥ Bs nothing) ⟩
    f >>= maybe (λ x → g x then _·_ h else zeroM) zeroM              ≡⟨ P.refl ⟩
    f ⊙M (Λ g ⊙M h)                                                  ∎

  infixr  _⊙M-cong_

  _⊙M-cong_ : ∀ {As₁ As₂ Bs} {B C : Set ℓ}
                {f₁ : As₁ ⇨M B} {g₁ : B ∷ Bs ⇨M C}
                {f₂ : As₂ ⇨M B} {g₂ : B ∷ Bs ⇨M C} →
              f₁ ≈ f₂ → g₁ ≈ g₂ → f₁ ⊙M g₁ ≈ f₂ ⊙M g₂
  f₁≈f₂ ⊙M-cong g₁≈g₂ =
    f₁≈f₂ then (λ x → g₁≈g₂ ·-cong H.refl) else-cong zeroM-cong
    where open LawsM

  identityM-cong : {A : Set ℓ} → identityM {A = A} ≈ identityM {A = A}
  identityM-cong = refl identityM