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