A simple inlining pass in tagless-final style.

inlining-tagless-final.ml

(* Inlining for the λ-calculus, implemented in tagless final style. This seems
 * like it _must_ be related to Normalization By Evaluation (see eg [1,2]),
 * since they both amount to β-reduction (plus η-expansion in extensional NBE),
 * and the techniques have a similar "flavor", but I don't immediately see a
 * formal connection.
 *
 * [1] https://gist.github.com/rntz/f2f15f6cb4b8690c3599119225a57d33
 * [2] http://homepages.inf.ed.ac.uk/slindley/nbe/nbe-cambridge2016.pdf
 *)

(* Symbols: string names paired with unique ids. *)
type sym = {uid: int; name: string}
module Sym = struct
  type t = sym
  let compare x y = x.uid - y.uid
  let next_uid = ref 0
  let gen (name: string): t =
    let uid = !next_uid in next_uid := uid + 1; { uid; name }
end

(* Contexts: maps from symbols to things. *)
module Cx = Map.Make(Sym)
type 'a cx = 'a Cx.t

(* Signature for interpretations of the λ calculus. *)
module type LANG = sig
  type term
  val var: sym -> term
  val app: term -> term -> term
  val lam: sym -> term -> term
end

(* Precedence-aware printing. *)
module Pretty: LANG with type term = int * string = struct
  type term = int * string
  let paren (outer: int) (inner, text) =
    if inner <= outer then text else "(" ^ text ^ ")"
  let var x = 0, x.name
  let app m n = 1, paren 1 m ^ " " ^ paren 0 n
  let lam x body = 2, "\\" ^ x.name ^ ". " ^ paren 2 body
end

(* Inlining. Note that inlining might not terminate, eg. (λx.xx)(λx.xx).
 * But if you restrict to simply-typeable terms, it will. *)
module Inline(Next: LANG) = struct
  (* Thinking of inlining as a compiler pass, Next is the remaining backend or
   * the "next stage". `next` is the type of terms in the next stage. *)
  type next = Next.term

  (* A value takes a list of argument values and produces a next-stage term. *)
  type value = value' list -> next
  and value' = {apply: value}
  (* A term yields a value given a substitution mapping variables to values. *)
  type term = value cx -> value

  let appnext (f: next) (a: value'): next = Next.app f (a.apply [])
  (* `varval x` produces a value corresponding to the bare variable `x`. *)
  let varval (x: sym): value = List.fold_left appnext (Next.var x)

  let var x cx = try Cx.find x cx with Not_found -> varval x
  (* Shove `e` on the argument stack and evaluate `f`. *)
  let app f e cx args = f cx ({apply = e cx} :: args)
  let lam x body cx = function
    (* If we're given an argument, inline it and apply body to remaining
       arguments. *)
    | arg::args -> body (Cx.add x arg.apply cx) args
    (* Otherwise, produce a lambda. We freshen the parameter `x` to avoid
       capture, since we might be substituting terms into `body` which `x` as a
       free variable. *)
    | [] ->
       let x' = Sym.gen x.name in
       let cx' = Cx.add x (varval x') cx in
       Next.lam x' (body cx' [])
end


module Test = struct
  module I = Inline(Pretty)
  open I
  let x,y,z,foo,bar,baz =
    Sym.(gen "x", gen "y", gen "z", gen "foo", gen "bar", gen "baz")
  let vx,vy,vz,vfoo,vbar,vbaz = var x, var y, var z, var foo, var bar, var baz
  let (@) = app

  (* x y z   ==>   x y z *)
  let t1 = (vx @ vy) @ vz

  (* (\x.x) (y x)   ==>   y x *)
  let t2 = (lam x vx @ (vy @ vx))

  (* (\x y. y x) foo   ==>   \y. y foo *)
  let t3 = (lam x (lam y (vy @ vx))) @ vfoo

  (* (\x y. y x) foo bar   ==>   bar foo *)
  let t4 = ((lam x (lam y (vy @ vx))) @ vfoo) @ vbar

  (* (\x y. y x) foo (\x. x bar)   ==>   foo bar *)
  let t5 = ((lam x (lam y (vy @ vx))) @ vfoo) @ (lam x (vx @ vbar))

  (* TODO: moar tests. *)
  let run (x: I.term) = print_endline (snd (x Cx.empty []))
  let run_tests () = run t1; run t2; run t3; run t4; run t5
end