implement closure conversion on Ir.id's

bin/main.ml

@@ -5,7 +5,16 @@ let () =
   Logs.set_level (Some Logs.Debug);
 
   try
-    let ast = parse "val two = 2  fun twice(x) x*two  val r = twice(8)" in
+    let ast = parse "
+                val two = 2
+                val one = 1
+                fun twice(x) {
+                  fun f(x) (x - one) * two
+                  fun g() x + one
+                  f(g())
+                }
+                println(twice(4))
+ " in
     Logs.debug (fun m -> m "[AST] %a" Ast.pp_modl ast);
     let prog = compile ast in
     Logs.debug (fun m -> Code.dump (m "[BC] %s") prog.main);

lib/compile/bcc.ml

@@ -14,7 +14,7 @@ let undef_method =
   Value.Native_function
     (fun _ -> failwith "BUG: method undefined")
 
-let rec compile_lambda (lam : Ir.lambda) =
+let rec compile_lambda ?clos_map (lam : Ir.lambda) =
   let entrypoint = Code.make_block () in
   let currb = ref entrypoint in
   let emit i = Code.extend !currb i in
@@ -22,12 +22,13 @@ let rec compile_lambda (lam : Ir.lambda) =
 
   let reg_of_id = Hashtbl.create 128 in
   let set_reg id r =
-    if Hashtbl.mem reg_of_id id then failwith "BUG: id reassigned";
+    if Hashtbl.mem reg_of_id id then
+      Fmt.failwith "BUG: '%a' reassigned" Ir.pp_id id;
     Hashtbl.add reg_of_id id r
   in
   let get_reg id =
     try Hashtbl.find reg_of_id id with
-      Not_found -> failwith "BUG: id unassigned"
+      Not_found -> Fmt.failwith "BUG: '%a' unassigned" Ir.pp_id id
   in
 
   let rec emit_exp sp = function
@@ -103,23 +104,50 @@ let rec compile_lambda (lam : Ir.lambda) =
       in
       emit (CAL (sp, obj, mth, args))
 
-    | Ir.Obj { vals; funs } ->
-      let n_slots = List.length vals in
-      let elems = Hashtbl.create (List.length vals + List.length funs) in
-      let mthds = Array.make (List.length funs) undef_method in
+    | Ir.Obj { vals; funs; clos } ->
+      (* assign each captured id to a slot *)
+      let clos_map = Hashtbl.create 64 in
+      let n_slots =
+        List.fold_left
+          (fun n id ->
+             Hashtbl.add clos_map id n;
+             n + 1)
+          (List.length vals)
+          clos
+      in
 
+      (* assign each val to a slot *)
+      let elems = Hashtbl.create 64 in
       List.iteri
         (fun i name ->
            Hashtbl.add elems name (Value.Field i))
         vals;
 
+      (* compile methods and assign to an index *)
+      let mthds = Array.make (List.length funs) undef_method in
       List.iteri
         (fun i (name, lambda) ->
            Hashtbl.add elems name (Value.Method i);
-           mthds.(i) <- Code.Function (compile_lambda lambda))
+           mthds.(i) <- Code.Function (compile_lambda lambda ~clos_map))
         funs;
 
