Bytecode: linear computation of closure environments

Author: lthls

Changes

@@ -163,6 +163,11 @@ Working version
   type to be ignored (see tests/typing-misc/constraints.ml).
   (Jacques Garrigue, report by Richard Eisenberg, review by Leo White)
 
+- #12207, #12222: Make closure computation linear in the number of recursive
+  functions instead of quadratic
+  (Vincent Laviron, report by François Pottier, review by Nathanaëlle Courant
+  and Gabriel Scherer)
+
 
 OCaml 5.1.0
 ---------------

bytecomp/bytegen.ml

@@ -34,20 +34,53 @@ let new_label () =
 (**** Operations on compilation environments. ****)
 
 let empty_env =
-  { ce_stack = Ident.empty; ce_heap = Ident.empty; ce_rec = Ident.empty }
+  { ce_stack = Ident.empty; ce_closure = Not_in_closure }
 
 (* Add a stack-allocated variable *)
 
 let add_var id pos env =
   { ce_stack = Ident.add id pos env.ce_stack;
-    ce_heap = env.ce_heap;
-    ce_rec = env.ce_rec }
+    ce_closure = env.ce_closure }
 
 let rec add_vars idlist pos env =
   match idlist with
     [] -> env
   | id :: rem -> add_vars rem (pos + 1) (add_var id pos env)
 
+(* Compute the closure environment *)
+
+let rec add_positions entries pos_to_entry ~pos ~delta = function
+  | [] -> entries, pos
+  | id :: rem ->
+    let entries =
+      Ident.add id (pos_to_entry pos) entries
+    in
+    add_positions entries pos_to_entry ~pos:(pos + delta) ~delta rem
+
+type function_definition =
+  | Single_non_recursive
+  | Multiple_recursive of Ident.t list
+
+let closure_entries fun_defs fvs =
+  let funct_entries, pos_end_functs =
+    match fun_defs with
+    | Single_non_recursive ->
+      (* No need to store the function in the environment, but we still need to
+         reserve a slot in the closure block *)
+      Ident.empty, 3
+    | Multiple_recursive functs ->
+      add_positions Ident.empty (fun pos -> Function pos) ~pos:0 ~delta:3 functs
+  in
+  (* Note: [pos_end_functs] is the position where we would store the next
+     function if there was one, and points after an eventual infix tag.
+     Since that was the last function, we don't need the last infix tag
+     and start storing free variables at [pos_end_functs - 1]. *)
+  let all_entries, _end_pos =
+    add_positions funct_entries (fun pos -> Free_variable pos)
+      ~pos:(pos_end_functs - 1) ~delta:1 fvs
+  in
+  all_entries
+
 (**** Examination of the continuation ****)
 
 (* Return a label to the beginning of the given continuation.
@@ -370,9 +403,8 @@ type function_to_compile =
   { params: Ident.t list;               (* function parameters *)
     body: lambda;                       (* the function body *)
     label: label;                       (* the label of the function entry *)
-    free_vars: Ident.t list;            (* free variables of the function *)
-    num_defs: int;            (* number of mutually recursive definitions *)
-    rec_vars: Ident.t list;             (* mutually recursive fn names *)
+    entries: closure_entry Ident.tbl;   (* the offsets for the free variables
+                                           and mutually recursive functions *)
     rec_pos: int }                      (* rank in recursive definition *)
 
 let functions_to_compile  = (Stack.create () : function_to_compile Stack.t)
@@ -570,15 +602,18 @@ let rec comp_expr stack_info env exp sz cont =
         let pos = Ident.find_same id env.ce_stack in
         Kacc(sz - pos) :: cont
       with Not_found ->
-      try
-        let pos = Ident.find_same id env.ce_heap in
-        Kenvacc(pos) :: cont
-      with Not_found ->
-      try
-        let ofs = Ident.find_same id env.ce_rec in
-        Koffsetclosure(ofs) :: cont
-      with Not_found ->
+      let not_found () =
         fatal_error ("Bytegen.comp_expr: var " ^ Ident.unique_name id)
+      in
+      match env.ce_closure with
+      | Not_in_closure -> not_found ()
+      | In_closure { entries; env_pos } ->
+        match Ident.find_same id entries with
+        | Free_variable pos ->
+          Kenvacc(pos - env_pos) :: cont
+        | Function pos ->
+          Koffsetclosure(pos - env_pos) :: cont
+        | exception Not_found -> not_found ()
       end
   | Lconst cst ->
       Kconst cst :: cont
@@ -627,9 +662,10 @@ let rec comp_expr stack_info env exp sz cont =
       let cont = add_pseudo_event loc !compunit_name cont in
       let lbl = new_label() in
       let fv = Ident.Set.elements(free_variables exp) in
+      let entries = closure_entries Single_non_recursive fv in
       let to_compile =
         { params = List.map fst params; body = body; label = lbl;
-          free_vars = fv; num_defs = 1; rec_vars = []; rec_pos = 0 } in
+          entries = entries; rec_pos = 0 } in
       Stack.push to_compile functions_to_compile;
       comp_args stack_info env (List.map (fun n -> Lvar n) fv) sz
         (Kclosure(lbl, List.length fv) :: cont)
@@ -646,14 +682,16 @@ let rec comp_expr stack_info env exp sz cont =
         let fv =
           Ident.Set.elements (free_variables (Lletrec(decl, lambda_unit))) in
         let rec_idents = List.map (fun (id, _lam) -> id) decl in
