flambda-backend: add the %get_header primitive (#1539)

Author: riaqn

asmcomp/cmmgen.ml

@@ -642,7 +642,7 @@ let rec transl env e =
          | Pasrbint _ | Pbintcomp (_, _) | Pstring_load _ | Pbytes_load _
          | Pbytes_set _ | Pbigstring_load _ | Pbigstring_set _
          | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
-         | Pbbswap _), _)
+         | Pbbswap _ | Pget_header _), _)
         ->
           fatal_error "Cmmgen.transl:prim"
       end
@@ -960,6 +960,8 @@ and transl_prim_1 env p arg dbg =
   | Pbswap16 ->
       tag_int (bswap16 (ignore_high_bit_int (untag_int
         (transl env arg) dbg)) dbg) dbg
+  | Pget_header m ->
+      box_int dbg Pnativeint m (get_header (transl env arg) dbg)
   | (Pfield_computed | Psequand | Psequor
     | Paddint | Psubint | Pmulint | Pandint
     | Porint | Pxorint | Plslint | Plsrint | Pasrint
@@ -1155,7 +1157,7 @@ and transl_prim_2 env p arg1 arg2 dbg =
   | Pnegbint _ | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _)
   | Pbigarraydim _ | Pbytes_set _ | Pbigstring_set _ | Pbbswap _
   | Pprobe_is_enabled _
-  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
+  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _ | Pget_header _
     ->
       fatal_errorf "Cmmgen.transl_prim_2: %a"
         Printclambda_primitives.primitive p
@@ -1216,7 +1218,7 @@ and transl_prim_3 env p arg1 arg2 arg3 dbg =
   | Pbigarrayref (_, _, _, _) | Pbigarrayset (_, _, _, _) | Pbigarraydim _
   | Pstring_load _ | Pbytes_load _ | Pbigstring_load _ | Pbbswap _
   | Pprobe_is_enabled _
-  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
+  | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _ | Pget_header _
     ->
       fatal_errorf "Cmmgen.transl_prim_3: %a"
         Printclambda_primitives.primitive p

bytecomp/bytegen.ml

@@ -114,6 +114,7 @@ let preserve_tailcall_for_prim = function
       true
   | Pbytes_to_string | Pbytes_of_string
   | Parray_to_iarray | Parray_of_iarray
+  | Pget_header _
   | Pignore
   | Pgetglobal _ | Psetglobal _ | Pgetpredef _
   | Pmakeblock _ | Pmakefloatblock _
@@ -531,6 +532,7 @@ let comp_primitive p args =
   | Pbytes_of_string -> Kccall("caml_bytes_of_string", 1)
   | Parray_to_iarray -> Kccall("caml_iarray_of_array", 1)
   | Parray_of_iarray -> Kccall("caml_array_of_iarray", 1)
+  | Pget_header _ -> Kccall("caml_get_header", 1)
   | Pobj_dup -> Kccall("caml_obj_dup", 1)
   (* The cases below are handled in [comp_expr] before the [comp_primitive] call
      (in the order in which they appear below),

lambda/lambda.ml

@@ -246,6 +246,7 @@ type primitive =
   (* Jane Street extensions *)
   | Parray_to_iarray
   | Parray_of_iarray
+  | Pget_header of alloc_mode
 
 and integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge
@@ -1462,7 +1463,7 @@ let primitive_may_allocate : primitive -> alloc_mode option = function
      Some alloc_heap
   | Pstring_load_16 _ | Pbytes_load_16 _ -> None
   | Pstring_load_32 (_, m) | Pbytes_load_32 (_, m)
-  | Pstring_load_64 (_, m) | Pbytes_load_64 (_, m) -> Some m
+  | Pstring_load_64 (_, m) | Pbytes_load_64 (_, m) | Pget_header m -> Some m
   | Pbytes_set_16 _ | Pbytes_set_32 _ | Pbytes_set_64 _ -> None
   | Pbigstring_load_16 _ -> None
   | Pbigstring_load_32 (_,m) | Pbigstring_load_64 (_,m) -> Some m
@@ -1576,6 +1577,7 @@ let primitive_result_layout (p : primitive) =
       (* CR ncourant: use an unboxed int64 here when it exists *)
       layout_any_value
   | (Parray_to_iarray | Parray_of_iarray) -> layout_any_value
+  | Pget_header _ -> layout_boxedint Pnativeint
 
 let compute_expr_layout free_vars_kind lam =
   let rec compute_expr_layout kinds = function

lambda/lambda.mli

@@ -199,6 +199,11 @@ type primitive =
                         one; O(1) *)
   | Parray_of_iarray (* Unsafely reinterpret an immutable array as a mutable
                         one; O(1) *)
+  | Pget_header of alloc_mode
+  (* Get the header of a block. This primitive is invalid if provided with an
+     immediate value.
+     Note: The GC color bits in the header are not reliable except for checking
+     if the value is locally allocated *)
 
 and integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge

lambda/printlambda.ml

