Normalization by Evaluation for the simply typed lambda calculus, in Agda

nbe.agda

-- Normalization by Evaluation for STLC, in Agda
--
-- Partially based on "NbE for CBPV and polarized lambda calculus" (Abel, 2019)

open import Function

infixr 5 _⇒_
infixl 4 _,_
infix 2 _⊃_

data Ty : Set where
  base : Ty
  _⇒_ : Ty → Ty → Ty

variable
  τ σ : Ty


-- Contexts form a category under extension
data Ctx : Set where
  ∙ : Ctx
  _,_ : Ctx → Ty → Ctx

variable
  Γ Δ Π : Ctx

data _⊃_ : Ctx → Ctx → Set where
  ⊃-base : ∙ ⊃ ∙
  ⊃-keep : Γ ⊃ Δ → Γ , τ ⊃ Δ , τ
  ⊃-extend : Γ ⊃ Δ → Γ , τ ⊃ Δ

⊃-refl : Γ ⊃ Γ
⊃-refl {∙} = ⊃-base
⊃-refl {Γ , x} = ⊃-keep ⊃-refl

⊃-trans : Δ ⊃ Π → Γ ⊃ Δ → Γ ⊃ Π
⊃-trans x ⊃-base = x
⊃-trans (⊃-keep x) (⊃-keep y) = ⊃-keep (⊃-trans x y)
⊃-trans (⊃-extend x) (⊃-keep y) = ⊃-extend (⊃-trans x y)
⊃-trans x (⊃-extend y) = ⊃-extend (⊃-trans x y)


-- Data types will be parametrized by their free variables, making them
-- presheaves Ctxᵒᵖ → Set
record PSh (A : Ctx → Set) : Set where
  field
    extend-ctx : Γ ⊃ Δ → A Δ → A Γ
    -- some functorality conditions...
    -- refl-does-nothing : ∀ x → x ≡ extend-ctx ⊃-refl x
    -- composes-ok : ∀ x p₁ p₂ →
    --   extend-ctx p₁ (extend-ctx p₂ x) ≡ extend-ctx (⊃-trans p₁ p₂) x
open PSh ⦃...⦄


-- Variables in the context!
data Var (τ : Ty) : Ctx → Set where
  z : Var τ (Γ , τ)
  s : Var τ Γ → Var τ (Γ , σ)

instance
  VarPSh : PSh (Var τ)
  VarPSh .extend-ctx (⊃-keep pf) z = z
  VarPSh .extend-ctx (⊃-keep pf) (s v) = s (extend-ctx pf v)
  VarPSh .extend-ctx (⊃-extend pf) v = s (extend-ctx pf v)


-- Intrinsically-typed raw syntax of the STLC
data Term : Ty → Ctx → Set where
  var : Var τ Γ → Term τ Γ
  app : Term (σ ⇒ τ) Γ → Term σ Γ → Term τ Γ
  abs : Term τ (Γ , σ) → Term (σ ⇒ τ) Γ


-- Normal forms!
data Nf : Ty → Ctx → Set
data Ne : Ty → Ctx → Set

data Nf where
  ne : Ne τ Γ → Nf τ Γ
  abs : Nf τ (Γ , σ) → Nf (σ ⇒ τ) Γ
data Ne where
  var : Var τ Γ → Ne τ Γ
  app : Ne (σ ⇒ τ) Γ → Nf σ Γ → Ne τ Γ

instance
  PShNf : PSh (Nf τ)
  PShNe : PSh (Ne τ)

  PShNf .extend-ctx pf (ne n) = ne (extend-ctx pf n)
  PShNf .extend-ctx pf (abs nf) = abs (extend-ctx (⊃-keep pf) nf)

  PShNe .extend-ctx pf (var x) = var (extend-ctx pf x)
  PShNe .extend-ctx pf (app n x) = app (extend-ctx pf n) (extend-ctx pf x)

-- The semantic domain!
Val : Ty → Ctx → Set
Val base Γ = Ne base Γ
Val (σ ⇒ τ) Γ = ∀ {Δ} → Δ ⊃ Γ → Val σ Δ → Val τ Δ

instance
  PShVal : PSh (Val τ)
  PShVal {base} .extend-ctx pf v = extend-ctx ⦃ PShNe ⦄ pf v
  PShVal {σ ⇒ τ} .extend-ctx pf v pf′ arg = v (⊃-trans pf pf′) arg

Env : Ctx → Ctx → Set
Env Γ Δ = ∀ {τ} → Var τ Γ → Val τ Δ

-- eval, aka unquote
eval : Env Γ Δ → Term τ Γ → Val τ Δ
eval ρ (var x) = ρ x
eval ρ (app e₁ e₂) = eval ρ e₁ ⊃-refl (eval ρ e₂)
eval ρ (abs e) {Π} pf arg = eval ρ′ e
  where ρ′ : Env _ _
        ρ′ {._} z = arg
        ρ′ (s x) = extend-ctx pf (ρ x)

-- reify, aka quote
reify : Val τ Γ → Nf τ Γ
reflect : Ne τ Γ → Val τ Γ

reify {base} x = ne x
reify {σ ⇒ τ} f = abs (reify (f (⊃-extend ⊃-refl) (reflect (var z))))

reflect {base} n = n
reflect {σ ⇒ τ} n = λ pf arg → reflect (app (extend-ctx pf n) (reify arg))

normalize : Term τ Γ → Nf τ Γ
normalize = reify ∘ eval ρ
  where ρ : Env Γ Γ
        ρ = reflect ∘ var


-- Some things to normalize!
nat : Ty
nat = (base ⇒ base) ⇒ (base ⇒ base)

mult : Term (nat ⇒ nat ⇒ nat) ∙
mult = abs (abs (abs (app (var (s (s z))) (app (var (s z)) (var z)))))

three four : Term nat ∙
three = abs (abs (f (f (f (var z)))))
  where f : _ → _
        f = app (var (s z))
four = abs (abs (f (f (f (f (var z))))))
  where f : _ → _
        f = app (var (s z))

-- Try: C-n twelve
twelve : Nf nat ∙
twelve = normalize (app (app mult three) four)