+        let entries =
+          closure_entries (Multiple_recursive rec_idents) fv
+        in
         let rec comp_fun pos = function
             [] -> []
           | (_id, Lfunction{params; body}) :: rem ->
               let lbl = new_label() in
               let to_compile =
                 { params = List.map fst params; body = body; label = lbl;
-                  free_vars = fv; num_defs = ndecl; rec_vars = rec_idents;
-                  rec_pos = pos} in
+                  entries = entries; rec_pos = pos} in
               Stack.push to_compile functions_to_compile;
               lbl :: comp_fun (pos + 1) rem
           | _ -> assert false in
@@ -1119,13 +1157,15 @@ let comp_block env exp sz cont =
 
 let comp_function tc cont =
   let arity = List.length tc.params in
-  let rec positions pos delta = function
-      [] -> Ident.empty
-    | id :: rem -> Ident.add id pos (positions (pos + delta) delta rem) in
+  let ce_stack, _last_pos =
+    add_positions Ident.empty Fun.id ~pos:arity ~delta:(-1) tc.params
+  in
   let env =
-    { ce_stack = positions arity (-1) tc.params;
-      ce_heap = positions (3 * (tc.num_defs - tc.rec_pos) - 1) 1 tc.free_vars;
-      ce_rec = positions (-3 * tc.rec_pos) 3 tc.rec_vars } in
+    { ce_stack;
+      ce_closure =
+        In_closure { entries = tc.entries; env_pos = 3 * tc.rec_pos }
+    }
+  in
   let cont =
     comp_block env tc.body arity (Kreturn arity :: cont) in
   if arity > 1 then

bytecomp/instruct.ml

@@ -15,10 +15,20 @@
 
 open Lambda
 
+type closure_entry =
+  | Free_variable of int
+  | Function of int
+
+type closure_env =
+  | Not_in_closure
+  | In_closure of {
+      entries: closure_entry Ident.tbl;
+      env_pos: int;
+    }
+
 type compilation_env =
   { ce_stack: int Ident.tbl;
-    ce_heap: int Ident.tbl;
-    ce_rec: int Ident.tbl }
+    ce_closure: closure_env }
 
 type debug_event =
   { mutable ev_pos: int;                (* Position in bytecode *)

bytecomp/instruct.mli

@@ -19,21 +19,35 @@ open Lambda
 
 (* Structure of compilation environments *)
 
+type closure_entry =
+  | Free_variable of int
+  | Function of int
+
+type closure_env =
+  | Not_in_closure
+  | In_closure of {
+      entries: closure_entry Ident.tbl; (* Offsets of the free variables and
+                                           recursive functions from the start of
+                                           the block *)
+      env_pos: int;                     (* Offset of the current function from
+                                           the start of the block *)
+    }
+
 type compilation_env =
-  { ce_stack: int Ident.tbl; (* Positions of variables in the stack *)
-    ce_heap: int Ident.tbl;  (* Structure of the heap-allocated env *)
-    ce_rec: int Ident.tbl }  (* Functions bound by the same let rec *)
+  { ce_stack: int Ident.tbl;  (* Positions of variables in the stack *)
+    ce_closure: closure_env } (* Structure of the heap-allocated env *)
 
 (* The ce_stack component gives locations of variables residing
    in the stack. The locations are offsets w.r.t. the origin of the
    stack frame.
-   The ce_heap component gives the positions of variables residing in the
-   heap-allocated environment.
-   The ce_rec component associates offsets to identifiers for functions
-   bound by the same let rec as the current function.  The offsets
-   are used by the OFFSETCLOSURE instruction to recover the closure
-   pointer of the desired function from the env register (which
-   points to the closure for the current function). *)
+   The ce_closure component gives the positions of variables residing in the
+   heap-allocated environment. The env_pos component gives the position of
+   the current function from the start of the closure block, and the entries
+   component gives the positions of free variables and functions bound by the
+   same let rec as the current function, from the start of the closure block.
+   These are used by the ENVACC and OFFSETCLOSURE instructions to recover the
+   relevant value from the env register (which points to the current function).
+*)
 
 (* Debugging events *)
 

debugger/eval.ml

@@ -51,20 +51,28 @@ let rec address path event = function
             with Symtable.Error _ -> raise(Error(Unbound_identifier id))
           end
       | None ->
+        let not_found () =
+          raise(Error(Unbound_identifier id))
+        in
         begin match event with
           Some {ev_ev = ev} ->
             begin try
               let pos = Ident.find_same id ev.ev_compenv.ce_stack in
               Debugcom.Remote_value.local (ev.ev_stacksize - pos)
             with Not_found ->
-            try
-              let pos = Ident.find_same id ev.ev_compenv.ce_heap in
-              Debugcom.Remote_value.from_environment pos
-            with Not_found ->
-              raise(Error(Unbound_identifier id))
+            match ev.ev_compenv.ce_closure with
+            | Not_in_closure -> not_found ()
+            | In_closure { entries; env_pos } ->
+              match Ident.find_same id entries with
+              | Free_variable pos ->
+                Debugcom.Remote_value.from_environment (pos - env_pos)
+              | Function _pos ->
+                (* Recursive functions seem to be unhandled *)
+                not_found ()
+              | exception Not_found -> not_found ()
             end
         | None ->
-            raise(Error(Unbound_identifier id))
+            not_found ()
         end
     end
   | Env.Adot(root, pos) ->