flambda-backend: Preallocation of blocks in Translmod for value let rec w/ flambda2 (#59)

Author: mshinwell

lambda/translmod.ml

@@ -212,6 +212,34 @@ let record_primitive = function
       primitive_declarations := p :: !primitive_declarations
   | _ -> ()
 
+(* Helper for compiling value "let rec".  This is only used for Flambda 2
+   at present, which uses the new [Dissect_letrec] module, planned to be
+   upstreamed.  At that point this helper can move into that module. *)
+
+let preallocate_letrec ~bindings ~body =
+  assert Config.flambda2;
+  let caml_update_dummy_prim =
+    Primitive.simple ~name:"caml_update_dummy" ~arity:2 ~alloc:true
+  in
+  let update_dummy var expr =
+    Lprim (Pccall caml_update_dummy_prim, [Lvar var; expr], Loc_unknown)
+  in
+  let bindings = List.rev bindings in
+  let body_with_initialization =
+    List.fold_left
+      (fun body (id, def, _size) -> Lsequence (update_dummy id def, body))
+      body bindings
+  in
+  List.fold_left
+    (fun body (id, _def, size) ->
+       let desc =
+         Primitive.simple ~name:"caml_alloc_dummy" ~arity:1 ~alloc:true
+       in
+       let size : lambda = Lconst (Const_base (Const_int size)) in
+       Llet (Strict, Pgenval, id,
+             Lprim (Pccall desc, [size], Loc_unknown), body))
+    body_with_initialization bindings
+
 (* Utilities for compiling "module rec" definitions *)
 
 let mod_prim = Lambda.transl_prim "CamlinternalMod"
@@ -419,12 +447,14 @@ let compile_recmodule ~scopes compile_rhs bindings cont =
 
 (* Code to translate class entries in a structure *)
 
+let class_block_size = 4
+
 let transl_class_bindings ~scopes cl_list =
   let ids = List.map (fun (ci, _) -> ci.ci_id_class) cl_list in
   (ids,
    List.map
      (fun ({ci_id_class=id; ci_expr=cl; ci_virt=vf}, meths) ->
-       (id, transl_class ~scopes ids id meths cl vf))
+       (id, transl_class ~scopes ids id meths cl vf, class_block_size))
      cl_list)
 
 (* Compile one or more functors, merging curried functors to produce
@@ -709,7 +739,13 @@ and transl_structure ~scopes loc fields cc rootpath final_env = function
             transl_structure ~scopes loc (List.rev_append ids fields)
               cc rootpath final_env rem
           in
-          Lletrec(class_bindings, body), size
+          if Config.flambda2 then
+            preallocate_letrec ~bindings:class_bindings ~body, size
+          else
+            let class_bindings =
+              List.map (fun (id, lam, _) -> id, lam) class_bindings
+            in
+            Lletrec(class_bindings, body), size
       | Tstr_include incl ->
           let ids = bound_value_identifiers incl.incl_type in
           let modl = incl.incl_mod in
@@ -1155,8 +1191,17 @@ let transl_store_structure ~scopes glob map prims aliases str =
                            (add_idents true ids subst) cont rem))
         | Tstr_class cl_list ->
             let (ids, class_bindings) = transl_class_bindings ~scopes cl_list in
+            let body = store_idents Loc_unknown ids in
             let lam =
-              Lletrec(class_bindings, store_idents Loc_unknown ids)
+              if Config.flambda2 then
+                preallocate_letrec
+                  ~bindings:class_bindings
+                  ~body
+              else
+                let class_bindings =
+                  List.map (fun (id, lam, _) -> id, lam) class_bindings
+                in
+                Lletrec(class_bindings, body)
             in
             Lsequence(Lambda.subst no_env_update subst lam,
                       transl_store ~scopes rootpath (add_idents false ids subst)
@@ -1524,7 +1569,14 @@ let transl_toplevel_item ~scopes item =
          be a value named identically *)
       let (ids, class_bindings) = transl_class_bindings ~scopes cl_list in
       List.iter set_toplevel_unique_name ids;
-      Lletrec(class_bindings, make_sequence toploop_setvalue_id ids)
+      let body = make_sequence toploop_setvalue_id ids in
+      if Config.flambda2 then
+        preallocate_letrec ~bindings:class_bindings ~body
+      else
+        let class_bindings =
+          List.map (fun (id, lam, _) -> id, lam) class_bindings
+        in
+        Lletrec(class_bindings, body)
   | Tstr_include incl ->
       let ids = bound_value_identifiers incl.incl_type in
       let modl = incl.incl_mod in