add the %get_header primitive (#1539)

backend/cmmgen.ml

@@ -722,7 +722,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
@@ -1036,6 +1036,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
@@ -1231,7 +1233,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
@@ -1292,7 +1294,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

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/flambda2/from_lambda/closure_conversion.ml

@@ -755,7 +755,7 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
       | Pbigstring_set_64 _ | Pctconst _ | Pbswap16 | Pbbswap _
       | Pint_as_pointer _ | Popaque _ | Pprobe_is_enabled _ | Pobj_dup
       | Pobj_magic _ | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
-        ->
+      | Pget_header _ ->
         (* Inconsistent with outer match *)
         assert false
     in

middle_end/flambda2/from_lambda/lambda_to_flambda.ml

@@ -1004,7 +1004,7 @@ let primitive_can_raise (prim : Lambda.primitive) =
   | Pbigstring_set_64 true
   | Pctconst _ | Pbswap16 | Pbbswap _ | Pint_as_pointer _ | Popaque _
   | Pprobe_is_enabled _ | Pobj_dup | Pobj_magic _ | Pbox_float _ | Punbox_float
-  | Punbox_int _ | Pbox_int _ ->
+  | Punbox_int _ | Pbox_int _ | Pget_header _ ->
     false
 
 type cps_continuation =

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -385,6 +385,10 @@ let string_like_load_unsafe ~access_size kind mode string index ~current_region
   in
   wrap (Binary (String_or_bigstring_load (kind, access_size), string, index))
 
+let get_header obj mode ~current_region =
+  let wrap hd = box_bint Pnativeint mode hd ~current_region in
+  wrap (Unary (Get_header, obj))
+
 let string_like_load_safe ~dbg ~size_int ~access_size kind mode str index
     ~current_region =
   match (kind : P.string_like_value) with
@@ -1257,6 +1261,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
   | Pprobe_is_enabled { name }, [] ->
     [tag_int (Nullary (Probe_is_enabled { name }))]
   | Pobj_dup, [[v]] -> [Unary (Obj_dup, v)]
+  | Pget_header m, [[obj]] -> [get_header obj m ~current_region]
   | ( ( Pmodint Unsafe
       | Pdivbint { is_safe = Unsafe; size = _; mode = _ }
       | Pmodbint { is_safe = Unsafe; size = _; mode = _ }
@@ -1279,7 +1284,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       | Pduparray _ | Pfloatfield _ | Pcvtbint _ | Poffsetref _ | Pbswap16
       | Pbbswap _ | Pisint _ | Pint_as_pointer _ | Pbigarraydim _ | Pobj_dup
       | Pobj_magic _ | Punbox_float | Pbox_float _ | Punbox_int _ | Pbox_int _
-        ),
+      | Pget_header _ ),
       ([] | _ :: _ :: _ | [([] | _ :: _ :: _)]) ) ->
     Misc.fatal_errorf
       "Closure_conversion.convert_primitive: Wrong arity for unary primitive \

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -539,7 +539,8 @@ let unop env (op : Flambda_primitive.unary_primitive) : Fexpr.unop =
   | String_length string_or_bytes -> String_length string_or_bytes
   | Boolean_not -> Boolean_not
   | Int_as_pointer _ | Duplicate_block _ | Duplicate_array _ | Bigarray_length _
-  | Float_arith _ | Reinterpret_int64_as_float | Is_boxed_float | Obj_dup ->
+  | Float_arith _ | Reinterpret_int64_as_float | Is_boxed_float | Obj_dup
+  | Get_header ->
     Misc.fatal_errorf "TODO: Unary primitive: %a"
       Flambda_primitive.Without_args.print
       (Flambda_primitive.Without_args.Unary op)

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -592,6 +592,12 @@ let simplify_obj_dup dbg dacc ~original_term ~arg ~arg_ty ~result_var =
     | Proved ((Heap_or_local | Local), _) | Unknown ->
       SPR.create_unknown dacc ~result_var K.value ~original_term)
 
