-- Code supplement for Case Study on Lambda Terms
-- (C) 2020 Patricia Johann, Andrew Polonsky
-- This information is free; you can redistribute and/or modify it under the terms of CC BY-SA 4.0.

module Lambda where

open import Data.Maybe renaming (map to Mmap)
open import Function

data Lam (A : Set) : Set where
  var : A → Lam A
  app : Lam A → Lam A → Lam A
  abs : Lam (Maybe A) → Lam A

-- Induction, Fold, Map

Lind : ∀ {l} (F : ∀ (X : Set) → Lam X → Set l)
         → (Fvar : ∀ {X} (x : X) → F X (var x))
         → (Fapp : ∀ {X} {s : Lam X} {t : Lam X} → F X s → F X t → F X (app s t))
         → (Fabs : ∀ {X} {r : Lam (Maybe X)} → F (Maybe X) r → F X (abs r))
         → ∀ A → ∀ (t : Lam A) → F A t
Lind F Fvar Fapp Fabs A (var a) = Fvar a
Lind F Fvar Fapp Fabs A (app t₁ t₂) = Fapp (Lind F Fvar Fapp Fabs A t₁)
                                           (Lind F Fvar Fapp Fabs A t₂)
Lind F Fvar Fapp Fabs A (abs t) = Fabs (Lind F Fvar Fapp Fabs (Maybe A) t)

Lfold : ∀ {l} (F : Set → Set l)
        → (∀ {X} → X → F X)
        → (∀ {X} → F X → F X → F X)
        → (∀ {X} → F (Maybe X) → F X)
        → ∀ X → Lam X → F X
Lfold F Fvar Fapp Fabs = Lind (λ X _ → F X) Fvar Fapp Fabs

Lmap : ∀ {A B : Set} → (A → B) → Lam A → Lam B
Lmap {A} {B} f s = Lfold (λ X → ∀ Y → (X → Y) → Lam Y)
                        ((λ x   Y g → var (g x)))
                        ((λ p q Y g → app (p Y g) (q Y g)))
                        ((λ   r Y g → abs (r (Maybe Y) (Mmap g))))
                        A s B f

-- Distributive law from Maybe to Lam
ML⇒LM : ∀ {A} → Maybe (Lam A) → Lam (Maybe A)
ML⇒LM = maybe (Lmap just) (var nothing)

-- Monad structure on the Lam datatype, given in terms of bind
Lbind : ∀ A (t : Lam A) → ∀ Y → (A → Lam Y) → Lam Y
Lbind = Lind (λ X _ → ∀ (Y : Set) → (X → Lam Y) → Lam Y)
             (λ x   Y f → f x)
             (λ p q Y f → app (p Y f) (q Y f))
             (λ r   Y f → abs (r (Maybe Y) (ML⇒LM ∘ Mmap f)))

-- Multiplication derived from bind
LL⇒L : ∀ {A} → Lam (Lam A) → Lam A
LL⇒L {A} t = Lbind (Lam A) t A id

