A simple CPS conversion as in "Compiling with Continuations", Andrew W. Appel

a-main.ml

(* A variable identifier of the lambda language [term]. *)
type var = string [@@deriving eq]

(* The lambda language; direct style. *)
type term =
  | Var of var
  | Fix of (var * var list * term) list * term
  | Appl of term * term list
  | Record of term list
  | Select of term * int

type cps_var =
  (* Taken from the lambda term during CPS conversion. *)
  | CLamVar of var
  (* Generated uniquely during CPS conversion. *)
  | CGenVar of int
[@@deriving eq]

(* The resulting CPS term. *)
type cps_term =
  | CFix of (cps_var * cps_var list * cps_term) list * cps_term
  | CAppl of cps_var * cps_var list
  | CRecord of cps_var list * binder
  | CSelect of cps_var * int * binder
  | CHalt of cps_var
[@@deriving eq]

(* Binds a unique [cps_var] within [cps_term]. *)
and binder = cps_var * cps_term

(* Generates a unique CPS variable given the current [i]. *)
let gensym i =
  let x = CGenVar !i in
  i := !i + 1;
  x

(* Converts [term] to [cps_term], applying [finish] to the resulting
   CPS variable. *)
let rec convert gen finish = function
  | Var x -> finish (CLamVar x)
  | Fix (defs, m) -> CFix (List.map (convert_def gen) defs, convert gen finish m)
  | Appl (f, args) ->
      let ret_k = gensym gen in
      let ret_k_x = gensym gen in
      CFix
        ( [ (ret_k, [ ret_k_x ], finish ret_k_x) ],
          f
          |> convert gen (fun f_cps ->
                 args
                 |> convert_list gen (fun args_cps ->
                        CAppl (f_cps, args_cps @ [ ret_k ]))) )
  | Record fields ->
      fields
      |> convert_list gen (fun fields_cps ->
             let x = gensym gen in
             CRecord (fields_cps, (x, finish x)))
  | Select (m, i) ->
      m
      |> convert gen (fun m_cps ->
             let x = gensym gen in
             CSelect (m_cps, i, (x, finish x)))

(* Converts [term list] to [cps_var list] and applies [finish] to it. *)
and convert_list gen finish =
  let rec go acc = function
    | [] -> finish (List.rev acc)
    | x :: xs -> x |> convert gen (fun x_cps -> go (x_cps :: acc) xs)
  in
  go []

(* Converts a single function definition to its CPS form. *)
and convert_def gen (f, params, m) =
  let k = gensym gen in
  ( CLamVar f,
    List.map (fun x -> CLamVar x) params @ [ k ],
    m |> convert gen (fun m_cps -> CAppl (k, [ m_cps ])) )

(* Test CPS conversion. *)
let () =
  let assert_convert t expected =
    let convert = convert (ref 0) (fun x -> CHalt x) in
    assert (equal_cps_term (convert t) expected)
  in

  (* [Var] *)
  assert_convert (Var "a") (CHalt (CLamVar "a"));

  (* [Fix] *)
  assert_convert
    (Fix ([ ("f", [ "x" ], Var "x"); ("g", [ "y" ], Var "y") ], Var "a"))
    (CFix
       ( [
           ( CLamVar "f",
             [ CLamVar "x"; CGenVar 0 ],
             CAppl (CGenVar 0, [ CLamVar "x" ]) );
           ( CLamVar "g",
             [ CLamVar "y"; CGenVar 1 ],
             CAppl (CGenVar 1, [ CLamVar "y" ]) );
         ],
         CHalt (CLamVar "a") ));

  (* [Appl] *)
  assert_convert
    (Appl (Var "a", [ Var "b"; Var "c" ]))
    (CFix
       ( [ (CGenVar 0, [ CGenVar 1 ], CHalt (CGenVar 1)) ],
         CAppl (CLamVar "a", [ CLamVar "b"; CLamVar "c"; CGenVar 0 ]) ));

  (* [Appl] list conversion. *)
  assert_convert
    (Appl (Var "a", [ Select (Var "b", 2); Select (Var "c", 3) ]))
    (CFix
       ( [ (CGenVar 0, [ CGenVar 1 ], CHalt (CGenVar 1)) ],
         CSelect
           ( CLamVar "b",
             2,
             ( CGenVar 2,
               CSelect
                 ( CLamVar "c",
                   3,
                   ( CGenVar 3,
                     CAppl (CLamVar "a", [ CGenVar 2; CGenVar 3; CGenVar 0 ]) )
                 ) ) ) ));

  (* [Record] *)
  assert_convert
    (Record [ Var "a"; Var "b"; Var "c" ])
    (CRecord
       ( [ CLamVar "a"; CLamVar "b"; CLamVar "c" ],
         (CGenVar 0, CHalt (CGenVar 0)) ));

  (* [Record] list conversion. *)
  assert_convert
    (Record [ Select (Var "a", 2); Select (Var "b", 3); Select (Var "c", 4) ])
    (CSelect
       ( CLamVar "a",
         2,
         ( CGenVar 0,
           CSelect
             ( CLamVar "b",
               3,
               ( CGenVar 1,
                 CSelect
                   ( CLamVar "c",
                     4,
                     ( CGenVar 2,
                       CRecord
                         ( [ CGenVar 0; CGenVar 1; CGenVar 2 ],
                           (CGenVar 3, CHalt (CGenVar 3)) ) ) ) ) ) ) ));

  (* [Select] *)
  assert_convert
    (Select (Var "a", 2))
    (CSelect (CLamVar "a", 2, (CGenVar 0, CHalt (CGenVar 0))));

  (* A more complicated example. *)
  assert_convert
    (Fix
       ( [
           ("f", [ "a"; "b" ], Select (Var "a", 2));
           ("g", [ "a"; "b" ], Select (Var "b", 3));
         ],
         Record [ Var "x"; Appl (Var "f", [ Var "y"; Select (Var "z", 4) ]) ] ))
    (CFix
       ( [
           ( CLamVar "f",
             [ CLamVar "a"; CLamVar "b"; CGenVar 4 ],
             CSelect
               (CLamVar "a", 2, (CGenVar 5, CAppl (CGenVar 4, [ CGenVar 5 ])))
           );
           ( CLamVar "g",
             [ CLamVar "a"; CLamVar "b"; CGenVar 6 ],
             CSelect
               (CLamVar "b", 3, (CGenVar 7, CAppl (CGenVar 6, [ CGenVar 7 ])))
           );
         ],
         CFix
           ( [
               ( CGenVar 0,
                 [ CGenVar 1 ],
                 CRecord
                   ([ CLamVar "x"; CGenVar 1 ], (CGenVar 3, CHalt (CGenVar 3)))
               );
             ],
             CSelect
               ( CLamVar "z",
                 4,
                 ( CGenVar 2,
                   CAppl (CLamVar "f", [ CLamVar "y"; CGenVar 2; CGenVar 0 ]) )
               ) ) ))

dune

(executable
 (public_name cps_conv)
 (name main)
 (libraries cps_conv)
 (preprocess
  (pps ppx_deriving.eq)))