+let simplify_get_header ~original_prim dacc ~original_term ~arg:_ ~arg_ty:_
+    ~result_var =
+  SPR.create_unknown dacc ~result_var
+    (P.result_kind' original_prim)
+    ~original_term
+
 let simplify_unary_primitive dacc original_prim (prim : P.unary_primitive) ~arg
     ~arg_ty dbg ~result_var =
   let min_name_mode = Bound_var.name_mode result_var in
@@ -642,5 +648,6 @@ let simplify_unary_primitive dacc original_prim (prim : P.unary_primitive) ~arg
     | Begin_try_region -> simplify_begin_try_region
     | End_region -> simplify_end_region
     | Obj_dup -> simplify_obj_dup dbg
+    | Get_header -> simplify_get_header ~original_prim
   in
   simplifier dacc ~original_term ~arg ~arg_ty ~result_var

middle_end/flambda2/terms/code_size.ml

@@ -336,6 +336,7 @@ let unary_prim_size prim =
   | Begin_try_region -> 1
   | End_region -> 1
   | Obj_dup -> alloc_extcall_size + 1
+  | Get_header -> 2
 
 let binary_prim_size prim =
   match (prim : Flambda_primitive.binary_primitive) with

middle_end/flambda2/terms/flambda_primitive.ml

@@ -730,14 +730,15 @@ type unary_primitive =
   | Begin_try_region
   | End_region
   | Obj_dup
+  | Get_header
 
 (* Here and below, operations that are genuine projections shouldn't be eligible
    for CSE, since we deal with projections through types. *)
 let unary_primitive_eligible_for_cse p ~arg =
   match p with
   | Duplicate_array _ -> false
   | Duplicate_block { kind = _ } -> false
-  | Is_int _ | Get_tag -> true
+  | Is_int _ | Get_tag | Get_header -> true
   | Array_length -> true
   | Bigarray_length _ -> false
   | String_length _ -> true
@@ -790,6 +791,7 @@ let compare_unary_primitive p1 p2 =
     | Begin_try_region -> 22
     | End_region -> 23
     | Obj_dup -> 24
+    | Get_header -> 25
   in
   match p1, p2 with
   | ( Duplicate_array
@@ -867,7 +869,7 @@ let compare_unary_primitive p1 p2 =
       | Array_length | Bigarray_length _ | Unbox_number _ | Box_number _
       | Untag_immediate | Tag_immediate | Project_function_slot _
       | Project_value_slot _ | Is_boxed_float | Is_flat_float_array
-      | Begin_try_region | End_region | Obj_dup ),
+      | Begin_try_region | End_region | Obj_dup | Get_header ),
       _ ) ->
     Stdlib.compare (unary_primitive_numbering p1) (unary_primitive_numbering p2)
 
@@ -921,6 +923,7 @@ let print_unary_primitive ppf p =
   | Begin_try_region -> Format.pp_print_string ppf "Begin_try_region"
   | End_region -> Format.pp_print_string ppf "End_region"
   | Obj_dup -> Format.pp_print_string ppf "Obj_dup"
+  | Get_header -> Format.pp_print_string ppf "Get_header"
 
 let arg_kind_of_unary_primitive p =
   match p with
@@ -945,6 +948,7 @@ let arg_kind_of_unary_primitive p =
   | Begin_try_region -> K.region
   | End_region -> K.region
   | Obj_dup -> K.value
+  | Get_header -> K.value
 
 let result_kind_of_unary_primitive p : result_kind =
   match p with
@@ -971,6 +975,7 @@ let result_kind_of_unary_primitive p : result_kind =
   | Begin_try_region -> Singleton K.region
   | End_region -> Singleton K.value
   | Obj_dup -> Singleton K.value
+  | Get_header -> Singleton K.naked_nativeint
 
 let effects_and_coeffects_of_unary_primitive p : Effects_and_coeffects.t =
   match p with
