flambda-backend: Add is_last flag to closinfo words (ocaml-flambda#938)

Co-authored-by: Guillaume Bury <guillaume.bury@gmail.com>

asmcomp/cmm_helpers.ml

@@ -80,17 +80,20 @@ let caml_int64_ops = "caml_int64_ops"
 let pos_arity_in_closinfo = 8 * size_addr - 8
        (* arity = the top 8 bits of the closinfo word *)
 
-let closure_info ~arity ~startenv =
+let closure_info ~arity ~startenv ~is_last =
   let arity =
     match arity with
     | Lambda.Tupled, n -> -n
     | Lambda.Curried _, n -> n
   in
   assert (-128 <= arity && arity <= 127);
-  assert (0 <= startenv && startenv < 1 lsl (pos_arity_in_closinfo - 1));
+  assert (0 <= startenv && startenv < 1 lsl (pos_arity_in_closinfo - 2));
   Nativeint.(add (shift_left (of_int arity) pos_arity_in_closinfo)
-                 (add (shift_left (of_int startenv) 1)
-                      1n))
+    (add
+      (shift_left
+          (Bool.to_int is_last |> Nativeint.of_int)
+          (pos_arity_in_closinfo - 1))
+        (add (shift_left (of_int startenv) 1) 1n)))
 
 let alloc_float_header mode dbg =
   match mode with
@@ -102,8 +105,8 @@ let alloc_closure_header ~mode sz dbg =
   | Alloc_heap -> Cconst_natint (white_closure_header sz, dbg)
   | Alloc_local -> Cconst_natint (local_closure_header sz, dbg)
 let alloc_infix_header ofs dbg = Cconst_natint (infix_header ofs, dbg)
-let alloc_closure_info ~arity ~startenv dbg =
-  Cconst_natint (closure_info ~arity ~startenv, dbg)
+let alloc_closure_info ~arity ~startenv ~is_last dbg =
+  Cconst_natint (closure_info ~arity ~startenv ~is_last, dbg)
 let alloc_boxedint32_header mode dbg =
   match mode with
   | Lambda.Alloc_heap -> Cconst_natint (boxedint32_header, dbg)
@@ -2161,7 +2164,7 @@ let rec intermediate_curry_functions ~nlocal ~arity num =
                [alloc_closure_header ~mode 5 (dbg ());
                 Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1), dbg ());
                 alloc_closure_info ~arity:(curried (arity - num - 1))
-                                   ~startenv:3 (dbg ());
+                                   ~startenv:3 ~is_last:true (dbg ());
                 Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1) ^ "_app",
                   dbg ());
                 Cvar arg; Cvar clos],
@@ -2170,7 +2173,8 @@ let rec intermediate_curry_functions ~nlocal ~arity num =
            Cop(Calloc mode,
                 [alloc_closure_header ~mode 4 (dbg ());
                  Cconst_symbol(name1 ^ "_" ^ Int.to_string (num+1), dbg ());
-                 alloc_closure_info ~arity:(curried 1) ~startenv:2 (dbg ());
+                 alloc_closure_info ~arity:(curried 1) ~startenv:2
+                   ~is_last:true (dbg ());
                  Cvar arg; Cvar clos],
                 dbg ());
       fun_codegen_options = [];
@@ -2915,32 +2919,34 @@ let emit_constant_closure ((_, global_symb) as symb) fundecls clos_vars cont =
       let rec emit_others pos = function
           [] -> clos_vars @ cont
       | (f2 : Clambda.ufunction) :: rem ->
+          let is_last = match rem with [] -> true | _ :: _ -> false in
           match f2.arity with
           | Curried _, (0|1) as arity ->
             Cint(infix_header pos) ::
             (closure_symbol f2) @
             Csymbol_address f2.label ::
-            Cint(closure_info ~arity ~startenv:(startenv - pos)) ::
+            Cint(closure_info ~arity ~startenv:(startenv - pos) ~is_last) ::
             emit_others (pos + 3) rem
           | arity ->
             Cint(infix_header pos) ::
             (closure_symbol f2) @
             Csymbol_address(curry_function_sym f2.arity) ::
-            Cint(closure_info ~arity ~startenv:(startenv - pos)) ::
+            Cint(closure_info ~arity ~startenv:(startenv - pos) ~is_last) ::
             Csymbol_address f2.label ::
             emit_others (pos + 4) rem in
+      let is_last = match remainder with [] -> true | _ :: _ -> false in
       Cint(black_closure_header (fundecls_size fundecls
                                  + List.length clos_vars)) ::
       cdefine_symbol symb @
       (closure_symbol f1) @
       match f1.arity with
       | Curried _, (0|1) as arity ->
         Csymbol_address f1.label ::