-      emit (CON (sp, { n_slots; elems; mthds }))
+      (* construct object and save captured id's *)
+      emit (CON (sp, { n_slots; elems; mthds }));
+      Hashtbl.iter
+        (fun id idx ->
+           let obj = sp in
+           let loc = suc sp in
+           emit (LDI (loc, Value.of_int idx));
+           emit (SET (get_reg id, obj, loc)))
+        clos_map
+
+    | Ir.Open id ->
+      let idx = try Hashtbl.find (Option.get clos_map) id
+        with Not_found -> failwith "BUG: %S not captured"
+           | Invalid_argument _ -> failwith "BUG: no captured variables"
+      in
+      emit (LDI (sp, Value.of_int idx));
+      emit (GET (sp, get_reg lam.self, sp))
 
     | ir ->
       let rv = emit_exp sp ir in
@@ -133,6 +161,8 @@ let rec compile_lambda (lam : Ir.lambda) =
 
   in
 
+  (* R0 = self *)
+  (* R(i+1) = args[i] *)
   set_reg lam.self (Code.R 0);
   let sp =
     List.fold_left

lib/compile/ir.ml

@@ -7,7 +7,17 @@ let compile_error f =
   Fmt.kstr (fun msg -> raise (Error msg)) f
 
 type imm = Value.t
-type id = Id of int [@@unboxed]
+
+type id = Id of int * string
+(* type id = Id of int [@@unboxed] *)
+
+let make_id_dispenser () =
+  let i = ref (-1) in
+  fun x -> (incr i; Id (!i, x))
+  (* fun _ -> (incr i; Id !i) *)
+
+let pp_id ppf (Id (n, x)) = Fmt.pf ppf "%s_%d" x n
+(* let pp_id ppf (Id n) = Fmt.pf ppf "_%d" n *)
 
 type uop =
   | Not
@@ -34,13 +44,14 @@ type exp =
   | Bop of bop * exp * exp
   | Call of path * exp list
   | Obj of obj
+  | Open of id
 
 and path = id * string
 
 and obj = {
   vals : string list;
   funs : (string * lambda) list;
-  (* clos : id list; *)
+  clos : id list;
 }
 
 and lambda = {
@@ -52,38 +63,47 @@ and lambda = {
 
 (* lower *)
 
-let make_id_dispenser () =
-  let i = ref (-1) in fun () -> (incr i; Id !i)
-
 module Env = struct
   type t =
+    | Empty
     | Obj of {
+        pred : t;
         self : id;
-        elems : string list;
+        elems : (string, string) Hashtbl.t;
       }
     | Fun of {
+        pred : t;
         args : (string * id) list;
-        (* clos : ??? *)
+        clos : (id, unit) Hashtbl.t;
       }
-    | Cons of t * t
 
   let rec find name = function
-    | Fun { args } ->
-      (* TODO: closure conversion *)
-      List.assoc name args, None
+    | Empty ->
+      raise Not_found
 
-    | Obj { self; elems } ->
-      if not (List.mem name elems) then
-        raise Not_found;
-      self, Some name
+    | Fun { pred; args; clos } ->
+      begin match List.assoc name args with
+        | id -> id, None
+        | exception Not_found ->
+          let id, fld = find name pred in
+          (* mark id's from pred env as needing capture *)
+          Hashtbl.replace clos id ();
+          id, fld
+      end
 
-    | Cons (e1, e2) ->
-      try find name e1 with
-        Not_found -> find name e2
+    | Obj { pred; self; elems } ->
+      begin match Hashtbl.find elems name with
+        | elem -> self, Some elem
+        | exception Not_found -> find name pred
+      end
 end
 
 let seq_r a b = Seq (b, a)
 
+let union xs ys =
+  List.sort_uniq compare
+    (List.rev_append ys xs)
+
 let lower ~lib (modl : Ast.modl) =
   let new_id = make_id_dispenser () in
 
@@ -142,7 +162,7 @@ let lower ~lib (modl : Ast.modl) =
     match path with
     | Ast.Ele (obj, fld) ->
       let rhs = lower_exp env obj in
-      let lhs = new_id () in
+      let lhs = new_id "get" in
       Let (lhs, rhs, k (`Get (lhs, fld)))
 
     | Ast.Var name ->
@@ -150,55 +170,58 @@ let lower ~lib (modl : Ast.modl) =
       | id, None -> k (`Var id)
       | obj, Some fld -> k (`Get (obj, fld))
       | exception Not_found ->
-        compile_error "unbound variable %S" name
+        compile_error "%S not in scope" name
 
   and lower_block ?(is_scope = false) env items =
     (* collect names of bindings to form the new environment; also check if a scope ends
        with an expression, if not then it is an error *)
-    let elems, ends_with_exp =
+    let elems = Hashtbl.create 32 in
+    let ends_with_exp =
       List.fold_left
-        (fun (elems, _) -> function
-           | Ast.Item_exp _ -> elems, true
+        (fun _ -> function
+           | Ast.Item_exp _ -> true
            | Ast.Item_val (name, _)
            | Ast.Item_obj (name, _)
-           | Ast.Item_fun (name, _, _) -> name :: elems, false)
-        ([], false)
+           | Ast.Item_fun (name, _, _) ->
+             if Hashtbl.mem elems name then
+               compile_error "multiple definitions of %S" name;
+             Hashtbl.add elems name name;
+             false)
+        false
         items
     in
     if is_scope && not ends_with_exp then
       compile_error "scope does not end in expression";
 
-    let self = new_id () in
-    let env = Env.Cons (Obj { self; elems }, env) in
+    let self = new_id "obj" in
+    let env = Env.Obj { self; elems; pred = env } in
 
-    let funs_r, vals_r, inits_r =
+    let funs_r, vals_r, inits_r, clos =
       List.fold_left
-        (fun (fns, vls, ins) -> function
+        (fun (fns, vls, ins, clos) -> function
            | Ast.Item_exp exp ->
              let init = lower_exp env exp in
-             fns, vls, init :: ins
+             fns, vls, init :: ins, clos
 
            | Ast.Item_val (name, exp) ->
              let init = Set ((self, name), lower_exp env exp) in
-             fns, name :: vls, init :: ins
+             fns, name :: vls, init :: ins, clos
 
            | Ast.Item_obj (name, items) ->
              (* TODO: it would be ideal if we could construct the empty versions of obj's
                 in a sort of "pre-init" phase, before assigning field values. but for now,
                 obj items are identical to val's where the rhs is an obj expression. *)
              let init = Set ((self, name), lower_block env items) in
-             fns, name :: vls, init :: ins
+             fns, name :: vls, init :: ins, clos
 
            | Ast.Item_fun (name, args, body) ->
-             let fn = name, compile_lambda self env args body in
-             fn :: fns, vls, ins)
-        ([], [], [])
+             let lam, clos' = lower_lambda self env args body in
+             (name, lam) :: fns, vls, ins, union clos clos')
+        ([], [], [], [])
         items
     in
 
-    (* TODO: closure conversion *)
-
-    (* if [is_scope], return the last expr, otherwise return the object itself *)
+    (* if [is_scope], return the last expr, otherwise return the object (self) *)
     let ret, inits_r = match is_scope, inits_r with
       | true, init :: inits -> init, inits
       | _,    inits         -> Var self, inits
@@ -210,6 +233,7 @@ let lower ~lib (modl : Ast.modl) =
       Obj {
         funs = List.rev funs_r;
         vals = List.rev vals_r;
+        clos;
       },
       List.fold_left
         (fun a b -> Seq (b, a))
@@ -217,16 +241,34 @@ let lower ~lib (modl : Ast.modl) =
         inits_r      
     )
 
-  and compile_lambda self env args body =
-    let args = List.map (fun a -> a, new_id ()) args in
-    let env = Env.Cons (Fun { args }, env) in
+  and lower_lambda self env args body =
+    let args = List.map (fun a -> a, new_id a) args in
+    let clos = Hashtbl.create 32 in
+    let env = Env.Fun { args; clos; pred = env } in
     let body = lower_exp env body in
+
+    (* wrap body in let bindings to read from the closure *)
+    let body, clos =
+      Hashtbl.fold
+        (fun id () (ir, clos) ->
+           if id = self then
+             (* [self] isn't "captured"; it IS the closure! *)
+             ir, clos
+           else
+             Let (id, Open id, ir), id :: clos)
+        clos
+        (body, [])
+    in
+
     let args = List.map snd args in
-    { self; args; body }
+    { self; args; body }, clos
 
   in
-  let self = new_id () in
-  let env = Env.Obj { self; elems = List.map fst lib } in
+
+  let self = new_id "lib" in
+  let elems = Hashtbl.create 128 in
+  List.iter (fun (name, _) -> Hashtbl.add elems name name) lib;
+  let env = Env.Obj { self; elems; pred = Empty } in
   let args = [] in
   let body = lower_block env modl.items in
   { self; args; body }