module EqualityAndCongruences where

  -- standard facts about equality

  data _≡_ {A : Set} : A → A → Set where
    refl : ∀ a → a ≡ a

  infix 5 _≡_

  tran : ∀ {A} {x y z : A} → x ≡ y → y ≡ z → x ≡ z
  tran (refl a) (refl .a) = refl a

  var-≡ : ∀ {A} {a1 a2 : A} → a1 ≡ a2 → var a1 ≡ var a2
  var-≡ {A} {.a} {.a} (refl a) = refl (var a)

  app-≡ : ∀ {A} {s1 s2 : Lam A} (s* : s1 ≡ s2) {t1 t2} (t* : t1 ≡ t2) → app s1 t1 ≡ app s2 t2
  app-≡ {A} {s1} {.s1} (refl .s1) {t1} {.t1} (refl .t1) = refl _

  abs-≡ : ∀ {A} {r1 r2 : Lam (Maybe A)} → r1 ≡ r2 → abs r1 ≡ abs r2
  abs-≡ {A} {r} {.r} (refl .r) = refl _

  just-≡ : ∀ {A} {x y : A} → x ≡ y → _≡_ {Maybe A} (just x) (just y)
  just-≡ {A} {.a} {.a} (refl a) = refl _

  symm : ∀ {A} {x y : A} → x ≡ y → y ≡ x
  symm (refl a) = refl a

  _≃_ : ∀ {A B : Set} → (A → B) → (A → B) → Set
  f ≃ g = ∀ x → f x ≡ g x

  infix 5 _≃_

  fun-refl : ∀ {A B : Set} → (f : A → B) → f ≃ f
  fun-refl f x = refl (f x)

  fun-tran : ∀ {A B : Set} {f g h : A → B} → f ≃ g → g ≃ h → f ≃ h
  fun-tran p q x = tran (p x) (q x)

  fun-symm : ∀ {A B : Set} {f g : A → B} → f ≃ g → g ≃ f
  fun-symm p = symm ∘ p

  fun-app : ∀ {A B : Set} {f g : A → B} → f ≃ g → ∀ {x y : A} → x ≡ y → f x ≡ g y
  fun-app e (refl a) = e a

  ∘-≃ : ∀ {A B C : Set} {f1 f2 : B → C} {g1 g2 : A → B}
           →  f1 ≃ f2  →  g1 ≃ g2  →  f1 ∘ g1 ≃ f2 ∘ g2
  ∘-≃ f* g* = fun-app f* ∘ g*

  -- functoriality properties of Maybe

  Mmap-≃ : ∀ {A B : Set} {f g : A → B} → f ≃ g → Mmap f ≃ Mmap g
  Mmap-≃ e (just x) = just-≡ (e x)
  Mmap-≃ e nothing = refl _

  Mmap-id : ∀ {A : Set} → Mmap {A = A} id ≃ id
  Mmap-id {A} (just x) = refl _
  Mmap-id {A} nothing = refl _

  Mmap-∘ : ∀ {A B C : Set} (f : B → C) (g : A → B) → Mmap (f ∘ g) ≃ Mmap f ∘ Mmap g
  Mmap-∘ f g (just x) = refl _
  Mmap-∘ f g nothing = refl _

  Mmap-∘-≃ : ∀ {A B1 B2 C : Set} {f1 : B1 → C} {f2 : A → B1} {g1 : B2 → C} {g2 : A → B2}
                 →  f1 ∘ f2 ≃ g1 ∘ g2  →  Mmap f1 ∘ Mmap f2 ≃ Mmap g1 ∘ Mmap g2
  Mmap-∘-≃ h = fun-tran (fun-symm (Mmap-∘ _ _)) (fun-tran (Mmap-≃ h) (Mmap-∘ _ _))

  -- functoriality properties of Lam

  Lmap-≃ : ∀ {A B : Set} {f g : A → B} → f ≃ g → Lmap f ≃ Lmap g
  Lmap-≃ p t = Lind (λ X t → ∀ {Y} {f g : X → Y} → f ≃ g → Lmap f t ≡ Lmap g t)
                    (λ x   f≃g → var-≡ (f≃g x))
                    (λ p q f≃g → app-≡ (p f≃g) (q f≃g))
                    (λ r   f≃g → abs-≡ (r (Mmap-≃ f≃g)))
                    _ t p

  Lmap-id : ∀ {A} → Lmap {A} id ≃ id
  Lmap-id = Lind (λ X t → Lmap id t ≡ t)
                 (λ x         → var-≡ (refl x))
                 (λ p q       → app-≡ p q)
                 (λ {X} {u} r → abs-≡ (tran (Lmap-≃ Mmap-id u) r)) _

  Lmap-∘ : ∀ {A B C : Set} (f : B → C) (g : A → B) → Lmap (f ∘ g) ≃ Lmap f ∘ Lmap g
  Lmap-∘ f₀ g₀ t =
    Lind (λ X t → ∀ {Y Z} (f : Y → Z) (g : X → Y) → Lmap (f ∘ g) t ≡ Lmap f (Lmap g t))
         (λ x         → λ f g → var-≡ (refl (f (g x))))
         (λ p q       → λ f g → app-≡ (p f g) (q f g))
         (λ {X} {u} r → λ f g → abs-≡ (tran (Lmap-≃ (Mmap-∘ f g) u) (r (Mmap f) (Mmap g))))
         _ t f₀ g₀

  Lmap-∘-≃ : ∀ {A B1 B2 C : Set} {f1 : B1 → C} {f2 : A → B1} {g1 : B2 → C} {g2 : A → B2}
               →  f1 ∘ f2  ≃  g1 ∘ g2  →  Lmap f1 ∘ Lmap f2  ≃  Lmap g1 ∘ Lmap g2
  Lmap-∘-≃ h = fun-tran (fun-symm (Lmap-∘ _ _))
                        (fun-tran (Lmap-≃ h) (Lmap-∘ _ _))

open EqualityAndCongruences

Lfold-in : ∀ (A : Set) → Lam A → Lam A
Lfold-in = Lfold Lam var app abs

