Support unboxed ints in mixed blocks

backend/cmm_helpers.ml

@@ -1246,6 +1246,41 @@ let get_field_computed imm_or_ptr mutability ~block ~index dbg =
   Cop
     (Cload { memory_chunk; mutability; is_atomic = false }, [field_address], dbg)
 
+(* Getters for unboxed int fields *)
+
+let get_field_unboxed_int32 mutability ~block ~index dbg =
+  (* CR mixed blocks: Could a backend person check the "signed"? *)
+  let memory_chunk = Thirtytwo_signed in
+  (* CR layouts v5.1: We'll need to vary log2_size_addr to efficiently pack
+   * int32s *)
+  let field_address = array_indexing log2_size_addr block index dbg in
+  Cop
+    (Cload { memory_chunk; mutability; is_atomic = false }, [field_address], dbg)
+
+let get_field_unboxed_int64_or_nativeint mutability ~block ~index dbg =
+  let memory_chunk = Word_int in
+  let field_address = array_indexing log2_size_addr block index dbg in
+  Cop
+    (Cload { memory_chunk; mutability; is_atomic = false }, [field_address], dbg)
+
+(* Setters for unboxed int fields *)
+
+let setfield_unboxed_int32 arr ofs newval dbg =
+  (* CR layouts v5.1: We will need to vary log2_size_addr when int32 fields are
+     efficiently packed. *)
+  return_unit dbg
+    (Cop
+       ( Cstore (Thirtytwo_signed, Assignment),
+         [array_indexing log2_size_addr arr ofs dbg; newval],
+         dbg ))
+
+let setfield_unboxed_int64_or_nativeint arr ofs newval dbg =
+  return_unit dbg
+    (Cop
+       ( Cstore (Word_int, Assignment),
+         [array_indexing log2_size_addr arr ofs dbg; newval],
+         dbg ))
+
 (* String length *)
 
 (* Length of string block *)
@@ -1525,16 +1560,19 @@ let make_mixed_alloc ~mode dbg tag shape args =
       match flat_suffix.(idx - value_prefix_len) with
       | Imm -> int_array_set arr ofs newval dbg
       | Float | Float64 -> float_array_set arr ofs newval dbg
+      | Bits32 -> setfield_unboxed_int32 arr ofs newval dbg
+      | Bits64 | Word -> setfield_unboxed_int64_or_nativeint arr ofs newval dbg
   in
   let size =
-    let values, floats = Lambda.count_mixed_block_values_and_floats shape in
-    if size_float <> size_addr
-    then
-      Misc.fatal_error
-        "Unable to compile mixed blocks on a platform where a float is not the \
-         same width as a value.";
-    values + floats
+    (* CR layouts 5.1: When we pack int32s more efficiently, this code will need
+       to change. *)
+    value_prefix_len + Array.length flat_suffix
   in
+  if size_float <> size_addr
+  then
+    Misc.fatal_error
+      "Unable to compile mixed blocks on a platform where a float is not the \
+       same width as a value.";
   make_alloc_generic ~scannable_prefix:(Scan_prefix value_prefix_len) ~mode
     set_fn dbg tag size args
 

backend/cmm_helpers.mli

@@ -972,3 +972,31 @@ val unboxed_int64_or_nativeint_array_set :
   new_value:expression ->
   Debuginfo.t ->
   expression
+
+(** {2 Getters and setters for unboxed int fields of mixed blocks} *)
+
+(** The argument structure for getters is parallel to [get_field_computed]. *)
+
+val get_field_unboxed_int32 :
+  Asttypes.mutable_flag ->
+  block:expression ->
+  index:expression ->
+  Debuginfo.t ->
+  expression
+
+val get_field_unboxed_int64_or_nativeint :
+  Asttypes.mutable_flag ->
+  block:expression ->
+  index:expression ->
+  Debuginfo.t ->
+  expression
+
+(** The argument structure for setters is parallel to [setfield_computed].
+   [immediate_or_pointer] is not needed as the layout is implied from the name,
+   and [initialization_or_assignment] is not needed as unboxed ints can always be
+   assigned without caml_modify (etc.).
+ *)
+
+val setfield_unboxed_int32 : ternary_primitive
+
+val setfield_unboxed_int64_or_nativeint : ternary_primitive

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -991,13 +991,16 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       List.mapi
         (fun i arg ->
           match Lambda.get_mixed_block_element shape i with
-          | Value_prefix | Flat_suffix (Float64 | Imm) -> arg
+          | Value_prefix | Flat_suffix (Float64 | Imm | Bits32 | Bits64 | Word)
+            ->
+            arg
           | Flat_suffix Float -> unbox_float arg)
         args
     in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
     let mutability = Mutability.from_lambda mutability in
     let tag = Tag.Scannable.create_exn tag in