-        Cint(closure_info ~arity ~startenv) ::
+        Cint(closure_info ~arity ~startenv ~is_last) ::
         emit_others 3 remainder
       | arity ->
         Csymbol_address(curry_function_sym f1.arity) ::
-        Cint(closure_info ~arity ~startenv) ::
+        Cint(closure_info ~arity ~startenv ~is_last) ::
         Csymbol_address f1.label ::
         emit_others 4 remainder
 

asmcomp/cmm_helpers.mli

@@ -59,12 +59,13 @@ val boxedint64_header : nativeint
 val boxedintnat_header : nativeint
 
 (** Closure info for a closure of given arity and distance to environment *)
-val closure_info : arity:Clambda.arity -> startenv:int -> nativeint
+val closure_info : arity:Clambda.arity -> startenv:int -> is_last:bool
+  -> nativeint
 
 (** Wrappers *)
 val alloc_infix_header : int -> Debuginfo.t -> expression
 val alloc_closure_info :
-      arity:(Lambda.function_kind * int) -> startenv:int ->
+      arity:(Lambda.function_kind * int) -> startenv:int -> is_last:bool ->
       Debuginfo.t -> expression
 
 (** Integers *)

asmcomp/cmmgen.ml

@@ -395,19 +395,20 @@ let rec transl env e =
           [] ->
             List.map (transl env) clos_vars
         | f :: rem ->
+            let is_last = match rem with [] -> true | _::_ -> false in
             Cmmgen_state.add_function f;
             let dbg = f.dbg in
             let without_header =
               match f.arity with
               | Curried _, (1|0) as arity ->
                 Cconst_symbol (f.label, dbg) ::
                 alloc_closure_info ~arity
-                                   ~startenv:(startenv - pos) dbg ::
+                                   ~startenv:(startenv - pos) ~is_last dbg ::
                 transl_fundecls (pos + 3) rem
               | arity ->
                 Cconst_symbol (curry_function_sym f.arity, dbg) ::
                 alloc_closure_info ~arity
-                                   ~startenv:(startenv - pos) dbg ::
+                                   ~startenv:(startenv - pos) ~is_last dbg ::
                 Cconst_symbol (f.label, dbg) ::
                 transl_fundecls (pos + 4) rem
             in

runtime/alloc.c

@@ -229,7 +229,7 @@ CAMLprim value caml_alloc_dummy_infix(value vsize, value voffset)
      block contains no pointers into the heap.  However, the block
      cannot be marshaled or hashed, because not all closinfo fields
      and infix header fields are correctly initialized. */
