Terminating untyped NbE with a configurable limit

terminating-NbE.ml

module Make_term (S : sig
  type 'a t [@@deriving show]
end) =
struct
  type t = def S.t

  and def = Lam of t | Var of int | Appl of t * t
  [@@deriving show { with_path = false }]

  let _ = S.show
end

module T = Make_term (struct
  type 'a t = 'a [@@deriving show]
end)

(* A term annotated with a number indicating how many more times it can be
   passed to [eval] (defined below). *)
module Nbe_term = Make_term (struct
  type 'a t = 'a * int ref [@@deriving show]
end)

type value = VClosure of value list * Nbe_term.t | VNt of neutral
and neutral = NVar of int | NAppl of neutral * value

let vvar lvl = VNt (NVar lvl)
let vappl (m, n) = VNt (NAppl (m, n))

(* The maximum number of times [eval] is permitted to evaluate a single term
   (not accounting for different environments). *)
let limit = 1000

let tank : T.def -> Nbe_term.t =
  let open Nbe_term in
  let rec go term = (go_def term, ref limit)
  and go_def = function
    | T.Lam m -> Lam (go m)
    | T.Var idx -> Var idx
    | T.Appl (m, n) -> Appl (go m, go n)
  in
  go

exception Out_of_fuel of Nbe_term.def

let check_limit_exn (term, fuel) =
  if !fuel > 0 then decr fuel else raise (Out_of_fuel term)

(* Since [eval] can only be called with some term [M] at most [limit] times,
   and the set of terms is finite (they are all drawn from the input program),
   [eval] will always terminate. *)
let rec eval ~rho ((term : Nbe_term.def), fuel) =
  let open Nbe_term in
  check_limit_exn (term, fuel);
  match term with
  | Lam m -> VClosure (rho, m)
  | Var idx -> List.nth rho idx
  | Appl (m, n) -> (
      let m_val = eval ~rho m in
      let n_val = eval ~rho n in
      match m_val with
      | VClosure (rho, m) -> eval ~rho:(n_val :: rho) m
      | VNt neut -> vappl (neut, n_val))

let rec quote ~lvl : value -> T.def = function
  | VClosure (rho, m) ->
      let m_nf = normalize_at ~lvl ~rho m in
      T.Lam m_nf
  | VNt neut -> quote_neut ~lvl neut

and quote_neut ~lvl : neutral -> T.def = function
  | NVar var -> T.Var (lvl - var - 1)
  | NAppl (neut, n_val) ->
      let m_nf = quote_neut ~lvl neut in
      let n_nf = quote ~lvl n_val in
      T.Appl (m_nf, n_nf)

(* [normalize] will terminate too by the same reasoning. *)
and normalize ~lvl ~rho (term : Nbe_term.t) : T.def =
  quote ~lvl (eval ~rho term)

and normalize_at ~lvl ~rho (term : Nbe_term.t) : T.def =
  normalize ~lvl:(lvl + 1) ~rho:(vvar lvl :: rho) term

let test term =
  try
    (* Try to find the beta normal form of the input term... *)
    let _nf = normalize ~lvl:0 ~rho:[] (tank term) in
    print_endline "Ok."
  with Out_of_fuel term ->
    print_endline ("Out of fuel: " ^ Nbe_term.show_def term)

let zero = T.(Lam (Lam (Var 0)))
let succ = T.(Lam (Lam (Lam (Appl (Var 1, Appl (Appl (Var 2, Var 1), Var 0))))))

let mul =
  T.(Lam (Lam (Lam (Lam (Appl (Appl (Var 3, Appl (Var 2, Var 1)), Var 0))))))

let appl (f, list) = List.fold_left (fun m n -> T.Appl (m, n)) f list
let one = appl (succ, [ zero ])
let two = appl (succ, [ one ])
let three = appl (succ, [ two ])
let four = appl (succ, [ three ])
let five = appl (succ, [ four ])
let ten = appl (mul, [ five; two ])
let hundred = appl (mul, [ ten; ten ])
let thousand = appl (mul, [ hundred; ten ])

let () =
  let id_appl = T.(Appl (Lam (Var 0), Lam (Var 0))) in
  (* Ok. *)
  test id_appl;
  let n5k = appl (mul, [ thousand; five ]) in
  (* Ok. *)
  test n5k;
  let self_appl = T.(Lam (Appl (Var 0, Var 0))) in
  let omega = T.Appl (self_appl, self_appl) in
  (* Out of fuel: (Appl (((Var 0), ref (0)), ((Var 0), ref (0)))) *)
  test omega

Blog post: https://hirrolot.github.io/posts/fueled-evaluation.html.