flambda-backend: Add Obj functions for accessing the scannable prefix len of a mixed block (#2636)

* Add scannable prefix externals

* Mark as noalloc

* bytecode/native changes

* Address review comment

Author: ncik-roberts

runtime/obj.c

@@ -330,3 +330,19 @@ struct queue_chunk {
   struct queue_chunk *next;
   value entries[ENTRIES_PER_QUEUE_CHUNK];
 };
+
+/* Return 0 for uniform blocks and 1+n for a mixed block with scannable prefix
+   len n.
+ */
+CAMLprim value caml_succ_scannable_prefix_len (value v) {
+#ifdef NATIVE_CODE
+  return Val_long(Reserved_val(v));
+#else
+  reserved_t reserved = Reserved_val(v);
+  if (reserved == Faux_mixed_block_sentinel) {
+    return Val_long(Scannable_wosize_val(v) + 1);
+  } else {
+    return Val_long(0);
+  }
+#endif /* NATIVE_CODE */
+}

runtime4/obj.c

@@ -365,3 +365,19 @@ struct queue_chunk {
   struct queue_chunk *next;
   value entries[ENTRIES_PER_QUEUE_CHUNK];
 };
+
+/* Return 0 for uniform blocks and 1+n for a mixed block with scannable prefix
+   len n.
+*/
+CAMLprim value caml_succ_scannable_prefix_len (value v) {
+#ifdef NATIVE_CODE
+  return Val_long(Reserved_val(v));
+#else
+  reserved_t reserved = Reserved_val(v);
+  if (reserved == Faux_mixed_block_sentinel) {
+    return Val_long(Scannable_wosize_val(v) + 1);
+  } else {
+    return Val_long(0);
+  }
+#endif /* NATIVE_CODE */
+}

stdlib/obj.ml

@@ -197,3 +197,31 @@ module Ephemeron = struct
   external blit_data : t -> t -> unit = "caml_ephe_blit_data"
 
 end
+
+module Uniform_or_mixed = struct
+  type obj_t = t
+
+  (* In native code, the raw reserved header bits, which is either 0 if the
+     block is uniform or n+1 if the block has a scannable prefix of length n.
+     In bytecode, this will be size+1 for "faux mixed blocks" representing
+     mixed records, and otherwise 0.
+   *)
+  type t = int
+
+  external of_block : obj_t -> t = "caml_succ_scannable_prefix_len" [@@noalloc]
+
+  type repr =
+    | Uniform
+    | Mixed of { scannable_prefix_len : int }
+
+  let repr = function
+    | 0 -> Uniform
+    | n -> Mixed { scannable_prefix_len = n - 1 }
+
+  let is_uniform t = t = 0
+  let is_mixed t = not (is_uniform t)
+  let mixed_scannable_prefix_len_exn t =
+    if is_uniform t
+    then invalid_arg "Uniform_or_mixed.mixed_scannable_prefix_len_exn";
+    t - 1
+end

stdlib/obj.mli

@@ -172,3 +172,39 @@ module Ephemeron: sig
   (** Maximum length of an ephemeron, ie the maximum number of keys an
       ephemeron could contain *)
 end
+
+module Uniform_or_mixed : sig
+  (** Blocks with a nominally scannable tag can still have a suffix of
+      unscanned objects; such a block is "mixed". This contrasts with
+      "uniform" blocks which are either all-scanned or all-unscanned.
+
+      Note that this module can return different results for the scannable
+      prefix len of a mixed block in native code vs. bytecode. That's
+      because more fields are scanned in bytecode.
+  *)
+
+  type obj_t := t
+
+  type t [@@immediate]
+
+  type repr =
+    | Uniform
+    (** The block is tagged as not scannable or the block is tagged as scannable
+        and all fields can be scanned. *)
+    | Mixed of { scannable_prefix_len : int }
+    (** The block is tagged as scannable but some fields can't be scanned. *)
+
+  val repr : t -> repr
+
+  external of_block : obj_t -> t = "caml_succ_scannable_prefix_len" [@@noalloc]
+
+  val is_uniform : t -> bool
+  (** Equivalent to [repr] returning [Uniform]. *)
+
+  val is_mixed : t -> bool
+  (** Equivalent to [repr] returning [Mixed _]. *)
+
+  val mixed_scannable_prefix_len_exn : t -> int
+  (** Returns the [scannable_prefix_len] without materializing the return
+      value of [repr]. Raises if [is_mixed] is [false]. *)
+end

testsuite/tests/lib-obj/uniform_or_mixed.bytecode.reference

@@ -0,0 +1,7 @@
+t_uniform1: uniform
+t_uniform2.0: uniform
+t_uniform2.1: uniform
+t_uniform3: uniform
+t_mixed0: mixed (scannable_prefix_len = 1)
+t_mixed1: mixed (scannable_prefix_len = 2)
+t_mixed2: mixed (scannable_prefix_len = 3)

testsuite/tests/lib-obj/uniform_or_mixed.ml

@@ -0,0 +1,53 @@
+(* TEST
+   flags = "-extension layouts_beta";
+   {
+     reference = "${test_source_directory}/uniform_or_mixed.native.reference";
+     compiler_reference2 = "${test_source_directory}/uniform_or_mixed.compiler.reference";
+     native;
+   }{
+     reference = "${test_source_directory}/uniform_or_mixed.bytecode.reference";
+     compiler_reference2 = "${test_source_directory}/uniform_or_mixed.compiler.reference";
+     bytecode;
+   }
+*)
+
+(* Bytecode and native code have slightly different outputs for the scannable prefix
+   of mixed records. The fields of mixed records must be scanned in bytecode.
+   (They are "faux mixed blocks".)
+*)
+
+type t_uniform1 = { x : int }
+type t_uniform2 = floatarray
+type t_uniform3 = int -> int
+
+type t_mixed0 = { x : int64# }
+type t_mixed1 = { x : string; y : int64# }
+type t_mixed2 = { x : string; y : string; z : int64# }
+
+let run (type a) test_name (obj : a) =
+  let obj = Obj.repr obj in
+  let uniform_or_mixed = Obj.Uniform_or_mixed.of_block obj in
+  match Obj.Uniform_or_mixed.repr uniform_or_mixed with
+  | Uniform ->
+      assert (Obj.Uniform_or_mixed.is_uniform uniform_or_mixed);
+      assert (not (Obj.Uniform_or_mixed.is_mixed uniform_or_mixed));
+      Printf.printf "%s: uniform\n" test_name
+  | Mixed { scannable_prefix_len } ->
+      assert (Obj.Uniform_or_mixed.is_mixed uniform_or_mixed);
+      assert (not (Obj.Uniform_or_mixed.is_uniform uniform_or_mixed));
+      assert
+        (Obj.Uniform_or_mixed.mixed_scannable_prefix_len_exn uniform_or_mixed =
+            scannable_prefix_len);
+      Printf.printf "%s: mixed (scannable_prefix_len = %d)\n"
+        test_name scannable_prefix_len;
+;;
+
+let () = run "t_uniform1" ({ x = 3 } : t_uniform1)
+let () = run "t_uniform2.0" (Float.Array.create 0 : t_uniform2)
+let () = run "t_uniform2.1" (Float.Array.create 1 : t_uniform2)
+let () = run "t_uniform3" ((fun x -> x) : t_uniform3)
+
+let () = run "t_mixed0" ({ x = #4L } : t_mixed0)
+let () = run "t_mixed1" ({ x = ""; y = #5L } : t_mixed1)
+let () = run "t_mixed2" ({ x = ""; y = ""; z = #5L }: t_mixed2)
+

testsuite/tests/lib-obj/uniform_or_mixed.native.reference

@@ -0,0 +1,7 @@
+t_uniform1: uniform
+t_uniform2.0: uniform
+t_uniform2.1: uniform
+t_uniform3: uniform
+t_mixed0: mixed (scannable_prefix_len = 0)
+t_mixed1: mixed (scannable_prefix_len = 1)
+t_mixed2: mixed (scannable_prefix_len = 2)