Hyperprover in 150 LOC

nbe.ml

type expr =
  | EVar of string
  | EApp of expr * expr
  | EPi of string * expr * expr
  | ESig of string * expr * expr
  | ELam of string * expr * expr
  | EPair of expr * expr
  | EFst of expr
  | ESnd of expr
  | ESet of int

let rec showExpr : expr -> string = function
  | EVar x -> x
  | EPi (x, a, b) -> Printf.sprintf "(Π (%s : %s), %s)" x (showExpr a) (showExpr b)
  | ESig (x, a, b) -> Printf.sprintf "(Σ (%s : %s), %s)" x (showExpr a) (showExpr b)
  | EApp (f, x) -> Printf.sprintf "(%s %s)" (showExpr f) (showExpr x)
  | ELam (x, t, e) -> Printf.sprintf "(λ (%s : %s), %s)" x (showExpr t) (showExpr e)
  | EPair (a, b) -> Printf.sprintf "(%s, %s)" (showExpr a) (showExpr b)
  | EFst x -> Printf.sprintf "%s.1" (showExpr x)
  | ESnd x -> Printf.sprintf "%s.2" (showExpr x)
  | ESet x -> Printf.sprintf "(U %d)" x

type value =
  | VVar of string
  | VApp of value * value
  | VPi of value * (value -> value)
  | VSig of value * (value -> value)
  | VLam of value * (value -> value)
  | VPair of value * value
  | VFst of value
  | VSnd of value
  | VSet of int

module Env = Map.Make(String)
type ctx = value Env.t

let rec synt (ctx : ctx) : expr -> value = function
  | EVar x -> begin
    match Env.find_opt x ctx with
    | Some v -> v
    | None -> failwith (Printf.sprintf "unbound variable “%s”" x)
  end
  | EPi (x, t, e) -> VPi (synt ctx t, fun y -> synt (Env.add x y ctx) e)
  | ESig (x, t, e) -> VSig (synt ctx t, fun y -> synt (Env.add x y ctx) e)
  | ELam (x, t, e) -> VLam (synt ctx t, fun y -> synt (Env.add x y ctx) e)
  | EApp (f, x) -> VApp (synt ctx f, synt ctx x)
  | EPair (a, b) -> VPair (synt ctx a, synt ctx b)
  | EFst x -> VFst (synt ctx x)
  | ESnd x -> VSnd (synt ctx x)
  | ESet x -> VSet x

let gen : int ref = ref 0
let fresh () = gen := !gen + 1; "#" ^ string_of_int !gen

let freshVar () = VVar (fresh ())

let rec norm : value -> value = function
  | VVar x -> VVar x
  | VPi (t, f) -> VPi (norm t, f)
  | VSig (t, f) -> VSig (norm t, f)
  | VLam (t, f) -> VLam (norm t, f)
  | VApp (f, x) -> app f x
  | VFst x -> vfst x
  | VSnd x -> vsnd x
  | VPair (a, b) -> VPair (norm a, norm b)
  | VSet x -> VSet x
and app f x =
  match norm f with
  | VLam (t, g) -> norm (g (norm x))
  | g           -> VApp (norm g, norm x)
and vfst x =
  match norm x with
  | VApp (a, _) -> a
  | y           -> VFst y
and vsnd x =
  match norm x with
  | VApp (_, b) -> b
  | y           -> VSnd y

let rec reify : value -> expr = function
  | VVar x -> EVar x
  | VPi (t, f) -> let x = fresh () in EPi (x, reify t, reify (norm (f (VVar x))))
  | VSig (t, f) -> let x = fresh () in ESig (x, reify t, reify (norm (f (VVar x))))
  | VLam (t, f) -> let x = fresh () in ELam (x, reify t, reify (norm (f (VVar x))))
  | VApp (f, x) -> EApp (reify f, reify x)
  | VPair (a, b) -> EPair (reify a, reify b)
  | VFst x -> EFst (reify x)
  | VSnd x -> ESnd (reify x)
  | VSet x -> ESet x

let eval ctx tau = norm (synt ctx tau)
let showValue v = showExpr (reify v)