-  Closinfo_val(v) = Make_closinfo(0, wosize);
+  Closinfo_val(v) = Make_closinfo(0, wosize, 1);
   if (offset > 0) {
     v += Bsize_wsize(offset);
     Hd_val(v) = Make_header(offset, Infix_tag, Caml_white);

runtime/caml/mlvalues.h

@@ -245,20 +245,28 @@ CAMLextern value caml_get_public_method (value obj, value tag);
 #define Code_val(val) (((code_t *) (val)) [0])     /* Also an l-value. */
 #define Closinfo_val(val) Field((val), 1)          /* Arity and start env */
 /* In the closure info field, the top 8 bits are the arity (signed).
+   The next least significant bit is set iff the current closure is the
+   last one to occur in the block.  (This is used in the compactor.)
    The low bit is set to one, to look like an integer.
-   The remaining bits are the field number for the first word of the
-   environment, or, in other words, the offset (in words) from the closure
-   to the environment part. */
+   The remaining bits are the field number for the first word of the scannable
+   part of the environment, or, in other words, the offset (in words) from the
+   closure to the scannable part of the environment.
+   The non-scannable part of the environment lives between the end of the
+   last closure and the start of the scannable environment within the block. */
 #ifdef ARCH_SIXTYFOUR
 #define Arity_closinfo(info) ((intnat)(info) >> 56)
-#define Start_env_closinfo(info) (((uintnat)(info) << 8) >> 9)
-#define Make_closinfo(arity,delta) \
-  (((uintnat)(arity) << 56) + ((uintnat)(delta) << 1) + 1)
+#define Start_env_closinfo(info) (((uintnat)(info) << 9) >> 10)
+#define Is_last_closinfo(info) (((uintnat)(info) << 8) >> 63)
+#define Make_closinfo(arity,delta,is_last) \
+  (((uintnat)(arity) << 56) + ((uintnat)(is_last) << 55) \
+    + ((uintnat)(delta) << 1) + 1)
 #else
 #define Arity_closinfo(info) ((intnat)(info) >> 24)
-#define Start_env_closinfo(info) (((uintnat)(info) << 8) >> 9)
-#define Make_closinfo(arity,delta) \
-  (((uintnat)(arity) << 24) + ((uintnat)(delta) << 1) + 1)
+#define Start_env_closinfo(info) (((uintnat)(info) << 9) >> 10)
+#define Is_last_closinfo(info) (((uintnat)(info) << 8) >> 31)
+#define Make_closinfo(arity,delta,is_last) \
+  (((uintnat)(arity) << 24) + ((uintnat)(is_last) << 23) \
+    + ((uintnat)(delta) << 1) + 1)
 #endif
 
 /* This tag is used (with Forward_tag) to implement lazy values.

runtime/compact.c

@@ -261,16 +261,18 @@ static void do_compaction (intnat new_allocation_policy)
 
           if (t == Closure_tag){
             /* Revert the infix pointers to this block. */
-            mlsize_t i, startenv;
+            mlsize_t i;
             value v;
 
             v = Val_hp (p);
-            startenv = Start_env_closinfo (Closinfo_val (v));
             i = 0;
             while (1){
-              int arity = Arity_closinfo (Field (v, i+1));
+              int arity;
+              uintnat closinfo = Field (v, i+1);
+              if (Is_last_closinfo (closinfo)) break;
+              arity = Arity_closinfo (closinfo);
               i += 2 + (arity != 0 && arity != 1);
-              if (i >= startenv) break;
+              CAMLassert (i < Start_env_closinfo (Closinfo_val (v)));
 
               /* Revert the inverted list for infix header at offset [i]. */
               q = Field (v, i);

runtime/interp.c

@@ -568,7 +568,7 @@ value caml_interprete(code_t prog, asize_t prog_size)
         Field(accu, 2) = env;
         for (i = 0; i < num_args; i++) Field(accu, i + 3) = sp[i];
         Code_val(accu) = pc - 3; /* Point to the preceding RESTART instr. */
-        Closinfo_val(accu) = Make_closinfo(0, 2);
+        Closinfo_val(accu) = Make_closinfo(0, 2, 1);
         sp += num_args;
         pc = (code_t)(sp[0]);
         env = sp[1];
@@ -596,7 +596,7 @@ value caml_interprete(code_t prog, asize_t prog_size)
       /* The code pointer is not in the heap, so no need to go through
          caml_initialize. */
       Code_val(accu) = pc + *pc;
-      Closinfo_val(accu) = Make_closinfo(0, 2);
+      Closinfo_val(accu) = Make_closinfo(0, 2, 1);
       pc++;
       sp += nvars;
       Next;
@@ -628,13 +628,14 @@ value caml_interprete(code_t prog, asize_t prog_size)
       *--sp = accu;
       p = &Field(accu, 0);
       *p++ = (value) (pc + pc[0]);
-      *p++ = Make_closinfo(0, envofs);
+      *p++ = Make_closinfo(0, envofs, nfuncs < 2);
       for (i = 1; i < nfuncs; i++) {
         *p++ = Make_header(i * 3, Infix_tag, Caml_white); /* color irrelevant */
         *--sp = (value) p;
         *p++ = (value) (pc + pc[i]);
         envofs -= 3;
-        *p++ = Make_closinfo(0, envofs);
+        CAMLassert(i <= nfuncs - 1);
+        *p++ = Make_closinfo(0, envofs, i == nfuncs - 1);
       }
       pc += nfuncs;
       Next;

runtime/meta.c

@@ -126,7 +126,7 @@ CAMLprim value caml_reify_bytecode(value ls_prog,
 
   clos = caml_alloc_small (2, Closure_tag);
   Code_val(clos) = (code_t) prog;
-  Closinfo_val(clos) = Make_closinfo(0, 2);
+  Closinfo_val(clos) = Make_closinfo(0, 2, 1);
   bytecode = caml_alloc_small (2, Abstract_tag);
   Bytecode_val(bytecode)->prog = prog;
   Bytecode_val(bytecode)->len = len;

runtime/obj.c

@@ -108,7 +108,7 @@ CAMLprim value caml_obj_block(value tag, value size)
     /* Closinfo_val is the second field, so we need size at least 2 */
     if (sz < 2) caml_invalid_argument ("Obj.new_block");
     res = caml_alloc(sz, tg);
-    Closinfo_val(res) = Make_closinfo(0, 2); /* does not allocate */
+    Closinfo_val(res) = Make_closinfo(0, 2, 1); /* does not allocate */
     break;
   }
   case String_tag: {

utils/config.mlp

@@ -231,6 +231,7 @@ let configuration_variables =
 
   p_bool "flambda_backend" flambda_backend;
   p_bool "probes" probes;
+  p_bool "stack_allocation" stack_allocation;
 ]
 
 let print_config_value oc = function