@@ -484,6 +484,7 @@ let primitive ppf = function
 
   | Parray_to_iarray -> fprintf ppf "array_to_iarray"
   | Parray_of_iarray -> fprintf ppf "array_of_iarray"
+  | Pget_header m -> fprintf ppf "get_header%s" (alloc_kind m)
 
 let name_of_primitive = function
   | Pbytes_of_string -> "Pbytes_of_string"
@@ -598,6 +599,7 @@ let name_of_primitive = function
   | Pbox_int _ -> "Pbox_int"
   | Parray_of_iarray -> "Parray_of_iarray"
   | Parray_to_iarray -> "Parray_to_iarray"
+  | Pget_header _ -> "Pget_header"
 
 let check_attribute ppf check =
   let check_property = function

lambda/tmc.ml

@@ -921,6 +921,7 @@ let rec choice ctx t =
     | Pbytes_set_16 _ | Pbytes_set_32 _ | Pbytes_set_64 _
     | Pbigstring_load_16 _ | Pbigstring_load_32 _ | Pbigstring_load_64 _
     | Pbigstring_set_16 _ | Pbigstring_set_32 _ | Pbigstring_set_64 _
+    | Pget_header _
     | Pctconst _
     | Pbswap16
     | Pbbswap _

lambda/translprim.ml

@@ -418,6 +418,7 @@ let lookup_primitive loc poly pos p =
     | "%array_of_iarray" -> Primitive (Parray_of_iarray, 1)
     | "%unbox_float" -> Primitive(Punbox_float, 1)
     | "%box_float" -> Primitive(Pbox_float mode, 1)
+    | "%get_header" -> Primitive (Pget_header mode, 1)
     | s when String.length s > 0 && s.[0] = '%' ->
        raise(Error(loc, Unknown_builtin_primitive s))
     | _ -> External p
@@ -966,7 +967,7 @@ let lambda_primitive_needs_event_after = function
   | Pbytes_load_64 _ | Pbytes_set_16 _ | Pbytes_set_32 _ | Pbytes_set_64 _
   | Pbigstring_load_16 _ | Pbigstring_load_32 _ | Pbigstring_load_64 _
   | Pbigstring_set_16 _ | Pbigstring_set_32 _ | Pbigstring_set_64 _
-  | Pbbswap _ | Pobj_dup -> true
+  | Pbbswap _ | Pobj_dup | Pget_header _ -> true
 
   | Pbytes_to_string | Pbytes_of_string
   | Parray_to_iarray | Parray_of_iarray

middle_end/clambda_primitives.ml

@@ -128,6 +128,7 @@ type primitive =
   | Pbox_float of alloc_mode
   | Punbox_int of boxed_integer
   | Pbox_int of boxed_integer * alloc_mode
+  | Pget_header of alloc_mode
 
 and integer_comparison = Lambda.integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge

middle_end/clambda_primitives.mli

@@ -131,6 +131,7 @@ type primitive =
   | Pbox_float of alloc_mode
   | Punbox_int of boxed_integer
   | Pbox_int of boxed_integer * alloc_mode
+  | Pget_header of alloc_mode
 
 and integer_comparison = Lambda.integer_comparison =
     Ceq | Cne | Clt | Cgt | Cle | Cge

middle_end/convert_primitives.ml

@@ -158,6 +158,7 @@ let convert (prim : Lambda.primitive) : Clambda_primitives.primitive =
   | Pbox_float m -> Pbox_float m
   | Punbox_int bi -> Punbox_int bi
   | Pbox_int (bi, m) -> Pbox_int (bi, m)
+  | Pget_header m -> Pget_header m
   | Pobj_magic _
   | Pbytes_to_string
   | Pbytes_of_string

middle_end/internal_variable_names.ml

@@ -176,6 +176,7 @@ let pxorint = "Pxorint"
 let pprobe_is_enabled = "Pprobe_is_enabled"
 let parray_of_iarray = "Parray_of_iarray"
 let parray_to_iarray = "Parray_to_iarray"
+let pget_header = "Pget_header"
 let pabsfloat_arg = "Pabsfloat_arg"
 let paddbint_arg = "Paddbint_arg"
 let paddfloat_arg = "Paddfloat_arg"
@@ -285,6 +286,7 @@ let pxorint_arg = "Pxorint_arg"
 let pprobe_is_enabled_arg = "Pprobe_is_enabled_arg"
 let parray_of_iarray_arg = "Parray_of_iarray_arg"
 let parray_to_iarray_arg = "Parray_to_iarray_arg"
+let pget_header_arg = "Pget_header_arg"
 let raise = "raise"
 let raise_arg = "raise_arg"
 let read_mutable = "read_mutable"
@@ -439,6 +441,7 @@ let of_primitive : Lambda.primitive -> string = function
   | Pbox_int _ -> pbox_int
   | Parray_of_iarray -> parray_of_iarray
   | Parray_to_iarray -> parray_to_iarray
+  | Pget_header _ -> pget_header
 
 let of_primitive_arg : Lambda.primitive -> string = function
   | Pbytes_of_string -> pbytes_of_string_arg
