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