Implement %array_element_size_in_bytes

bytecomp/bytegen.ml

@@ -158,6 +158,7 @@ let preserve_tailcall_for_prim = function
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Pduprecord _
   | Pufloatfield _ | Psetufloatfield _ | Pmixedfield _ | Psetmixedfield _
   | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Parray_element_size_in_bytes _
   | Pccall _ | Praise _ | Pnot | Pnegint | Paddint | Psubint | Pmulint
   | Pdivint _ | Pmodint _ | Pandint | Porint | Pxorint | Plslint | Plsrint
   | Pasrint | Pintcomp _ | Poffsetint _ | Poffsetref _ | Pintoffloat _
@@ -414,6 +415,8 @@ let comp_primitive stack_info p sz args =
   | Pcompare_bints bi -> comp_bint_primitive bi "compare" args
   | Pfield (n, _ptr, _sem) -> Kgetfield n
   | Punboxed_product_field (n, _layouts) -> Kgetfield n
+  | Parray_element_size_in_bytes _array_kind ->
+      Kconst (Const_base (Const_int (Sys.word_size / 8)))
   | Pfield_computed _sem -> Kgetvectitem
   | Psetfield(n, _ptr, _init) -> Ksetfield n
   | Psetfield_computed(_ptr, _init) -> Ksetvectitem

lambda/lambda.ml

@@ -149,6 +149,7 @@ type primitive =
   (* Unboxed products *)
   | Pmake_unboxed_product of layout list
   | Punboxed_product_field of int * layout list
+  | Parray_element_size_in_bytes of array_kind
   (* Context switches *)
   | Prunstack
   | Pperform
@@ -1930,7 +1931,8 @@ let primitive_may_allocate : primitive -> locality_mode option = function
   | Patomic_cas
   | Patomic_fetch_add
   | Pdls_get
-  | Preinterpret_unboxed_int64_as_tagged_int63 -> None
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Parray_element_size_in_bytes _ -> None
   | Preinterpret_tagged_int63_as_unboxed_int64 ->
     if !Clflags.native_code then None
     else
@@ -2094,7 +2096,8 @@ let primitive_can_raise prim =
   | Patomic_exchange | Patomic_cas | Patomic_fetch_add | Patomic_load _ -> false
   | Prunstack | Pperform | Presume | Preperform -> true (* XXX! *)
   | Pdls_get | Ppoll | Preinterpret_tagged_int63_as_unboxed_int64
-  | Preinterpret_unboxed_int64_as_tagged_int63 ->
+  | Preinterpret_unboxed_int64_as_tagged_int63
+  | Parray_element_size_in_bytes _ ->
     false
 
 let constant_layout: constant -> layout = function
@@ -2206,6 +2209,7 @@ let primitive_result_layout (p : primitive) =
   | Pfield _ | Pfield_computed _ -> layout_value_field
   | Punboxed_product_field (field, layouts) -> (Array.of_list layouts).(field)
   | Pmake_unboxed_product layouts -> layout_unboxed_product layouts
+  | Parray_element_size_in_bytes _ -> layout_int
   | Pfloatfield _ -> layout_boxed_float Pfloat64
   | Pfloatoffloat32 _ -> layout_boxed_float Pfloat64
   | Pfloat32offloat _ -> layout_boxed_float Pfloat32
@@ -2479,3 +2483,34 @@ let rec try_to_find_location lam =
 
 let try_to_find_debuginfo lam =
   Debuginfo.from_location (try_to_find_location lam)
+
+let rec count_initializers_scannable
+    (scannable : scannable_product_element_kind) =
+  match scannable with
+  | Pint_scannable | Paddr_scannable -> 1
+  | Pproduct_scannable scannables ->
+    List.fold_left
+      (fun acc scannable -> acc + count_initializers_scannable scannable)
+      0 scannables
+
+let rec count_initializers_ignorable
+    (ignorable : ignorable_product_element_kind) =
+  match ignorable with
+  | Pint_ignorable | Punboxedfloat_ignorable _ | Punboxedint_ignorable _ -> 1
+  | Pproduct_ignorable ignorables ->
+    List.fold_left
+      (fun acc ignorable -> acc + count_initializers_ignorable ignorable)
+      0 ignorables
+
+let count_initializers_array_kind (lambda_array_kind : array_kind) =
+  match lambda_array_kind with
+  | Pgenarray | Paddrarray | Pintarray | Pfloatarray | Punboxedfloatarray _
+  | Punboxedintarray _ | Punboxedvectorarray _ -> 1
+  | Pgcscannableproductarray scannables ->
+    List.fold_left
+      (fun acc scannable -> acc + count_initializers_scannable scannable)
+      0 scannables
+  | Pgcignorableproductarray ignorables ->
+    List.fold_left
+      (fun acc ignorable -> acc + count_initializers_ignorable ignorable)
+      0 ignorables