Lfold-in-id : ∀ (A : Set) (t : Lam A) → t ≡ Lfold-in A t
Lfold-in-id = Lind (λ X t → t ≡ Lfold-in X t)
                   (λ x → refl (var x))
                   (λ p q → app-≡ p q)
                   (λ r → abs-≡ r)

-- Naturality and other properties proved by induction

just-nat : ∀ {A B : Set} (f : A → B)  →  Mmap f ∘ just ≃ just ∘ f
just-nat f x = refl (just (f x))

ML⇒LM-nat : ∀ {A B : Set} (f : A → B) → Lmap (Mmap f) ∘ ML⇒LM  ≃  ML⇒LM ∘ Mmap (Lmap f)
ML⇒LM-nat {A} {B} f = maybe (Lmap-∘-≃ (just-nat f)) (refl (var nothing))

Lbind-≃ : ∀ {A} t {B} {f g : A → Lam B} → f ≃ g → Lbind A t B f ≡ Lbind A t B g
Lbind-≃ = Lind (λ A t → ∀ {B} {f g : A → Lam B} → f ≃ g → Lbind A t B f ≡ Lbind A t B g)
               (λ x e → e x)
               (λ p q e → app-≡ (p e) (q e))
               (λ r e → abs-≡ (r (∘-≃ (fun-refl ML⇒LM) (Mmap-≃ e)))) _

Lbind-≡ : ∀ {A B} {s t} (e : s ≡ t) f → Lbind A s B f ≡ Lbind A t B f
Lbind-≡ {A} {B} {.a} {.a} (refl a) f = refl _

Lbind-nat-pred : ∀ X → Lam X → Set₁
Lbind-nat-pred X t = ∀ {Y Z : Set} (f : X → Y) (g : Y → Lam Z)
                     → Lbind Y (Lmap f t) Z g ≡ Lbind X t Z (g ∘ f)

Lbind-nat : ∀ {A} (t : Lam A) → Lbind-nat-pred A t
Lbind-nat = Lind Lbind-nat-pred
                 (λ x         u v → refl _)
                 (λ       p q u v → app-≡ (p u v) (q u v))
                 (λ {X} {t} p u v → abs-≡ (tran (p (Mmap u) (ML⇒LM ∘ Mmap v))
                    (Lbind-≃ t (∘-≃ (fun-refl ML⇒LM) (fun-symm (Mmap-∘ v u)))))) _

ML⇒LM-var : ∀ {A} → ML⇒LM ∘ Mmap var ≃ id ∘ var {Maybe A}
ML⇒LM-var (just x) = refl _
ML⇒LM-var nothing = refl _

Lbind-map : ∀ {A} {B} (f : A → B) (t : Lam A) → Lbind (Lam A) (Lmap var t) B (Lmap f) ≡ Lmap f t
Lbind-map {A} {B} f₀ t₀ =
  Lind (λ X t → ∀ Y (f : X → Y) → Lbind (Lam X) (Lmap var t) Y (Lmap f) ≡ Lmap f t)
       (λ x Y f → refl (var (f x)))
       (λ {X} {s} {t} p q Y f → app-≡ (p Y f) (q Y f))
       (λ {X} {u} r Y f → abs-≡ (tran (tran (Lbind-nat u (Mmap var) (ML⇒LM ∘ Mmap (Lmap f)))
                                            (Lbind-≃ u λ x →
                                                tran (∘-≃ (fun-refl ML⇒LM)
                                                          (Mmap-∘-≃ (λ z → refl (var (f z)))) x)
                                                     (ML⇒LM-var (Mmap f x)) ))
                                      (tran (symm (Lbind-nat u var (Lmap (Mmap f))))
                                            (r (Maybe Y) (Mmap f)))))
       A t₀ B f₀

Lbind-var : ∀ A t → Lbind (Lam A) (Lmap var t) A id ≡ t
Lbind-var A t = tran (Lbind-≃ (Lmap var t) (fun-symm Lmap-id))
                     (tran (Lbind-map id t) (Lmap-id t))

-- the first triangle (unit) law
Lam-mult-unit : ∀ {A : Set} → (t : Lam A) → LL⇒L (var t) ≡ t
Lam-mult-unit _ = refl _

-- the second triangle (unit) law
Lam-unit-mult : ∀ {A : Set} → (t : Lam A) → LL⇒L (Lmap var t) ≡ t
Lam-unit-mult = Lbind-var _

LM-func : ∀ {A B C : Set} (f : B → C) (g : A → B)
          → Lmap (Mmap (f ∘ g)) ≃  Lmap (Mmap f) ∘ Lmap (Mmap g)
LM-func f g = fun-tran (Lmap-≃ (Mmap-∘ f g)) (Lmap-∘ (Mmap f) (Mmap g))

