Untyped lambda calculus with HOAS and De Bruijn levels

untyped-hoas.ml

type term = Lam of (term -> term) | Appl of term * term | FreeVar of int

let unfurl lvl f = f (FreeVar lvl)
let unfurl2 lvl (f, g) = (unfurl lvl f, unfurl lvl g)

let rec print lvl =
  let plunge f = print (lvl + 1) (unfurl lvl f) in
  function
  | Lam f -> "(λ" ^ plunge f ^ ")"
  | Appl (m, n) -> "(" ^ print lvl m ^ " " ^ print lvl n ^ ")"
  | FreeVar x -> string_of_int x

let rec eval = function
  | Lam f -> Lam (fun n -> eval (f n))
  | Appl (m, n) -> (
      match (eval m, eval n) with Lam f, n -> f n | m, n -> Appl (m, n))
  | FreeVar x -> FreeVar x

let rec equate lvl =
  let plunge (f, g) = equate (lvl + 1) (unfurl2 lvl (f, g)) in
  function
  | Lam f, Lam g -> plunge (f, g)
  | Appl (m, n), Appl (m', n') -> equate lvl (m, m') && equate lvl (n, n')
  | FreeVar x, FreeVar y -> x = y
  | _, _ -> false

let () =
  let k = Lam (fun x -> Lam (fun _y -> x)) in
  let s =
    Lam (fun x -> Lam (fun y -> Lam (fun z -> Appl (Appl (x, z), Appl (y, z)))))
  in
  assert (print 0 k = "(λ(λ0))");
  assert (print 0 s = "(λ(λ(λ((0 2) (1 2)))))");
  assert (print 0 @@ eval (Appl (Lam (fun x -> x), FreeVar 123)) = "123");
  assert (print 0 @@ eval (Lam (fun x -> Appl (Lam (fun x -> x), x))) = "(λ0)");
  assert (equate 0 (k, k));
  assert (equate 0 (s, s));
  assert (not @@ equate 0 (k, s));
  assert (not @@ equate 0 (s, k));
  assert (not @@ equate 0 (FreeVar 123, Appl (Lam (fun x -> x), FreeVar 123)));
  assert (equate 0 (FreeVar 123, eval @@ Appl (Lam (fun x -> x), FreeVar 123)));
  let i = Appl (Appl (s, k), k) in
  assert (not @@ equate 0 (Lam (fun x -> x), i));
  assert (equate 0 (Lam (fun x -> x), eval i))