+    let shape = P.Mixed_block_kind.from_lambda shape in
     [Variadic (Make_mixed_block (tag, shape, mutability, mode), args)]
   | Pmakearray (lambda_array_kind, mutability, mode), _ -> (
     let args = List.flatten args in
@@ -1420,19 +1423,17 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
         | Mread_flat_suffix read ->
           Flat_suffix
             (match read with
-            | Flat_read_imm -> Imm
-            | Flat_read_float _ -> Float
-            | Flat_read_float64 -> Float64)
+            | Flat_read flat_element ->
+              P.Mixed_block_flat_element.from_lambda flat_element
+            | Flat_read_float _ -> Float)
       in
       Mixed { tag = Unknown; field_kind; size = Unknown }
     in
     let block_access : H.expr_primitive =
       Binary (Block_load (block_access, mutability), arg, Simple field)
     in
     match read with
-    | Mread_value_prefix _
-    | Mread_flat_suffix (Flat_read_imm | Flat_read_float64) ->
-      [block_access]
+    | Mread_value_prefix _ | Mread_flat_suffix (Flat_read _) -> [block_access]
     | Mread_flat_suffix (Flat_read_float mode) ->
       [box_float mode block_access ~current_region])
   | ( Psetfield (index, immediate_or_pointer, initialization_or_assignment),
@@ -1484,15 +1485,18 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
             | Mwrite_value_prefix immediate_or_pointer ->
               Value_prefix
                 (convert_block_access_field_kind immediate_or_pointer)
-            | Mwrite_flat_suffix flat -> Flat_suffix flat);
+            | Mwrite_flat_suffix flat ->
+              Flat_suffix (P.Mixed_block_flat_element.from_lambda flat));
           size = Unknown;
           tag = Unknown
         }
     in
     let init_or_assign = convert_init_or_assign initialization_or_assignment in
     let value =
       match write with