let rec conv v1 v2 : bool = match v1, v2 with
  | VSet u, VSet v -> u = v
  | VVar a, VVar b -> a = b
  | VFst a, VFst b -> conv a b
  | VSnd a, VSnd b -> conv a b
  | VPair (a, b), VPair (c, d) -> conv a c && conv b d
  | VPair (a, b), v -> conv a (vfst v) && conv b (vsnd v)
  | v, VPair (a, b) -> conv (vfst v) a && conv (vsnd v) b
  | VPi (a, f), VPi (b, g)
  | VSig (a, f), VSig (b, g)
  | VLam (a, f), VLam (b, g) ->
    let x = freshVar () in conv a b && conv (f x) (g x)
  | VLam (a, f), b ->
    let x = freshVar () in conv (f x) (app b x)
  | a, VLam (b, f) ->
    let x = freshVar () in conv (app a x) (f x)
  | VApp (f, x), VApp (g, y) ->
    conv f g && conv x y
  | _, _ ->
    Printf.printf "FAIL: %s ≠ %s\n" (showValue v1) (showValue v2);
    false

let eqNf v1 v2 : unit =
  (*Printf.printf "EQNF: %s = %s\n" (showValue v1) (showValue v2);*)
  if conv v1 v2 then ()
  else failwith (Printf.sprintf "eqNf failed: %s ≠ %s" (showValue v1) (showValue v2))

let imax a b = match a, b with
  | VSet u, VSet v -> VSet (max u v)
  | _, _ -> failwith "expected type"

let extPiG : value -> value * (value -> value) = function
  | VPi (t, f) -> (t, f)
  | _          -> failwith "expected Π-type"

let extSigG : value -> value * (value -> value) = function
  | VSig (t, f) -> (t, f)
  | _           -> failwith "expected Σ-type"

let rec check gma rho (e : expr) (t : value) =
  (*Printf.printf "CHECK: %s : %s\n" (showExpr e) (showValue t);*)
  match e, t with
  | ELam (p, a, e), VPi (t, f) ->
    eqNf (eval rho a) t; let x = freshVar () in
    check (Env.add p t gma) (Env.add p x rho) e (f x)
  | EPair (a, b), VSig (t, f) ->
    check gma rho a t; check gma rho b (f (eval rho a))
  | e, t -> eqNf t (infer gma rho e)
and infer gma rho : expr -> value = function
  | EVar x -> begin
    match Env.find_opt x gma with
    | Some v -> v
    | None -> failwith (Printf.sprintf "unbound variable “%s”" x)
  end
  | ESet u -> VSet (u + 1)
  | EPi (p, a, e)
  | ESig (p, a, e) ->
    let x = freshVar () in
    let v = infer (Env.add p (eval rho a) gma) (Env.add p x rho) e in
    imax (infer gma rho a) v
  | EApp (f, x) ->
    let (t, g) = extPiG (infer gma rho f) in
    check gma rho x t; g (eval rho x)
  | EFst e -> fst (extSigG (infer gma rho e))
  | ESnd e -> let (_, f) = extSigG (infer gma rho e) in f (vfst (eval rho e))
  | x -> failwith (Printf.sprintf "cannot infer type of “%s”" (showExpr x))

let impl (a, b) = EPi ("_", a, b)
let ensure e t =
  Printf.printf "THEOREM: %s : %s\n" (showExpr e) (showExpr t);
  check Env.empty Env.empty e (eval Env.empty t)

(* Π (α β : U), α → β → α *)
let t = EPi ("α", ESet 0, EPi ("β", ESet 0, impl (EVar "α", impl (EVar "β", EVar "α"))))
let e = ELam ("A", ESet 0, ELam ("B", ESet 0,
  ELam ("x", EVar "A", ELam ("y", EVar "B", EVar "x"))))

let w1 = EVar "w"
let w2 = EPair (EFst (EVar "w"), ESnd (EVar "w"))

let sigma = ESig ("x", EVar "α", EApp (EVar "β", EVar "x"))

(* Π (α : U) (β : α → U) (π : Σ x, β x → U) (w : Σ x, β x), π w → π (w.1, w.2)*)
let extsig =
  EPi ("α", ESet 0, EPi ("β", impl (EVar "α", ESet 0),
    EPi ("π", impl (sigma, ESet 0),
      EPi ("w", sigma, impl (EApp (EVar "π", w1), EApp (EVar "π", w2))))))

let extsigproof =
  ELam ("α", ESet 0, ELam ("β", impl (EVar "α", ESet 0),
    ELam ("π", impl (sigma, ESet 0), ELam ("w", sigma,
      ELam ("p", EApp (EVar "π", w1), EVar "p")))))

let () = ensure e t; ensure extsigproof extsig