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

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

module ReverseComposition where

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 5 _→̂_

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

-- Application.

infixl 10 _·_

field _·_ : A → B

open _→̂_ public

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

infixr 4 _⇨_

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

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

infix 4 _≈_

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  3 _∎
infixr 2 _≈⟨_⟩_ _≡⟨_⟩_

_≈⟨_⟩_ : ∀ {ℓ 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 1 ▷_

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 10 _·-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

infixl 6 _>>=_

_>>=_ : ∀ {ℓ 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.

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 6 _>>=-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 9 _⊙_

_⊙_ : ∀ {ℓ 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 ≈-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 9 _⊙-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 5 _↑_⇨_

_↑_⇨_ : ∀ {ℓ} → 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 6 _>>=∥_

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

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 6 _>>=∥-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 9 _⊙∥_

_⊙∥_ : ∀ {ℓ 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 10 _⊛∥_

_⊛∥_ : ∀ {ℓ 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-∘ : ∀ {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                         ∎

-- 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                                       ∎

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))               ∎

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                          ∎

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

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                                                             ∎

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                                                        ∎

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 9 _⊙∥-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₂

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)

infixl 10 _⊛∥-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₂

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₂

-- 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                    ∎

------------------------------------------------------------------------
-- 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 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 4 _⇨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 6 _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 8 _⊕_

_⊕_ : ∀ {ℓ As} {B : Set ℓ} → As ⇨M B → As ⇨M B → As ⇨M B

-- 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

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 6 _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₂

zeroM-cong : ∀ {As} {B : Set ℓ} →
zeroM {As = As} {B = B} ≈ zeroM {As = As} {B = B}

infixr 8 _⊕-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

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

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 9 _⊙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 9 _⊙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
```