-      | Mwrite_value_prefix _ | Mwrite_flat_suffix (Imm | Float64) -> value
+      | Mwrite_value_prefix _
+      | Mwrite_flat_suffix (Imm | Float64 | Bits32 | Bits64 | Word) ->
+        value
       | Mwrite_flat_suffix Float -> unbox_float value
     in
     [ Ternary

middle_end/flambda2/parser/print_fexpr.ml

@@ -465,9 +465,9 @@ let block_access_kind ppf (access_kind : block_access_kind) =
     match field_kind with
     | Value_prefix Any_value -> ()
     | Value_prefix Immediate -> Format.fprintf ppf "@ imm"
-    | Flat_suffix Float -> Format.fprintf ppf "@ float"
-    | Flat_suffix Imm -> Format.fprintf ppf "@ imm"
-    | Flat_suffix Float64 -> Format.fprintf ppf "@ float64"
+    | Flat_suffix flat ->
+      Format.fprintf ppf "@ %s"
+        (Flambda_primitive.Mixed_block_flat_element.to_string flat)
   in
   match access_kind with
   | Values { field_kind; tag; size } ->

middle_end/flambda2/terms/code_size.ml

@@ -160,7 +160,11 @@ let block_set (kind : Flambda_primitive.Block_access_kind.t)
   | Values _, (Assignment Local | Initialization) -> 1 (* cadda + store *)
   | Naked_floats _, (Assignment _ | Initialization) -> 1
   | ( Mixed
-        { field_kind = Value_prefix _ | Flat_suffix (Imm | Float | Float64); _ },
+        { field_kind =
+            ( Value_prefix _
+            | Flat_suffix (Imm | Float | Float64 | Bits32 | Bits64 | Word) );
+          _
+        },
       (Assignment _ | Initialization) ) ->
     1
 

middle_end/flambda2/terms/flambda_primitive.ml

@@ -58,31 +58,95 @@ module Block_kind = struct
     match t with Values _ -> K.value | Naked_floats -> K.naked_float
 end
 
+module Mixed_block_flat_element = struct
+  type t =
+    | Imm
+    | Float
+    | Float64
+    | Bits32
+    | Bits64
+    | Word
+
+  let from_lambda : Lambda.flat_element -> t = function
+    | Imm -> Imm
+    | Float -> Float
+    | Float64 -> Float64
+    | Bits32 -> Bits32
+    | Bits64 -> Bits64
+    | Word -> Word
+
+  let to_lambda : t -> Lambda.flat_element = function
+    | Imm -> Imm
+    | Float -> Float
+    | Float64 -> Float64
+    | Bits32 -> Bits32
+    | Bits64 -> Bits64
+    | Word -> Word
+
+  let to_string = function
+    | Imm -> "Imm"
+    | Float -> "Float"
+    | Float64 -> "Float64"
+    | Bits32 -> "Bits32"
+    | Bits64 -> "Bits64"
+    | Word -> "Word"
+
+  let compare t1 t2 =
+    match t1, t2 with
+    | Imm, Imm
+    | Float, Float
+    | Float64, Float64
+    | Word, Word
+    | Bits32, Bits32
+    | Bits64, Bits64 ->
+      0
+    | Imm, _ -> -1
+    | _, Imm -> 1
+    | Float, _ -> -1
+    | _, Float -> 1
+    | Float64, _ -> -1
+    | _, Float64 -> 1
+    | Word, _ -> -1
+    | _, Word -> 1
+    | Bits32, _ -> -1
+    | _, Bits32 -> 1
+
+  let print ppf t = Format.fprintf ppf "%s" (to_string t)
+
+  let element_kind = function
+    | Imm -> K.value
+    | Float | Float64 -> K.naked_float
+    | Bits32 -> K.naked_int32
+    | Bits64 -> K.naked_int64
+    | Word -> K.naked_nativeint
+end
+
 module Mixed_block_kind = struct
-  type t = Lambda.mixed_block_shape
+  type t =
+    { value_prefix_len : int;
+      (* We use an array just so we can index into the middle. *)
+      flat_suffix : Mixed_block_flat_element.t array
+    }
+
+  let from_lambda { Lambda.value_prefix_len; flat_suffix } =
+    { value_prefix_len;
+      flat_suffix = Array.map Mixed_block_flat_element.from_lambda flat_suffix
+    }
 
-  let print_flat_element ppf (e : Lambda.flat_element) =
-    match e with
-    | Imm -> Format.fprintf ppf "Imm"
-    | Float -> Format.fprintf ppf "Float"
-    | Float64 -> Format.fprintf ppf "Float64"
+  let to_lambda { value_prefix_len; flat_suffix } : Lambda.mixed_block_shape =
+    { value_prefix_len;
+      flat_suffix = Array.map Mixed_block_flat_element.to_lambda flat_suffix
+    }
 
   let print ppf ({ value_prefix_len; flat_suffix } : t) =
     Format.fprintf ppf "[|@ ";
     Format.fprintf ppf "Value (x%d);@ " value_prefix_len;
     Array.iter
-      (fun elem -> Format.fprintf ppf "%a;@ " print_flat_element elem)
+      (fun elem ->
+        Format.fprintf ppf "%a;@ " Mixed_block_flat_element.print elem)
       flat_suffix;
     Format.fprintf ppf "|]"
 
-  let compare_flat_element e1 e2 =
-    match (e1 : Lambda.flat_element), (e2 : Lambda.flat_element) with
-    | Imm, Imm | Float, Float | Float64, Float64 -> 0
-    | Imm, _ -> -1
-    | _, Imm -> 1
-    | Float, _ -> -1
-    | _, Float -> 1
-
   let compare (t1 : t) (t2 : t) =
     let components (t : t) =
       let (({ value_prefix_len; flat_suffix } [@warning "+9"]) : t) = t in
@@ -91,19 +155,18 @@ module Mixed_block_kind = struct
     let v1, a1 = components t1 in
     let v2, a2 = components t2 in
     match Int.compare v1 v2 with
-    | 0 -> Misc.Stdlib.Array.compare compare_flat_element a1 a2
+    | 0 -> Misc.Stdlib.Array.compare Mixed_block_flat_element.compare a1 a2
     | cmp -> cmp
 
   let length ({ value_prefix_len; flat_suffix } : t) =
     value_prefix_len + Array.length flat_suffix
 
-  let element_kind i ({ value_prefix_len; flat_suffix } : t) =
+  let element_kind i { value_prefix_len; flat_suffix } =
+    if i < 0 then Misc.fatal_errorf "Negative index: %d" i;
     if i < value_prefix_len
     then K.value
     else
-      match flat_suffix.(i - value_prefix_len) with
-      | Imm -> K.value
-      | Float | Float64 -> K.naked_float
+      Mixed_block_flat_element.element_kind flat_suffix.(i - value_prefix_len)
 
   let fold_left f init t =
     let result = ref init in
@@ -375,7 +438,7 @@ end
 module Mixed_block_access_field_kind = struct
   type t =
     | Value_prefix of Block_access_field_kind.t
-    | Flat_suffix of Lambda.flat_element
+    | Flat_suffix of Mixed_block_flat_element.t
 
   let [@ocamlformat "disable"] print ppf t =
     match t with
@@ -390,16 +453,20 @@ module Mixed_block_access_field_kind = struct
           "@[<hov 1>(Flat_suffix \
            @[<hov 1>(flat_element@ %a)@]\
            )@]"
-          Printlambda.flat_element flat_element
+          Mixed_block_flat_element.print flat_element
 
   let compare t1 t2 =
     match t1, t2 with
     | Value_prefix field_kind1, Value_prefix field_kind2 ->
       Block_access_field_kind.compare field_kind1 field_kind2
     | Flat_suffix element_kind1, Flat_suffix element_kind2 ->
-      Stdlib.compare element_kind1 element_kind2
+      Mixed_block_flat_element.compare element_kind1 element_kind2
     | Value_prefix _, Flat_suffix _ -> -1
     | Flat_suffix _, Value_prefix _ -> 1
+
+  let to_element_kind = function
+    | Value_prefix _ -> K.value
+    | Flat_suffix flat -> Mixed_block_flat_element.element_kind flat
 end
 
 module Block_access_kind = struct
@@ -449,11 +516,8 @@ module Block_access_kind = struct
     match t with
     | Values _ -> K.value
     | Naked_floats _ -> K.naked_float
-    | Mixed { field_kind; _ } -> (
-      match field_kind with
-      | Value_prefix _ -> K.value
-      | Flat_suffix Imm -> K.value
-      | Flat_suffix (Float | Float64) -> K.naked_float)
+    | Mixed { field_kind; _ } ->
+      Mixed_block_access_field_kind.to_element_kind field_kind
 
   let element_subkind_for_load t =
     match t with
@@ -467,7 +531,10 @@ module Block_access_kind = struct
     | Mixed { field_kind = Flat_suffix field_kind; _ } -> (
       match field_kind with
       | Imm -> K.With_subkind.tagged_immediate
-      | Float | Float64 -> K.With_subkind.naked_float)
+      | Float | Float64 -> K.With_subkind.naked_float
+      | Bits32 -> K.With_subkind.naked_int32
+      | Bits64 -> K.With_subkind.naked_int64
+      | Word -> K.With_subkind.naked_nativeint)
 
   let element_kind_for_set = element_kind_for_load
 
@@ -1801,7 +1868,7 @@ type variadic_primitive =
   | Make_array of Array_kind.t * Mutability.t * Alloc_mode.For_allocations.t
   | Make_mixed_block of
       Tag.Scannable.t
-      * Lambda.mixed_block_shape
+      * Mixed_block_kind.t
       * Mutability.t
       * Alloc_mode.For_allocations.t