Bidirectional typechecking and normalisation by evaluation for a simply-typed lambda calculus, all in one

bidir-nbe.ml

(* Based on http://homepages.inf.ed.ac.uk/slindley/nbe/nbe-cambridge2016.pdf *)
exception TypeError

type tp = Base | Fn of tp * tp
let subtype (x: tp) (y: tp): bool = x = y

type sym = {name: string; id: int}

let next_id = ref 0
let nextId() = let x = !next_id in next_id := x + 1; x
let gensym name () = {name = name; id = nextId()}

type ('nf, 'ne) normF = [ `Lam of sym * 'nf ]
type ('nf, 'ne) neutF = [ `Var of sym | `App of 'ne * 'nf ]
type ('term, 'expr) exprF = [ `The of tp * 'term  | ('term,'expr) neutF ]

(* Normal and neutral terms. *)
type nf = [ (nf, ne) normF | (nf, ne) neutF ]
and  ne = (nf, ne) neutF

(* Checking terms & synthesizing expressions. *)
type term = [ (term, expr) normF  | (term, expr) exprF ]
and  expr = (term, expr) exprF

(* Typing contexts *)
type cx = (sym * tp) list
(* NBE "values", mixing neutral syntax & semantics. *)
type value = Neut of ne | Func of string * (value -> value)
type env = (sym * value) list
(* A term denotes a map from environments to values. *)
type den = env -> value

(* Maybe I should pass types to reflect, reify & app?
 * Then I could do extensional NBE! *)
let reflect (x: ne): value = Neut x

let rec reify (x: value): nf = match x with
  | Neut e -> (e :> nf)
  | Func (n,f) -> let x = gensym n () in
                  `Lam (x, reify (f (reflect (`Var x))))

let app (f: value) (a: value): value = match f with
  | Neut f -> reflect (`App (f, reify a))
  | Func (_,f) -> f a

let rec check (cx: cx) (expect: tp) (tm: term): den = match tm with
  | `Lam (x,body) ->
     (match expect with
      | Fn(a,b) -> let den = check ((x,a)::cx) b body in
                   fun env -> Func (x.name, fun v -> den ((x,v)::env))
      | _ -> raise TypeError)
  | #expr as e -> let (got,den) = infer cx e in
                  if subtype got expect then den
                  else raise TypeError

and infer (cx: cx) (tm: expr): tp * den = match tm with
  | `The (tp,term) -> (tp, check cx tp term)
  | `Var x -> (List.assoc x cx,
               fun env -> try List.assoc x env with Not_found -> Neut (`Var x))
  | `App (fncE, argT) ->
     (match infer cx fncE with
      | (Fn (a,b), fnc) ->
         let arg = check cx a argT in
         (b, fun env -> app (fnc env) (arg env))
      | _ -> raise TypeError)

let norm (d: den): nf = reify (d [])


(* Tests *)
let x = gensym "x" ();;
let y = gensym "y" ();;
let idfn = check [] (Fn (Base, Base)) (`Lam (x, `Var x));;
let (Base, boolx) = infer [x, Base] (`Var x);;
let (Fn (Base, Base), idfn2) = infer [] (`The (Fn (Base,Base), `Lam (x, `Var x)));;
let (Base, boolx2) = infer [x, Base] (`App (`The (Fn (Base,Base), `Lam (y, `Var y)), `Var x));;
(* Try: (norm idfn, norm boolx, norm idfn2, norm boolx2) *)