-- Normal forms and hereditary substitutions
-- for simply-typed lambda calculus.
-- Andreas Abel, 2017-12-02

-- Simple types.

data Ty : Set where
  *   : Ty
  _⇒_ : (a b : Ty)  Ty

-- Typing contexts.

data Cxt : Set where
  ε   : Cxt
  _,_ : (Γ : Cxt) (a : Ty)  Cxt

infixr 4 _⇒_
infixl 2 _,_

-- Well-typed normal forms.

mutual

  -- Well-typed variables.

  data Var : (Γ : Cxt) (a : Ty)  Set where
    vz : ∀{Γ a}  Var (Γ , a) a
    vs : ∀{Γ a b} (x : Var Γ a)  Var (Γ , b) a

  -- Normal forms gather applications into a spine.

  data Nf (Γ : Cxt) : (a : Ty)  Set where
    var : ∀{a b} (x : Var Γ a) (S : Spine Γ a b)  Nf Γ b
    abs : ∀{a b} (t : Nf (Γ , a) b)  Nf Γ (a  b)

  -- Spines.

  data Spine (Γ : Cxt) : (a b : Ty)  Set where
    ε   : ∀{a}  Spine Γ a a
    _∙_ : ∀{a b c} (u : Nf Γ a) (S : Spine Γ b c)  Spine Γ (a  b) c

appendSpine : ∀{Γ a b c}  Spine Γ a b  Spine Γ b c  Spine Γ a c
appendSpine ε       T = T
appendSpine (u  S) T = u  appendSpine S T

-- Renamings.

Ren : (Δ Γ : Cxt)  Set
Ren Δ Γ = ∀{a}  Var Γ a  Var Δ a

-- Lifting a renaming under a binder.

liftR : ∀{Γ Δ b} (ρ : Ren Δ Γ)  Ren (Δ , b) (Γ , b)
liftR ρ vz     = vz
liftR ρ (vs x) = vs (ρ x)

-- Renaming in normal forms.

mutual
  ren : ∀{Γ Δ b} (ρ : Ren Δ Γ) (t : Nf Γ b)  Nf Δ b
  ren ρ (var x S) = var (ρ x) (renSpine ρ S)
  ren ρ (abs t)   = abs (ren (liftR ρ) t)

  renSpine : ∀{Γ Δ a b} (ρ : Ren Δ Γ) (S : Spine Γ a b)  Spine Δ a b
  renSpine ρ ε       = ε
  renSpine ρ (u  S) = ren ρ u  renSpine ρ S

-- Weakening.

wk : ∀{Γ a b}  Nf Γ b  Nf (Γ , a) b
wk = ren vs

-- Hereditary substitution and application.

-- Single substitution.
-- Indexed by the type of the substituted variable.

data Sub a : (Δ Γ : Cxt)  Set where
  sg   : ∀{Γ} (u : Nf Γ a)  Sub a Γ (Γ , a)
  lift : ∀{Δ Γ b} (σ : Sub a Δ Γ)  Sub a (Δ , b) (Γ , b)

-- The result of looking up a variable in a single substitution
-- indexed by type "a" is either
--
--   * a normal form of type "a", or
--   * a variable of some other type.

data Res a Δ : (b : Ty)  Set where
  var : ∀{b} (x : Var Δ b)  Res a Δ b
  tm  : (u : Nf Δ a)  Res a Δ a

-- Weakening results.

wkRes : ∀{Δ a b c}  Res a Δ b  Res a (Δ , c) b
wkRes (var x) = var (vs x)
wkRes (tm u)  = tm (wk u)

-- The lookup function.

lookup : ∀{Δ Γ a b} (σ : Sub a Δ Γ) (x : Var Γ b)  Res a Δ b
lookup (sg u)   vz     = tm u
lookup (sg u)   (vs x) = var x
lookup (lift σ) vz     = var vz
lookup (lift σ) (vs x) = wkRes (lookup σ x)

-- The hereditary substitution algorithm.

mutual

  -- Normalizing substitution of a normal form into a normal form.

  hsubst : ∀{Γ Δ a b} (σ : Sub a Δ Γ) (t : Nf Γ b)   Nf Δ b
  hsubst σ (var x S) = rapps (lookup σ x) (hsubsts σ S)
  hsubst σ (abs t)   = abs (hsubst (lift σ) t)

  -- Normalizing substitution of a normal form into a spine.
  -- (Just pointwise substitution.)

  hsubsts : ∀{Γ Δ a b c} (σ : Sub a Δ Γ) (S : Spine Γ b c)   Spine Δ b c
  hsubsts σ ε      = ε
  hsubsts σ (u  S) = hsubst σ u  hsubsts σ S

  -- Applying a lookup result to a spine of nfs.

  rapps :  {Δ a b c}  Res a Δ b  Spine Δ b c  Nf Δ c
  rapps (var x) S = var x S
  rapps (tm t) S = apps t S

  -- Applying a normal form to a spine of normal forms.

  apps : ∀{Γ a b} (t : Nf Γ a) (S : Spine Γ a b)  Nf Γ b
  apps t         ε       = t
  apps (var x S) T       = var x (appendSpine S T)
  apps (abs t)   (u  S) = apps (hsubst (sg u) t) S

-- Normalizing application function.

appNf : ∀{Γ a b}  (t : Nf Γ (a  b)) (u : Nf Γ a)  Nf Γ b
appNf t u = apps t (u  ε)

-- Well-typed terms.

data Tm (Γ : Cxt) : (a : Ty)  Set where
  var : ∀{a}   (x : Var Γ a)  Tm Γ a
  abs : ∀{a b} (t : Tm (Γ , a) b)  Tm Γ (a  b)
  app : ∀{a b} (t : Tm Γ (a  b)) (u : Tm Γ a)  Tm Γ b

-- Normalization function.

nf : ∀{Γ a}  Tm Γ a  Nf Γ a
nf (var x)   = var x ε
nf (abs t)   = abs (nf t)
nf (app t u) = appNf (nf t) (nf u)