@@ -1058,6 +1063,7 @@ let effects_and_coeffects_of_unary_primitive p : Effects_and_coeffects.t =
     ( Only_generative_effects Mutable (* Mutable is conservative *),
       Has_coeffects,
       Strict )
+  | Get_header -> No_effects, No_coeffects, Strict
 
 let unary_classify_for_printing p =
   match p with
@@ -1072,6 +1078,7 @@ let unary_classify_for_printing p =
   | Project_function_slot _ | Project_value_slot _ -> Destructive
   | Is_boxed_float | Is_flat_float_array -> Neither
   | Begin_try_region | End_region -> Neither
+  | Get_header -> Neither
 
 let free_names_unary_primitive p =
   match p with
@@ -1092,7 +1099,7 @@ let free_names_unary_primitive p =
   | Boolean_not | Reinterpret_int64_as_float | Float_arith _ | Array_length
   | Bigarray_length _ | Unbox_number _ | Untag_immediate | Tag_immediate
   | Is_boxed_float | Is_flat_float_array | Begin_try_region | End_region
-  | Obj_dup ->
+  | Obj_dup | Get_header ->
     Name_occurrences.empty
 
 let apply_renaming_unary_primitive p renaming =
@@ -1107,7 +1114,7 @@ let apply_renaming_unary_primitive p renaming =
   | Boolean_not | Reinterpret_int64_as_float | Float_arith _ | Array_length
   | Bigarray_length _ | Unbox_number _ | Untag_immediate | Tag_immediate
   | Is_boxed_float | Is_flat_float_array | Begin_try_region | End_region
-  | Project_function_slot _ | Project_value_slot _ | Obj_dup ->
+  | Project_function_slot _ | Project_value_slot _ | Obj_dup | Get_header ->
     p
 
 let ids_for_export_unary_primitive p =
@@ -1119,7 +1126,7 @@ let ids_for_export_unary_primitive p =
   | Boolean_not | Reinterpret_int64_as_float | Float_arith _ | Array_length
   | Bigarray_length _ | Unbox_number _ | Untag_immediate | Tag_immediate
   | Is_boxed_float | Is_flat_float_array | Begin_try_region | End_region
-  | Project_function_slot _ | Project_value_slot _ | Obj_dup ->
+  | Project_function_slot _ | Project_value_slot _ | Obj_dup | Get_header ->
     Ids_for_export.empty
 
 type binary_int_arith_op =

middle_end/flambda2/terms/flambda_primitive.mli

@@ -323,6 +323,11 @@ type unary_primitive =
   | End_region
       (** Ending delimiter of local allocation region, accepting a region name. *)
   | Obj_dup  (** Corresponds to [Obj.dup]; see the documentation in obj.mli. *)
+  | Get_header
+      (** 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 *)
 
 (** Whether a comparison is to yield a boolean result, as given by a particular
     comparison operator, or whether it is to behave in the manner of "compare"

middle_end/flambda2/to_cmm/to_cmm_primitive.ml

@@ -632,6 +632,7 @@ let unary_primitive env res dbg f arg =
     None, res, C.eq ~dbg (C.get_tag arg dbg) (C.floatarray_tag dbg)
   | Begin_try_region -> None, res, C.beginregion ~dbg
   | End_region -> None, res, C.return_unit dbg (C.endregion ~dbg arg)
+  | Get_header -> None, res, C.get_header arg dbg
 
 let binary_primitive env dbg f x y =
   match (f : P.binary_primitive) with

middle_end/internal_variable_names.ml

@@ -180,6 +180,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"
@@ -289,6 +290,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"
@@ -444,6 +446,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
@@ -558,3 +561,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

@@ -156,6 +156,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
@@ -282,3 +283,4 @@ let may_locally_allocate (prim:Clambda_primitives.primitive) : bool =
   | Psequor ->
       false
   | Pprobe_is_enabled _ -> false
+  | Pget_header m -> is_local_alloc m

ocaml/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

ocaml/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),

ocaml/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