A short and sweet implementation of NbE for Thorin.

Thorin.ml

(** {1 Syntax} *)

(** Expressions. *)
type expr =
  | Var of cont * int
  (** Local references.

      Thorin is a "blockless" IR, so we can reference arguments of
      other nodes by adding a data dependency edge, which implicitly
      nests the two nodes. *)

  | Cont of cont
  (** Abstraction over a continuation. *)

  | Ctor of int * expr list
  (** Datatype constructors. *)

  | Struct of expr list
  (** Structures.

      For the sake of simplicitly, records do not have named fields. *)

  | Proj of expr * int
(** Projections out of records. *)

(** Continuations are de Bruijn levels pointing to a corresponding {!type:compute}. *)
and cont = int

(** Computations. *)
and compute =
  | App of expr * expr list
  (** Application. *)

  | Case of expr * cont list
  (** A "simple" case statement that only discriminates on a single layer
      of constructors.

      Each case consists of a continuation that is parameterized
      by all pattern variables. *)

(** {1 Values} *)

(** Values. *)
and value =
  | Clo of (env * cont)
  (** Closures. *)

  | Ctor of int * value list
  (** Datatype constructors. *)

  | Struct of value list
  (** Structures. *)

  | Neu of neu
  (** Neutrals. *)

(** Neutrals. *)
and neu = {
    hd : hd;
    (** The "head" of a neutral is what is causing the comptuation to get stuck. *)

    spine : frm list
    (** The "spine" of a neutral consists of a stack of computation "frames" that are
        stuck on the head of the neutral. *)

  }

(** Neutral heads. *)
and hd =
  | Var of cont * int
  (** A computation can get stuck on an unassigned variable. *)

and frm =
  | App of value list
  (** Applications blocked on a neutral. *)

  | Case of cont list
  (** Case statements blocked on a neutral. *)

  | Proj of int
  (** Projections blocked on a neutral. *)

(** Environments ala persistent arrays. *)
and env =
  | Base of (value list) array
  (** An underlying mutable array. *)

  | Set of env * cont * value list
  (** Update an environment entry with a new list of arguments. *)


module Neu = struct
  (** Push a frame to a neutral. *)
  let push (neu : neu) (frm : frm) =
    Neu { neu with spine = frm :: neu.spine }

  (** Construct a {!const:Var!} neutral with an empty spine. *)
  let var (k : cont) (i : int) =
    Neu { hd = Var (k, i); spine = [] }
end

module Env = struct
  type t = env

  (** Get the ith argument for a continuation from an environment. *)
  let rec get (env : t) (k : cont) (i : int) : value =
    match env with
    | Set (env, k', vs) ->
       if k == k' then List.nth vs i else get env k i
    | Base env ->
       List.nth (Array.get env k) i

  (** Update a continuations arguments in an environment. *)
  let set (env : t) (k : cont) (vs : value list) : env =
    Set (env, k, vs)

  (** Compress all {!const:Set!} constructors away in an environment.
      This makes lookups more efficient, but removes any sharing. *)
  let compress (env : t) : env =
    let rec go (env : t) : (value list) array =
      match env with
      | Set (env, k, vs) ->
         let env = go env in
         let () = Array.set env k vs in
         env
      | Base env -> env
    in Base (go env)

  (** Make an environment consisting of only neutrals. *)
  let neus (arities : int list) : env =
    Base (Array.of_list (List.mapi (fun k arity -> List.init arity (Neu.var k)) arities))
end

(** {1 Evaluation} *)
let rec eval_expr (env : env) (e : expr) : value =
  match e with
  | Var (k, i) ->
     Env.get env k i
  | Cont k ->
     Clo (env, k)
  | Ctor (tag, es) ->
     Ctor (tag, eval_exprs env es)
  | Struct es ->
     Struct (eval_exprs env es)
  | Proj (e, i) ->
     do_proj (eval_expr env e) i

and eval_exprs (env : env) (es : expr list) : value list =
  List.map (eval_expr env) es

and eval_cont (cs : compute list) (env : env) (k : cont) : value =
  match List.nth cs k with
  | App (e, es) ->
     do_app cs (eval_expr env e) (eval_exprs env es)
  | Case (e, ks) ->
     do_case cs env (eval_expr env e) ks

and do_app (cs : compute list) (v : value) (vs : value list) =
  match v with
  | Clo (env, k) ->
     do_cont cs env k vs
  | Neu neu ->
    Neu.push neu (App vs)
  | _ ->
     failwith "bad do_app"

and do_case (cs : compute list) (env : env) (v : value) (ks : cont list) =
  match v with
  | Ctor (tag, vs) ->
     do_cont cs env (List.nth ks tag) vs
  | Neu neu ->
     Neu.push neu (Case ks)
  | _ ->
     failwith "bad do_case"

and do_proj (v : value) (i : int) =
  match v with
  | Struct vs ->
     List.nth vs i
  | Neu neu ->
     Neu.push neu (Proj i)
  | _ ->
     failwith "bad do_proj"

and do_cont (cs : compute list) (env : env) (k : cont) (vs : value list) =
  eval_cont cs (Env.set env k vs) k

(** {1 Test} *)

(*
  add(m, n, ret) =
    case m of
      Zero -> add_zero
      Succ -> add_succ
*)
let add =
  Case (Var (0, 0), [1; 2])

(*
  add_zero() =
    add.ret(0)
*)
let add_zero =
  App (Var (0, 2), [Var (0, 1)])

(*
  add_succ(m) =
    add(m, Succ(add.n), ret)
*)
let add_succ =
  App (Cont 0, [Var (2, 0); Ctor (1, [Var (0, 1)]); Var (0, 2)])

let prog () =
  do_cont
    [add; add_zero; add_succ]
    (Env.neus [3; 0; 1])
    0
    [Ctor (1, [Ctor (1, [Ctor (0, [])])]); Ctor (1, [Ctor (1, [Ctor (0, [])])]); Neu.var 0 2]

(*
- : value =
Neu
 {hd = Var (0, 2);
  spine =
   [App [Ctor (1, [Ctor (1, [Ctor (1, [Ctor (1, [Ctor (0, [])])])])])]]}
*)