```------------------------------------------------------------------------
-- Wand's generalised composition operator
------------------------------------------------------------------------

open import Relation.Binary.PropositionalEquality as PropEq
using (_≡_; refl; sym; cong; cong₂; subst)

module MultiComposition
-- Assume extensionality; it is used in proofs below.
(ext : forall {a b : Set} {f g : a -> b} ->
(forall x -> f x ≡ g x) -> f ≡ g)
where

open PropEq.≡-Reasoning
open import Relation.Binary.HeterogeneousEquality as HetEq
using (_≅_; refl)
open import Function
open import Data.Nat
open import Algebra
import Data.Nat.Properties as Nat
private
module CS = CommutativeSemiring Nat.commutativeSemiring

------------------------------------------------------------------------
-- Type lists

-- I am using difference lists to get associativity "for free". This
-- makes it possible to state the right-assoc lemma below without
-- using any casts.

record TypeList : Set1 where
infixr 5 _⟶_
field
_⟶_    : Set -> Set
map    : forall {a b} -> (a -> b) -> _⟶_ a -> _⟶_ b
map-id : forall a -> map (id {A = a}) ≡ id
map-∘  : forall {a b c} (f : b -> c) (g : a -> b) ->
map (f ∘ g) ≡ map f ∘ map g

-- (The map-id field is not used for anything, but it feels nice
-- to include it.)

open TypeList

-- The empty list.

ε : TypeList
ε = record
{ _⟶_    = \c -> c
; map    = \f as -> f as
; map-id = \a -> refl
; map-∘  = \f g -> refl
}

-- Cons.

infixr 5 _▻_

_▻_ : Set -> TypeList -> TypeList
a ▻ bs = record
{ _⟶_    = \c -> a -> bs ⟶ c
; map    = \f k x -> map bs f (k x)
; map-id = \a -> ext \k -> ext \x ->
cong (\h -> h (k x)) (map-id bs a)
; map-∘  = \f g -> ext \k -> ext \x ->
cong (\h -> h (k x)) (map-∘ bs f g)
}

-- Append.

infixr 5 _▻▻_

_▻▻_ : TypeList -> TypeList -> TypeList
as ▻▻ bs = record
{ _⟶_    = \c -> as ⟶ bs ⟶ c
; map    = \f k -> map as (map bs f) k
; map-id = \a -> begin
map as (map bs id)
≡⟨ cong (map as) (map-id bs a) ⟩
map as id
≡⟨ map-id as (bs ⟶ a) ⟩
id
∎
; map-∘  = \f g -> begin
map as (map bs (f ∘ g))
≡⟨ cong (map as) (map-∘ bs f g) ⟩
map as (map bs f ∘ map bs g)
≡⟨ map-∘ as (map bs f) (map bs g) ⟩
map as (map bs f) ∘ map as (map bs g)
∎
}

-- Replicate.

infix 8 _^_

_^_ : Set -> ℕ -> TypeList
a ^ zero  = ε
a ^ suc n = a ▻ a ^ n

------------------------------------------------------------------------
-- Composition

-- This is Wand's combinator "B_k":
--
-- B_k f g = \x₁ … x_k -> f (g x₁ … x_k)

_⟪_⟫_ : forall {a b} ->
(a -> b) -> (ts : TypeList) -> ts ⟶ a -> ts ⟶ b
f ⟪ ts ⟫ g = map ts f g

-- Congruence.

_⟪_⟫-cong_ : forall {a b} {f₁ f₂ : a -> b} ->
f₁ ≡ f₂ -> (ts : TypeList) ->
{g₁ g₂ : ts ⟶ a} -> g₁ ≡ g₂ ->
f₁ ⟪ ts ⟫ g₁ ≡ f₂ ⟪ ts ⟫ g₂
eqf ⟪ ts ⟫-cong eqg = cong₂ (map ts) eqf eqg

-- Right associativity.

right-assoc : forall {a b t} ts₁ ts₂
(f : a -> b) (g : (t ▻ ts₂) ⟶ a) (h : ts₁ ⟶ t) ->
(f ⟪ t   ▻  ts₂ ⟫  g) ⟪ ts₁ ⟫ h  ≡
f ⟪ ts₁ ▻▻ ts₂ ⟫ (g  ⟪ ts₁ ⟫ h)
right-assoc ts₁ ts₂ f g h =
cong (\F -> F h) (map-∘ ts₁ (map ts₂ f) g)

------------------------------------------------------------------------
-- Boring lemmas

-- Lemmas about _^_.

shift : forall {A B} n -> A ^ suc n ⟶ B ≡ A ^ n ⟶ (A -> B)
shift         zero    = PropEq.refl
shift {A = A} (suc n) = PropEq.cong (\t -> A -> t) (shift n)

shift' : forall {A B} {C : Set} n ->
(C -> A ^ suc n ⟶ B) ≡ (C -> A ^ n ⟶ (A -> B))
shift' {C = C} n = PropEq.cong (\t -> C -> t) (shift n)

merge : forall {A B} m n -> A ^ m ⟶ A ^ n ⟶ B ≡ A ^ (m + n) ⟶ B
merge         zero    n = PropEq.refl
merge {A = A} (suc m) n = PropEq.cong (\t -> A -> t) (merge m n)

swap : forall {A B} m n -> A ^ (m + n) ⟶ B ≡ A ^ (n + m) ⟶ B
swap {A} {B} m n = PropEq.cong (\i -> A ^ i ⟶ B) (CS.+-comm m n)

-- Cast.

cast : {A B : Set} → A ≡ B → A → B
cast PropEq.refl = id

-- Cast commutes with well-typed function applications.

cast-commutes : {A B C : Set}
(f : A → B) (x : A)
(eq₂ : B ≡ C)
(eq₁ : (A → B) ≡ (A → C)) →
cast eq₁ f x ≡ cast eq₂ (f x)
cast-commutes f x refl refl = refl

-- A variant of the right associativity result. This lemma is used to
-- prove rot-lemma below.

-- (Note that it is not possible to pattern match on an equality proof
-- unless the two sides of the equality can be unified. This is the
-- reason for formulating the lemma in this way.)

right-assoc' : forall {a b k c ana anb amnc} m n
(f : k -> c) (g : k -> b) (h : ana)
(eq₀ : b ≡ a ^ m ⟶ k)
(eq₁ : ana ≡ a ^ n ⟶ k)
(eq₂ : a ^ n ⟶ b ≡ anb)
(eq₃ : anb ≡ a ^ (m + n) ⟶ k)
(eq₄ : a ^ (m + n) ⟶ c ≡ amnc)
(eq₅ : (k -> b) ≡ (k -> a ^ m ⟶ k))
(eq₆ : ana ≡ a ^ n ⟶ k)
(eq₇ : a ^ n ⟶ a ^ m ⟶ c ≡ amnc) ->
cast eq₄ (f ⟪ a ^ (m + n) ⟫
cast eq₃ (cast eq₂
(g ⟪ a ^ n ⟫ cast eq₁ h))) ≡
cast eq₇ ((f ⟪ k ▻ a ^ m ⟫ cast eq₅ g)
⟪ a ^ n ⟫ cast eq₆ h)
right-assoc' {a = a} {k = k} m n f g h
refl refl refl eq₃ refl refl refl eq₇ =
begin
f ⟪ a ^ (m + n) ⟫ cast eq₃ (g ⟪ a ^ n ⟫ h)
≡⟨ lem₁ m n f (g ⟪ a ^ n ⟫ h) eq₃ eq₇ ⟩
cast eq₇ (f ⟪ a ^ n ▻▻ a ^ m ⟫ (g ⟪ a ^ n ⟫ h))
≡⟨ cong (cast eq₇) (sym \$ right-assoc (a ^ n) (a ^ m) f g h) ⟩
cast eq₇ ((f ⟪ k ▻ a ^ m ⟫ g) ⟪ a ^ n ⟫ h)
∎
where
lem₂ : ∀ {A B C} m n
(f : B → C) (g : A ^ n ⟶ A ^ m ⟶ B)
(eq₁ : A ^ n ⟶ A ^ m ⟶ B ≡ A ^ (n + m) ⟶ B)
(eq₂ : A ^ n ⟶ A ^ m ⟶ C ≡ A ^ (n + m) ⟶ C) →
f ⟪ A ^ (n + m) ⟫ cast eq₁ g ≡
cast eq₂ (map (A ^ m) f ⟪ A ^ n ⟫ g)
lem₂     m zero    f g refl refl = refl
lem₂ {A} m (suc n) f g eq₁  eq₂  = ext (λ x → begin
map (A ^ (n + m)) f (cast eq₁ g x)
≡⟨ cong (map (A ^ (n + m)) f)
(cast-commutes g x (merge n m) eq₁) ⟩
map (A ^ (n + m)) f (cast (merge n m) (g x))
≡⟨ lem₂ m n f (g x) (merge n m) (merge n m) ⟩
cast (merge n m) (map (A ^ n) (map (A ^ m) f) (g x))
≡⟨ sym (cast-commutes (map (A ^ n) (map (A ^ m) f) ∘ g) x
(merge n m) eq₂) ⟩
cast eq₂ (map (A ^ n) (map (A ^ m) f) ∘ g) x
∎)

lem₁ : ∀ {A B C} m n (f : B → C) (g : A ^ n ⟶ A ^ m ⟶ B)
(eq₁ : A ^ n ⟶ A ^ m ⟶ B ≡ A ^ (m + n) ⟶ B)
(eq₂ : A ^ n ⟶ A ^ m ⟶ C ≡ A ^ (m + n) ⟶ C) →
f ⟪ A ^ (m + n) ⟫ cast eq₁ g ≡
cast eq₂ (map (A ^ m) f ⟪ A ^ n ⟫ g)
lem₁ m n f g eq₁ eq₂ with m + n | eq₁ | eq₂ | CS.+-comm m n
lem₁ m n f g eq₁ eq₂ | .(n + m) | eq₁′ | eq₂′ | refl =
lem₂ m n f g eq₁′ eq₂′

------------------------------------------------------------------------
-- Variants of _⟪_⟫_

-- Used to establish the correctness of the rotation function (see
-- Wand.agda).

B′ : forall {v c} n -> let k = v -> c in
(k -> c) -> (k -> v ^ suc n ⟶ c) -> (k -> v ^ n ⟶ c)
B′ {v} {c} n f g = f ⟪ k ▻ v ^ n ⟫ cast₁ g
where
k     = v -> c
cast₁ = cast (shift' n)

B″ : forall {v c} m n ->
((v -> c) -> v ^ m ⟶ c) -> v ^ suc n ⟶ c -> v ^ (m + n) ⟶ c
B″ {v} m n f g = cast₂ (f ⟪ v ^ n ⟫ cast₁ g)
where
cast₁ = cast (shift n)
cast₂ = cast (PropEq.trans (merge n m) (swap n m))

-- We still have a form of right associativity.

rot-lemma : forall {v c} m n (f : (v -> c) -> c) g h ->
B″ 0 (m + n) f (B″ (suc m) n g h) ≡ B″ m n (B′ m f g) h
rot-lemma m n f g h =
right-assoc' m n f g h
(shift m)
(shift n)
(PropEq.trans (merge n (suc m)) (swap n (suc m)))
(shift (m + n))
(PropEq.trans (merge (m + n) 0) (swap (m + n) 0))
(shift' m)
(shift n)
(PropEq.trans (merge n m) (swap n m))
```