L-ML⇒LM-nat : ∀ {A B : Set} (f : A → B)
              → Lmap (Lmap (Mmap f)) ∘ Lmap ML⇒LM  ≃  Lmap ML⇒LM ∘ Lmap (Mmap (Lmap f))
L-ML⇒LM-nat = Lmap-∘-≃ ∘ ML⇒LM-nat

Lbind-nat₂-pred : ∀ X → Lam X → Set₁
Lbind-nat₂-pred X t = ∀ {Y Z : Set} (f : X → Lam Y) (g : Y → Z)
                      → Lbind X t Z (Lmap g ∘ f) ≡ Lmap g (Lbind X t Y f)

Lbind-nat₂ : ∀ {A} t → Lbind-nat₂-pred A t
Lbind-nat₂ = Lind Lbind-nat₂-pred
  (λ         x u v → refl _)
  (λ       p q u v → app-≡ (p u v) (q u v))
  (λ {X} {r} p u v → let e x = tran (∘-≃ (fun-refl ML⇒LM) (Mmap-∘ (Lmap v) u) x)
                                    (symm (ML⇒LM-nat v (Mmap u x)))
                      in abs-≡ (tran (Lbind-≃ r e) (p (ML⇒LM ∘ Mmap u) (Mmap v)))) _
--
LL⇒L-nat : ∀ {A B : Set} (f : A → B) → LL⇒L ∘ Lmap (Lmap f) ≃ Lmap f ∘ LL⇒L
LL⇒L-nat {A} {B} f z = tran (Lbind-nat z (Lmap f) id) (Lbind-nat₂ z id f)

just-ML⇒LM : ∀ {A} → Lmap {A} just ≃ (ML⇒LM ∘ just)
just-ML⇒LM = Lind (λ X t → Lmap just t ≡ ML⇒LM (just t))
                  (λ _ → refl _)
                  (λ _ _ → refl _)
                  (λ _ → refl _) _

-- the main lemma concerns interaction between multiplication and distributive law

LL⇒L-ML⇒LM : ∀ {A} → ML⇒LM {A} ∘ Mmap LL⇒L ≃ LL⇒L ∘ Lmap ML⇒LM ∘ ML⇒LM
LL⇒L-ML⇒LM {A} = maybe (λ x → tran (symm (LL⇒L-nat just x))
                                   (∘-≃ (fun-refl LL⇒L) (Lmap-∘ ML⇒LM just) x))
                       (refl (var nothing))

Lbind-aux : ∀ (X Y Z : Set) (f : Y → Lam Z) (g : X → Lam Y) (m : Maybe X)
            → Lbind (Maybe Y) (ML⇒LM (Mmap g m)) (Maybe Z) (ML⇒LM ∘ Mmap f)
            ≡ ML⇒LM (Mmap (λ a → Lbind Y (g a) Z f) m)
Lbind-aux X Y Z f g =
  maybe (λ x → tran  (Lbind-nat (g x) just _)
                     (Lbind-nat₂ (g x) f just))
        (refl (var nothing))

Lbind-comm-pred : ∀ X → Lam X → Set₁
Lbind-comm-pred X t = ∀ Y Z (f : Y → Lam Z) (g : X → Lam Y)
                      → Lbind Y (Lbind X t Y g) Z f ≡ Lbind X t Z (λ a → Lbind Y (g a) Z f)
Lbind-comm : ∀ {A B C} (f : B → Lam C) (g : A → Lam B) (t : Lam A)
             → Lbind B (Lbind A t B g) C f ≡ Lbind A t C (λ a → Lbind B (g a) C f)

Lbind-comm {A} {B} {C} f₀ g₀ t₀ =
  Lind Lbind-comm-pred
       (λ x               Y Z f g → refl (Lbind Y (g _) Z f))
       (λ {X} {s} {t} p q Y Z f g → app-≡ (p Y Z f g) (q Y Z f g))
       (λ {X} {u} r       Y Z f g →
         let e = (r (Maybe Y) (Maybe Z) (ML⇒LM ∘ Mmap f) (ML⇒LM ∘ Mmap g))
         in abs-≡ (tran e (Lbind-≃ u λ m → Lbind-aux X Y Z f g m)) )
       A t₀ B C f₀ g₀

-- associativity of multiplication

LL⇒L-assoc : ∀ {A} → LL⇒L {A} ∘ LL⇒L ≃ LL⇒L ∘ Lmap LL⇒L
LL⇒L-assoc {A} t = tran (Lbind-comm id id t)
                        (tran (refl _) (symm (Lbind-nat t LL⇒L id)))