module BoundUniverse where

open import Function

open import Data.Unit
open import Data.Product
open import Data.Sum
open import Data.Nat
open import Data.Fin renaming (_+_ to _+F_)

open import Relation.Unary
open import Relation.Binary.PropositionalEquality


-- The categorical structures at play in this file are directly
-- inspired by the work of Hirschowitz & Maggesi's "Modules over
-- monads and initial semantics"
-- [http://web.math.unifi.it/users/maggesi/syn.pdf] and Ahrens &
-- Zsidó's "Initial Semantics for higher-order typed syntax in Coq"
-- [http://jfr.cib.unibo.it/article/view/2066]. I've merely adapted
-- their work to a relative monad, in a much less principled way than
-- Ahrens' "Modules over relative monads for syntax and semantics"
-- [http://arxiv.org/abs/1107.5252].


-- * Worlds

-- Worlds let us abstract over (non-dependent?) contexts (with W) and
-- indices into a context (with ⟦ Γ ⟧).

-- A World is a monoidal category (W, ⊗) equipped with a strict
-- monoidal functor ⟦_⟧ : W → Set/I

record World (I : Set)(W : Set) : Set₁ where
  field 
    -- W is monoidal:
    _⊗_ : W → W → W
    -- With a functor to Set/I:
    ⟦_⟧ : W → (I → Set)
    -- Which is strict:
    mon⟦⟧ : ∀{Ω Γ} → ⟦ Γ ⊗ Ω ⟧ ⊆ ⟦ Γ ⟧ ∪ ⟦ Ω ⟧
    mon⟦⟧⁻¹ : ∀{Ω Γ} → ⟦ Γ ⟧ ∪ ⟦ Ω ⟧ ⊆ ⟦ Γ ⊗ Ω ⟧

-- ⟦_⟧ being strict monoidal, we can map over the left-hand side of
-- the monoidal product (which we shall frequently do when going under
-- a bunch of binders):

module MapW {I W : Set}(w : World I W) where

  open World w
  
  mapL : ∀{Ω Γ Δ : W} → 
         (⟦ Γ ⟧ ⊆ ⟦ Δ ⟧) → ⟦ Γ ⊗ Ω ⟧ ⊆ ⟦ Δ ⊗ Ω ⟧
  mapL ρ k = mon⟦⟧⁻¹ (Data.Sum.map ρ id (mon⟦⟧ k))

-- Question: what is the relation with Tanaka & Power's "A unified
-- category-theoretic formulation of typed binding signatures"
-- [http://homepages.inf.ed.ac.uk/rpollack/TYPES/MERLIN_05/p13-tanaka.pdf]?

-- * Binding signature:

-- We give an intensional characterization of binding signatures with
-- a Martin-Löf universe. The idea is to reuse the standard machinery
-- for inductive families
-- [http://gallium.inria.fr/~pdagand/stuffs/thesis-2011-phd/model/html/Chapter5.IDesc.html]
-- (see also, Chapter 5 of my thesis
-- [http://gallium.inria.fr/~pdagand/stuffs/thesis-2011-phd/thesis.pdf]),
-- and extend the recursive constructor (here, `X) with a notion of
-- scope.

-- This should be (extensionally) equivalent to Ahrens' binding
-- signature. However, it supports first-order representations of
-- first-order datatypes: we get a working definitional equality over
-- these objects, as we shall see in the examples. This boils down to
-- the usual debate about the (inherently extensional) W-types
-- vs. more intensional presentations.

module BindingDesc {W I : Set}(w : World I W) where

  open World w

  -- Code:
  data BDesc : Set₁ where
    `⊤ : BDesc
    _`+_ _`×_ : (D₁ D₂ : BDesc) → BDesc
    `Σ `Π : (S : Set)(T : S → BDesc) → BDesc
    `X : (w : W)(i : I) → BDesc

  infixr  _`×_
  infixr  _`+_

  -- Alias for binders: Bind[ Ω ] X Γ i ≅ X (Γ ⊗ Ω) i
  data Bind[_] (Ω : W)(X : W → I → Set)(Γ : W)(i : I) : Set where
    sc : X (Γ ⊗ Ω) i → Bind[ Ω ] X Γ i

  -- Interpretation:
  ⟦_⟧B : BDesc → (W → I → Set) → W → Set
  ⟦ `⊤ ⟧B X w = ⊤
  ⟦ D₁ `+ D₂ ⟧B X w = ⟦ D₁ ⟧B X w ⊎ ⟦ D₂ ⟧B X w
  ⟦ D₁ `× D₂ ⟧B X w = ⟦ D₁ ⟧B X w × ⟦ D₂ ⟧B X w
  ⟦ `Σ S T ⟧B X w = Σ[ s ∈ S ] ⟦ T s ⟧B X w
  ⟦ `Π S T ⟧B X w = (s : S) → ⟦ T s ⟧B X w 
  ⟦ `X Ω i ⟧B X Γ = Bind[ Ω ] X Γ i 

  -- A binding signature is an I-indexed family of BDesc:
  record BSig : Set₁ where
    constructor mk
    field
      out : I → BDesc


-- * Free term algebra

-- From a presentation, we can build the syntax of terms (defined by
-- Tm) with bindings (defined by Scope[_]).

module Terms {W I : Set}{w : World I W}
             (D : BindingDesc.BSig w)  where

  open World w
  open BindingDesc w

  mutual
    -- Term of type I in context W:
    data Tm (w : W) : I → Set where
      `var : ∀{i} → ⟦ w ⟧ i → Tm w i
      `op : ∀{i} → ⟦ BSig.out D i ⟧B Tm w → Tm w i

    -- Binder of Ω-variables :
    Scope[_] : (Ω : W)(Γ : W)(i : I) → Set
    Scope[ Ω ] = Bind[ Ω ] Tm

  -- A relative strength over Tm (provided that Tm is a functor, which
  -- we will define) is defined as:

  Strength : Set 
  Strength = {Γ Δ : W} → (∀{Γ Δ} → (⟦ Γ ⟧ ⊆ ⟦ Δ ⟧) → Tm Γ ⊆ Tm Δ) → 
             Tm Γ ∪ ⟦ Δ ⟧ ⊆ Tm (Γ ⊗ Δ)


-- * Substitution structure

module Subst {I W : Set}{w : World I W}
             {D : BindingDesc.BSig w}(str : Terms.Strength D) where

  open World w
  open MapW w 
  open BindingDesc w
  open Terms D

  mutual

    -- Scope[ Ω ] : W → Set/I is a functor
    --   ie. for f ∈ W(Γ, Δ)
    --         we have a functorial map
    --           mapSc f ∈ Set/I(Scope Γ, Scope Δ)

    mapSc : ∀{Ω Γ Δ} →
            (⟦ Γ ⟧ ⊆ ⟦ Δ ⟧) → Scope[ Ω ] Γ ⊆ Scope[ Ω ] Δ
    mapSc σ (sc xs) = sc (mapTm (mapL σ) xs)

    -- Tm : W → Set/I is a functor
    --   ie. for f ∈ W(Γ, Δ), 
    --         we have a functorial map 
    --           mapTm f ∈ Set/I(Tm Γ, Tm Δ)

    mapTm : ∀{Γ Δ} → (⟦ Γ ⟧ ⊆ ⟦ Δ ⟧) → Tm Γ ⊆ Tm Δ
    mapTm σ {i} (`var x₁) = `var (σ {i} x₁)
    mapTm σ (`op x₁) = `op (⟦ BSig.out D _ ⟧map (λ {i} → σ {i}) x₁)
      where  ⟦_⟧map : ∀(D' : BDesc){Γ Δ} →
                      (⟦ Γ ⟧ ⊆ ⟦ Δ ⟧) → ⟦ D' ⟧B Tm Γ → ⟦ D' ⟧B Tm Δ
             ⟦ `⊤ ⟧map σ xs = tt
             ⟦ D₁ `× D₂ ⟧map σ (xs₁ , xs₂) = ⟦ D₁ ⟧map σ xs₁ , ⟦ D₂ ⟧map σ xs₂
             ⟦ D₁ `+ D₂ ⟧map σ (inj₁ xs₁) = inj₁ (⟦ D₁ ⟧map σ xs₁)
             ⟦ D₁ `+ D₂ ⟧map σ (inj₂ xs₂) = inj₂ (⟦ D₂ ⟧map σ xs₂)
             ⟦ `Σ S T ⟧map σ (s , xs) = s  , ⟦ T s ⟧map (λ {i} → σ {i}) xs
             ⟦ `Π S T ⟧map σ xs s = ⟦ T s ⟧map (λ {i} → σ {i}) (xs s)
             ⟦ `X w i ⟧map {Γ}{Δ} σ xs = mapSc σ xs

  -- Tm Σ : W → Set/I is a relative monad over ⟦_⟧ : W → Set/I
  --   ie. 
  --    1. we have a natural transformation
  --         return : ∀ Γ. Set/I(⟦ Γ ⟧, Tm Γ)

  return : ∀{Γ} → ⟦ Γ ⟧ ⊆ Tm Γ
  return = `var

  mutual
  --    2. we have a natural transformation
  --         sub : ∀ Γ Δ. Set/I(⟦ Γ ⟧, Tm Δ) → Set/I(Tm Γ, Tm Δ)

    sub : ∀{Γ Δ} → (⟦ Γ ⟧ ⊆ Tm Δ) → Tm Γ ⊆ Tm Δ
    sub ρ (`var x) = ρ x
    sub ρ (`op x) = `op (help {BSig.out D _} x ρ)
      where help : ∀{D' Γ Δ} → ⟦ D' ⟧B Tm Γ → (⟦ Γ ⟧ ⊆ Tm Δ) → ⟦ D' ⟧B Tm Δ
            help {`⊤} tt ρ = tt
            help {D₁ `× D₂} (xs₁ , xs₂) ρ = help {D₁} xs₁ ρ , help {D₂} xs₂ ρ
            help {D₁ `+ D₂} (inj₁ xs₁) ρ = inj₁ (help {D₁} xs₁ ρ)
            help {D₁ `+ D₂} (inj₂ xs₂) ρ = inj₂ (help {D₂} xs₂ ρ)
            help {`Σ S T} (s , xs) ρ = s , help {T s} xs ρ
            help {`Π S T} xs ρ s = help {T s} (xs s) ρ
            help {`X w i}{Γ} v ρ = weaken ρ v

  -- Scope[ Ω ] : W → Set/I is a Tm-module
  --   ie. we have a natural transformation
  --         scopeSub : ∀ Γ Δ. Set/I(⟦ Γ ⟧, Tm Δ) → Set/I(Scope[ Ω ] Γ, Scope[ Ω ] Δ)

    weaken : ∀{Ω Γ Δ} → (⟦ Γ ⟧ ⊆ Tm Δ) → Scope[ Ω ] Γ ⊆ Scope[ Ω ] Δ
    weaken {Ω}{Γ}{Δ} ρ {i} (sc tm) = sc (sub (λ v → str mapTm (Data.Sum.map ρ id (mon⟦⟧ v))) tm)

-- * Example: Well-scoped, untyped lambda-calculus

module LC where

  -- ** Untyped world

  monFin : ∀{n m} → Fin (n + m) → Fin m ⊎ Fin n
  monFin {zero} xs = inj₁ xs
  monFin {suc m} zero = inj₂ zero
  monFin {suc m} (suc xs) with monFin {m} xs
  monFin {suc m} (suc xs) | inj₁ x = inj₁ x
  monFin {suc m} (suc xs) | inj₂ y = inj₂ (suc y)

  monFin⁻¹ : ∀{n m} → Fin m ⊎ Fin n → Fin (n + m)
  monFin⁻¹ {n}{m} k = [ raise n , inject+ m ]′ k

  ℕW : World ⊤ ℕ
  ℕW = record { ⟦_⟧ = λ n tt → Fin n 
              ; _⊗_ = λ m n → n + m 
              ; mon⟦⟧ = monFin
              ; mon⟦⟧⁻¹ = monFin⁻¹ }

  open BindingDesc ℕW

  -- ** Terms

  Σλam : BSig
  Σλam = mk (λ tt → {- app: -} (`X  tt `× `X  tt) 
                 `+ {- abs: -} `X  tt )

  open Terms Σλam

  λam : ℕ → Set
  λam n = Tm n tt

  var : ∀{n} → Fin n → λam n
  var = `var

  app : ∀{n} → λam n → λam n → λam n
  app f s = `op (inj₁ (sc f , sc s))

  abs : ∀{n} → λam (suc n) → λam n 
  abs b = `op (inj₂ (sc b))


  -- ** Substitution

  str : Terms.Strength Σλam
  str {Δ = n} fmap (inj₁ tm) = fmap (raise n) tm
  str fmap (inj₂ k) = var (inject+ _ k)

  open Subst str

  subst⊢ : ∀{m n} → λam m → (Fin m → λam n) → λam n
  subst⊢ tm ρ = sub ρ tm

  -- ** Example

  tm1 : λam 
  tm1 = abs (var zero)

  tm2 : λam 
  tm2 = abs (app (var (suc zero)) (var zero))

  tm3 : λam 
  tm3 = subst⊢ tm2 ρ
    where ρ : Fin  → λam 
          ρ zero = tm1
          ρ (suc ())

  tm3' : λam 
  tm3' = abs (app (abs (var zero)) (var zero))

  test : tm3 ≡ tm3'
  test = refl  -- Definitional equality.

-- * Example: well-typed lambda-calculus

module STLC where

  data Typ : Set where
    `■ : Typ
    _`⇒_ : (A B : Typ) → Typ

  -- ** Simply-typed world

  data Ctxt : Set where
    ε : Ctxt 
    _▹_ : Ctxt → Typ → Ctxt
    
  data _`∈_ (T : Typ) : (Γ : Ctxt) → Set where
    ze : ∀{Γ} → T `∈ (Γ ▹ T)
    suc : ∀{Γ T'} → T `∈ Γ → T `∈ (Γ ▹ T') 

  _∷_ : Ctxt → Ctxt → Ctxt
  Γ ∷ ε = Γ
  Γ ∷ (Ω ▹ S) = (Γ ∷ Ω) ▹ S

  monW : ∀{Ω Γ T} → T `∈ (Γ ∷ Ω) → T `∈ Γ ⊎ T `∈ Ω
  monW {ε} k = inj₁ k
  monW {Ω ▹ T} ze = inj₂ ze
  monW {Ω ▹ x} (suc k) with monW {Ω} k 
  monW {Ω ▹ x₁} (suc k) | inj₁ x = inj₁ x
  monW {Ω ▹ x} (suc k) | inj₂ y = inj₂ (suc y)

  raiseW : ∀ {Γ T} Δ → T `∈ Γ → T `∈ (Γ ∷ Δ)
  raiseW ε k = k
  raiseW (Δ ▹ x) k = suc (raiseW Δ k)
  
  injectW : ∀ {Γ T} Δ → T `∈ Γ → T `∈ (Δ ∷ Γ)
  injectW Δ ze = ze
  injectW Δ (suc k) = suc (injectW Δ k)

  monW⁻¹ : ∀{Ω Γ T} → T `∈ Γ ⊎ T `∈ Ω → T `∈ (Γ ∷ Ω)
  monW⁻¹ {Ω}{Γ} = [ raiseW Ω , injectW Γ ]′ 

  CtxtW : World Typ Ctxt
  CtxtW = record { ⟦_⟧ = λ Γ T → T `∈ Γ 
                 ; _⊗_ = _∷_ 
                 ; mon⟦⟧ = monW
                 ; mon⟦⟧⁻¹ = monW⁻¹ }

  -- ** Terms

  open BindingDesc CtxtW

  Σλam : BSig
  Σλam = mk λ T →  {- app: -} (`Σ Typ λ S → `X ε (S `⇒ T) `× `X ε S)
                `+ {- abs: -} (`Σ Typ λ A → `Σ Typ λ B → `Σ (T ≡ A `⇒ B) λ _ → `X (ε ▹ A) B)

  open Terms Σλam

  _⊢_ : (Γ : Ctxt) → Typ → Set
  Γ ⊢ T = Tm Γ T

  var : ∀{Γ T} → T `∈ Γ → Γ ⊢ T
  var k = `var k

  app : ∀{Γ S T} → Γ ⊢ (S `⇒ T) → Γ ⊢ S → Γ ⊢ T
  app {S = S} f s = `op (inj₁ (S , sc f , sc s))

  abs : ∀{Γ S T} → (Γ ▹ S) ⊢ T → Γ ⊢ (S `⇒ T)
  abs {S = S}{T = T} b = `op (inj₂ (S , T , refl , sc b))

  -- ** Substitution

  str : Terms.Strength Σλam
  str {Δ = Δ} fmap (inj₁ tm) = fmap (raiseW Δ) tm
  str fmap (inj₂ k) = var (injectW _ k)

  open Subst str

  subst⊢ : ∀{Γ Δ T} → Γ ⊢ T → (∀{T} → T `∈ Γ → Δ ⊢ T) → Δ ⊢ T
  subst⊢ tm ρ = sub ρ tm

  -- ** Example

  tm1 : ε ⊢ (`■ `⇒ `■)
  tm1 = abs (var ze)

  tm2 : (ε ▹ (`■ `⇒ `■)) ⊢ (`■ `⇒ `■)
  tm2 = abs (app (var (suc ze)) (var ze))

  tm3 : ε ⊢ (`■ `⇒ `■)
  tm3 = subst⊢ tm2 ρ
    where ρ : ∀{T} → T `∈ (ε ▹ (`■ `⇒ `■)) → ε ⊢ T
          ρ ze = tm1
          ρ (suc ())

  test : tm3 ≡ abs (app (abs (var ze)) (var ze))
  test = refl