@@ -553,3 +556,4 @@ let of_primitive_arg : Lambda.primitive -> string = function
   | Pbox_int _ -> pbox_int_arg
   | Parray_of_iarray -> parray_of_iarray_arg
   | Parray_to_iarray -> parray_to_iarray_arg
+  | Pget_header _ -> pget_header_arg

middle_end/printclambda_primitives.ml

@@ -259,3 +259,4 @@ let primitive ppf (prim:Clambda_primitives.primitive) =
   | Pbox_int (bi, m) ->
     fprintf ppf "box_%s.%s" (boxed_integer_name bi) (alloc_kind m)
   | Punbox_int bi -> fprintf ppf "unbox_%s" (boxed_integer_name bi)
+  | Pget_header m -> fprintf ppf "get_header.%s" (alloc_kind m)

middle_end/semantics_of_primitives.ml

@@ -144,6 +144,7 @@ let for_primitive (prim : Clambda_primitives.primitive) =
   | Psequor ->
       (* Removed by [Closure_conversion] in the flambda pipeline. *)
       No_effects, No_coeffects
+  | Pget_header _ -> No_effects, No_coeffects
 
 type return_type =
   | Float
@@ -270,3 +271,4 @@ let may_locally_allocate (prim:Clambda_primitives.primitive) : bool =
   | Psequor ->
       false
   | Pprobe_is_enabled _ -> false
+  | Pget_header m -> is_local_alloc m

runtime/obj.c

@@ -73,6 +73,13 @@ CAMLprim value caml_obj_is_local(value blk)
   return Val_int(Is_block(blk) && Color_hd(Hd_val(blk)) == Local_unmarked);
 }
 
+CAMLprim value caml_get_header(value blk)
+{
+  // undefined behaviour if blk is not a block
+  intnat r = Hd_val(blk);
+  return caml_copy_nativeint(r);
+}
+
 /* [size] is a value encoding a number of blocks */
 CAMLprim value caml_obj_block(value tag, value size)
 {

testsuite/tests/lib-obj/get_header.byte.reference

@@ -0,0 +1,3 @@
+None false
+Some(wosize=1,color=0,tag=252) false
+Some(wosize=1,color=0,tag=0) false

testsuite/tests/lib-obj/get_header.ml

@@ -0,0 +1,79 @@
+(* TEST
+  * native
+    reference = "${test_source_directory}/get_header.opt.reference"
+  * bytecode
+    reference = "${test_source_directory}/get_header.byte.reference"
+*)
+
+external repr : ('a[@local_opt]) -> (Obj.t[@local_opt]) = "%identity"
+external get_header_unsafe : (Obj.t[@local_opt]) -> nativeint = "%get_header"
+external is_int : (Obj.t[@local_opt]) -> bool = "%obj_is_int"
+
+let get_header (local_ repr) =
+  if is_int repr then
+    None
+  else
+    Some (get_header_unsafe repr)
+
+type header = {
+  wosize : int;
+  color : int;
+  tag : int
+}
+
+let parse_header : nativeint -> header =
+  fun header ->
+    let wosize =
+      Nativeint.to_int
+        (Nativeint.shift_right_logical header 10)
+    in
+
+    let color =
+      0x3 land
+      (Nativeint.to_int (Nativeint.shift_right_logical header 8))
+    in
+
+    let tag = 0xff land (Nativeint.to_int header) in
+
+    {wosize; color; tag}
+
+let get_header_parsed repr =
+  Option.map parse_header (get_header repr)
+
+let print_header ppf header =
+  let {wosize; color; tag} = header in
+  Format.fprintf ppf "wosize=%i,color=%i,tag=%i" wosize color tag
+
+let print_maybe_header ppf header =
+  match header with
+  | None -> Format.fprintf ppf "None"
+  | Some header -> Format.fprintf ppf "Some(%a)" print_header header
+
+let is_local repr =
+  match get_header_parsed repr with
+  | None -> false
+  | Some {color; _} -> color = 2
+
+(* immediate *)
+let () =
+  let x = 42 in
+  let rp = repr x in
+  Format.printf "%a %a\n" print_maybe_header (get_header_parsed rp)
+    Format.pp_print_bool (is_local rp)
+
+(* global*)
+let () =
+  let s = "hello" in
+  let _r = ref s in
+  let rp = repr s in
+  Format.printf "%a %a\n" print_maybe_header (get_header_parsed rp)
+    Format.pp_print_bool (is_local rp)
+
+(* local *)
+let foo x =
+  let local_ s = ref x in
+  let rp = repr s in
+  Format.printf "%a %a\n" print_maybe_header (get_header_parsed rp)
+    Format.pp_print_bool (is_local rp)
+
+let () = foo 42

testsuite/tests/lib-obj/get_header.opt.reference

@@ -0,0 +1,3 @@
+None false
+Some(wosize=1,color=3,tag=252) false
+Some(wosize=1,color=2,tag=0) true