lambda/lambda.mli

@@ -137,6 +137,7 @@ type primitive =
   | Pmake_unboxed_product of layout list
   | Punboxed_product_field of int * (layout list)
       (* the [layout list] is the layout of the whole product *)
+  | Parray_element_size_in_bytes of array_kind
   (* Context switches *)
   | Prunstack
   | Pperform
@@ -1172,3 +1173,5 @@ val try_to_find_location : lambda -> scoped_location
 val try_to_find_debuginfo : lambda -> Debuginfo.t
 
 val primitive_can_raise : primitive -> bool
+
+val count_initializers_array_kind : array_kind -> int

lambda/printlambda.ml

@@ -606,6 +606,8 @@ let primitive ppf = function
       fprintf ppf "unboxed_product_field %d #(%a)" n
         (pp_print_list ~pp_sep:(fun ppf () -> fprintf ppf ", ") (layout' false))
         layouts
+  | Parray_element_size_in_bytes ak ->
+      fprintf ppf "array_element_size_in_bytes (%s)" (array_kind ak)
   | Pccall p -> fprintf ppf "%s" p.prim_name
   | Praise k -> fprintf ppf "%s" (Lambda.raise_kind k)
   | Psequand -> fprintf ppf "&&"
@@ -941,6 +943,7 @@ let name_of_primitive = function
   | Pduprecord _ -> "Pduprecord"
   | Pmake_unboxed_product _ -> "Pmake_unboxed_product"
   | Punboxed_product_field _ -> "Punboxed_product_field"
+  | Parray_element_size_in_bytes _ -> "Parray_element_size_in_bytes"
   | Pccall _ -> "Pccall"
   | Praise _ -> "Praise"
   | Psequand -> "Psequand"

lambda/tmc.ml

@@ -922,6 +922,7 @@ let rec choice ctx t =
 
     (* nor unboxed products *)
     | Pmake_unboxed_product _ | Punboxed_product_field _
+    | Parray_element_size_in_bytes _
 
     | Pobj_dup
     | Pobj_magic _

lambda/translprim.ml

@@ -536,6 +536,9 @@ let lookup_primitive loc ~poly_mode ~poly_sort pos p =
     | "%arrayblit" ->
       Language_extension.assert_enabled ~loc Layouts Language_extension.Alpha;
       Primitive (Parrayblit (gen_array_set_kind (get_third_arg_mode ())), 5)
+    | "%array_element_size_in_bytes" ->
+      (* The array kind will be filled in later *)
+      Primitive (Parray_element_size_in_bytes Pgenarray, 1)
     | "%obj_size" -> Primitive ((Parraylength Pgenarray), 1)
     | "%obj_field" -> Primitive ((Parrayrefu (Pgenarray_ref mode, Ptagged_int_index, Mutable)), 2)
     | "%obj_set_field" ->
@@ -1253,6 +1256,12 @@ let specialize_primitive env loc ty ~has_constant_constructor prim =
     in
     if st = array_type then None
     else Some (Primitive (Parrayblit array_type, arity))
+  | Primitive (Parray_element_size_in_bytes _, arity), p1 :: _ -> (
+      let array_kind =
+        array_type_kind ~elt_sort:None env (to_location loc) p1
+      in
+      Some (Primitive (Parray_element_size_in_bytes array_kind, arity))
+    )
   | Primitive (Pbigarrayref(unsafe, n, kind, layout), arity), p1 :: _ -> begin
       let (k, l) = bigarray_specialize_kind_and_layout env ~kind ~layout p1 in
       match k, l with
@@ -1772,6 +1781,7 @@ let lambda_primitive_needs_event_after = function
   | Pgetglobal _ | Pgetpredef _ | Pmakeblock _ | Pmakefloatblock _
   | Pmakeufloatblock _ | Pmakemixedblock _
   | Pmake_unboxed_product _ | Punboxed_product_field _
+  | Parray_element_size_in_bytes _
   | Pfield _ | Pfield_computed _ | Psetfield _
   | Psetfield_computed _ | Pfloatfield _ | Psetfloatfield _ | Praise _
   | Pufloatfield _ | Psetufloatfield _ | Pmixedfield _ | Psetmixedfield _

lambda/value_rec_compiler.ml

@@ -366,7 +366,8 @@ let compute_static_size lam =
     | Punboxed_float32_array_set_128 _
     | Punboxed_int32_array_set_128 _
     | Punboxed_int64_array_set_128 _
-    | Punboxed_nativeint_array_set_128 _ ->
+    | Punboxed_nativeint_array_set_128 _
+    | Parray_element_size_in_bytes _ ->
         Constant
 
     | Pmakeufloatblock (_, _)

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1041,10 +1041,10 @@ let close_primitive acc env ~let_bound_ids_with_kinds named
       | Punbox_vector _
       | Pbox_vector (_, _)
       | Punbox_int _ | Pbox_int _ | Pmake_unboxed_product _
-      | Punboxed_product_field _ | Pget_header _ | Prunstack | Pperform
-      | Presume | Preperform | Patomic_exchange | Patomic_cas
-      | Patomic_fetch_add | Pdls_get | Ppoll | Patomic_load _
-      | Preinterpret_tagged_int63_as_unboxed_int64
+      | Punboxed_product_field _ | Parray_element_size_in_bytes _
+      | Pget_header _ | Prunstack | Pperform | Presume | Preperform
+      | Patomic_exchange | Patomic_cas | Patomic_fetch_add | Pdls_get | Ppoll
+      | Patomic_load _ | Preinterpret_tagged_int63_as_unboxed_int64
       | Preinterpret_unboxed_int64_as_tagged_int63 ->
         (* Inconsistent with outer match *)
         assert false

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -1322,6 +1322,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     let mutability = Mutability.from_lambda mutability in
     [Variadic (Make_block (Values (tag, shape), mutability, mode), args)]
   | Pmake_unboxed_product layouts, _ ->
+    (* CR mshinwell: this should check the unarized lengths of [layouts] and
+       [args] (like [Parray_element_size_in_bytes] below) *)
     if List.compare_lengths layouts args <> 0
     then
       Misc.fatal_errorf "Pmake_unboxed_product: expected %d arguments, got %d"
@@ -1351,6 +1353,26 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       |> Array.to_list
     in
     List.map (fun arg : H.expr_primitive -> Simple arg) projected_args
+  | Parray_element_size_in_bytes array_kind, [_witness] ->
+    (* This is implemented as a unary primitive, but from our point of view it's
+       actually nullary. *)
+    let num_bytes =
+      match array_kind with
+      | Pgenarray | Paddrarray | Pintarray | Pfloatarray -> 8
+      | Punboxedfloatarray Pfloat32 ->
+        (* float32# arrays are packed *)
+        4
+      | Punboxedfloatarray Pfloat64 -> 8
+      | Punboxedintarray Pint32 ->
+        (* int32# arrays are packed *)
+        4
+      | Punboxedintarray (Pint64 | Pnativeint) -> 8
+      | Punboxedvectorarray Pvec128 -> 16
+      | Pgcscannableproductarray _ | Pgcignorableproductarray _ ->
+        (* All elements of unboxed product arrays are currently 8 bytes wide. *)
+        L.count_initializers_array_kind array_kind * 8
+    in
+    [Simple (Simple.const_int (Targetint_31_63.of_int num_bytes))]
   | Pmakefloatblock (mutability, mode), _ ->
     let args = List.flatten args in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
@@ -2348,7 +2370,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       | Punbox_int _ | Pbox_int _ | Punboxed_product_field _ | Pget_header _
       | Pufloatfield _ | Patomic_load _ | Pmixedfield _
       | Preinterpret_unboxed_int64_as_tagged_int63
-      | Preinterpret_tagged_int63_as_unboxed_int64 ),
+      | Preinterpret_tagged_int63_as_unboxed_int64
+      | Parray_element_size_in_bytes _ ),
       ([] | _ :: _ :: _ | [([] | _ :: _ :: _)]) ) ->
     Misc.fatal_errorf
       "Closure_conversion.convert_primitive: Wrong arity for unary primitive \

typing/primitive.ml

@@ -655,6 +655,12 @@ let prim_has_valid_reprs ~loc prim =
         is (Same_as_ocaml_repr C.value);
       ]
 
+    | "%array_element_size_in_bytes" ->
+      check [
+        any;
+        is (Same_as_ocaml_repr C.value);
+      ]
+
     | "%box_float" ->
       exactly [Same_as_ocaml_repr C.float64; Same_as_ocaml_repr C.value]
     | "%unbox_float" ->