flambda-backend: Support unboxed ints in mixed blocks (#2515)

* start work

* Implement support

* Broken attempt to type-check these

* Fix bugs

* Gets off the ground

* Generate new test cases

* A bad attempt to make the testsuite better

* Improve generated test

* Self-review

* Clean up unnecessary field

* Get some layouts tests working

* Add some mixed block tests and fix bug

* Move unboxed-float relevant mixed block tests back to typing-layouts-float64 directory

* Add tests for other kinds of mixed blocks

* Flesh out constructor arg tests

* Fix upstream build

* Mint separate mixed block shape type for flambda2 (#2530)

* Ban void in mixed products (#2531)

* Flesh out constructor arg tests to have more combos of flat suffixes

* Review: fix comments; add case for disallowed value in flat suffix

* Fix up both alpha and beta versions of tests according to review

* Review: add and clarify comments

* Remove stale comment (Max confirmed)

* Revert unintended change to test from review response

* Disallow big-endian for now

* Update comment to clarify why it would be better to use mutability (optimization, not soundness)

* make fmt

Author: ncik-roberts

lambda/lambda.ml

@@ -342,11 +342,11 @@ and layout =
 and block_shape =
   value_kind list option
 
-and flat_element = Types.flat_element = Imm | Float | Float64
+and flat_element = Types.flat_element =
+    Imm | Float | Float64 | Bits32 | Bits64 | Word
 and flat_element_read =
-  | Flat_read_imm
+  | Flat_read of flat_element (* invariant: not [Float] *)
   | Flat_read_float of alloc_mode
-  | Flat_read_float64
 and mixed_block_read =
   | Mread_value_prefix of immediate_or_pointer
   | Mread_flat_suffix of flat_element_read
@@ -472,13 +472,6 @@ let join_boxed_vector_layout v1 v2 =
   match v1, v2 with
   | Pvec128 v1, Pvec128 v2 -> Punboxed_vector (Pvec128 (join_vec128_types v1 v2))
 
-let equal_flat_element e1 e2 =
-  match e1, e2 with
-  | Imm, Imm -> true
-  | Float, Float -> true
-  | Float64, Float64 -> true
-  | (Imm | Float | Float64), _ -> false
-
 let rec equal_value_kind x y =
   match x, y with
   | Pgenval, Pgenval -> true
@@ -513,7 +506,7 @@ and equal_constructor_shape x y =
       List.length p1 = List.length p2
       && List.for_all2 equal_value_kind p1 p2
       && List.length s1 = List.length s2
-      && List.for_all2 equal_flat_element s1 s2
+      && List.for_all2 Types.equal_flat_element s1 s2
   | (Constructor_uniform _ | Constructor_mixed _), _ -> false
 
 let equal_layout x y =
@@ -1246,15 +1239,20 @@ let transl_prim mod_name name =
 let transl_mixed_product_shape : Types.mixed_product_shape -> mixed_block_shape =
   fun x -> x
 
-let count_mixed_block_values_and_floats =
-  Types.count_mixed_record_values_and_floats
-
 type mixed_block_element = Types.mixed_product_element =
   | Value_prefix
   | Flat_suffix of flat_element
 
 let get_mixed_block_element = Types.get_mixed_product_element
 
+let flat_read_non_float flat_element =
+  match flat_element with
+  | Float -> Misc.fatal_error "flat_element_read_non_float Float"
+  | Imm | Float64 | Bits32 | Bits64 | Word as flat_element ->
+      Flat_read flat_element
+
+let flat_read_float alloc_mode = Flat_read_float alloc_mode
+
 (* Compile a sequence of expressions *)
 
 let rec make_sequence fn = function
@@ -1635,7 +1633,7 @@ let primitive_may_allocate : primitive -> alloc_mode option = function
       match read with
       | Mread_value_prefix _ -> None
       | Mread_flat_suffix (Flat_read_float m) -> Some m
-      | Mread_flat_suffix (Flat_read_float64 | Flat_read_imm) -> None
+      | Mread_flat_suffix (Flat_read _) -> None
     end
   | Psetfloatfield _ -> None
   | Psetufloatfield _ -> None
@@ -1779,6 +1777,20 @@ let array_ref_kind_result_layout = function
   | Punboxedintarray_ref Pint64 -> layout_unboxed_int64
   | Punboxedintarray_ref Pnativeint -> layout_unboxed_nativeint
 
+let layout_of_mixed_field (kind : mixed_block_read) =
+  match kind with
+  | Mread_value_prefix _ -> layout_value_field
+  | Mread_flat_suffix (Flat_read_float (_ : alloc_mode)) ->
+      layout_boxed_float Pfloat64
+  | Mread_flat_suffix (Flat_read proj) ->
+      match proj with
+      | Imm -> layout_int
+      | Float64 -> layout_unboxed_float Pfloat64
+      | Bits32 -> layout_unboxed_int32
+      | Bits64 -> layout_unboxed_int64
+      | Word -> layout_unboxed_nativeint
+      | Float -> layout_boxed_float Pfloat64
+
 let primitive_result_layout (p : primitive) =
   assert !Clflags.native_code;
   match p with
@@ -1808,16 +1820,7 @@ let primitive_result_layout (p : primitive) =
   | Pbox_float (f, _) -> layout_boxed_float f
   | Pufloatfield _ -> Punboxed_float Pfloat64
   | Punbox_float float_kind -> Punboxed_float float_kind
-  | Pmixedfield (_, kind, _) -> begin
-      match kind with
-      | Mread_value_prefix _ -> layout_value_field
-      | Mread_flat_suffix proj -> begin
-          match proj with
-          | Flat_read_imm -> layout_int
-          | Flat_read_float _ -> layout_boxed_float Pfloat64
-          | Flat_read_float64 -> layout_unboxed_float Pfloat64
-        end
-    end
+  | Pmixedfield (_, kind, _) -> layout_of_mixed_field kind
   | Pccall { prim_native_repr_res = _, repr_res } -> layout_of_extern_repr repr_res
   | Praise _ -> layout_bottom
   | Psequor | Psequand | Pnot

lambda/lambda.mli

@@ -350,11 +350,11 @@ and layout =
 and block_shape =
   value_kind list option
 
-and flat_element = Types.flat_element = Imm | Float | Float64
-and flat_element_read =
-  | Flat_read_imm
+and flat_element = Types.flat_element =
+    Imm | Float | Float64 | Bits32 | Bits64 | Word
+and flat_element_read = private
+  | Flat_read of flat_element (* invariant: not [Float] *)
   | Flat_read_float of alloc_mode
-  | Flat_read_float64
 and mixed_block_read =
   | Mread_value_prefix of immediate_or_pointer
   | Mread_flat_suffix of flat_element_read
@@ -790,7 +790,6 @@ val transl_extension_path: scoped_location -> Env.t -> Path.t -> lambda
 val transl_class_path: scoped_location -> Env.t -> Path.t -> lambda
 
 val transl_mixed_product_shape: Types.mixed_product_shape -> mixed_block_shape
-val count_mixed_block_values_and_floats : mixed_block_shape -> int * int
 
 type mixed_block_element =
   | Value_prefix
@@ -799,6 +798,10 @@ type mixed_block_element =
 (** Raises if the int is out of bounds. *)
 val get_mixed_block_element : mixed_block_shape -> int -> mixed_block_element
 
+(** Raises if [flat_element] is float. *)
+val flat_read_non_float : flat_element -> flat_element_read
+val flat_read_float : alloc_mode -> flat_element_read
+
 val make_sequence: ('a -> lambda) -> 'a list -> lambda
 
 val subst:

lambda/matching.ml

@@ -100,7 +100,6 @@ module Scoped_location = Debuginfo.Scoped_location
 
 type error =
   | Void_layout
-  | Illegal_record_field of Jkind.const
 
 exception Error of Location.t * error
 
@@ -112,20 +111,6 @@ let jkind_layout_default_to_value_and_check_not_void loc jkind =
   | _ -> ()
 ;;
 
-(* CR layouts v5: This function is only used for sanity checking the
-   typechecker.  When we allow arbitrary layouts in structures, it will have
-   outlived its usefulness and should be deleted. *)
-let check_record_field_jkind lbl =
-  match Jkind.(get_default_value lbl.lbl_jkind), lbl.lbl_repres with
-  | (Value | Immediate | Immediate64 | Non_null_value), _ -> ()
-  | Float64, (Record_ufloat | Record_mixed _) -> ()
-  | Float64, (Record_boxed _ | Record_inlined _
-             | Record_unboxed | Record_float) ->
-    raise (Error (lbl.lbl_loc, Illegal_record_field Float64))
-  | (Any | Void | Word | Bits32 | Bits64) as c, _ ->
-    (* CR layouts v2.1: support unboxed ints here *)
-    raise (Error (lbl.lbl_loc, Illegal_record_field c))
-
 (*
    Compatibility predicate that considers potential rebindings of constructors
    of an extension type.
@@ -1824,12 +1809,11 @@ let get_expr_args_constr ~scopes head (arg, _mut, sort, layout) rem =
             | Flat_suffix flat ->
                 let flat_read =
                   match flat with
-                  | Imm -> Flat_read_imm
-                  | Float64 -> Flat_read_float64
                   | Float ->
                       Misc.fatal_error
                         "unexpected flat float of layout value in \
-                          constructor field"
+                         constructor field"
+                  | non_float -> flat_read_non_float non_float
                 in
                 Mread_flat_suffix flat_read
           in
@@ -2136,7 +2120,7 @@ let record_matching_line num_fields lbl_pat_list =
   List.iter (fun (_, lbl, pat) ->
     (* CR layouts v5: This void sanity check can be removed when we add proper
        void support (or whenever we remove `lbl_pos_void`) *)
-    check_record_field_jkind lbl;
+    jkind_layout_default_to_value_and_check_not_void pat.pat_loc lbl.lbl_jkind;
     patv.(lbl.lbl_pos) <- pat)
     lbl_pat_list;
   Array.to_list patv
@@ -2163,7 +2147,8 @@ let get_expr_args_record ~scopes head (arg, _mut, sort, layout) rem =
       rem
     else
       let lbl = all_labels.(pos) in
-      check_record_field_jkind lbl;
+      jkind_layout_default_to_value_and_check_not_void
+        head.pat_loc lbl.lbl_jkind;
       let ptr = Typeopt.maybe_pointer_type head.pat_env lbl.lbl_arg in
       let lbl_sort = Jkind.sort_of_jkind lbl.lbl_jkind in
       let lbl_layout = Typeopt.layout_of_sort lbl.lbl_loc lbl_sort in
@@ -2196,11 +2181,11 @@ let get_expr_args_record ~scopes head (arg, _mut, sort, layout) rem =
               else
                 let read =
                   match flat_suffix.(pos - value_prefix_len) with
-                  | Imm -> Flat_read_imm
-                  | Float64 -> Flat_read_float64
+                  | Imm | Float64 | Bits32 | Bits64 | Word as non_float ->
+                      flat_read_non_float non_float
                   | Float ->
                       (* TODO: could optimise to Alloc_local sometimes *)
-                      Flat_read_float alloc_heap
+                      flat_read_float alloc_heap
                 in
                 Mread_flat_suffix read
             in
@@ -4187,11 +4172,6 @@ let report_error ppf = function
       fprintf ppf
         "Void layout detected in translation:@ Please report this error to \
          the Jane Street compilers team."
-  | Illegal_record_field c ->
-      fprintf ppf
-        "Sort %s detected where value was expected in a record field:@ Please \
-         report this error to the Jane Street compilers team."
-        (Jkind.string_of_const c)
 
 let () =
   Location.register_error_of_exn

lambda/printlambda.ml

@@ -141,10 +141,9 @@ let constructor_shape print_value_kind ppf shape =
        | _ :: _ ->
            fprintf ppf ";@%a"
              (Format.pp_print_list ~pp_sep:(fun ppf () -> fprintf ppf ",@")
-              (fun ppf -> function
-                  | Imm -> fprintf ppf "[imm]"
-                  | Float -> fprintf ppf "[float]"
-                  | Float64 -> fprintf ppf "[float64]"))
+              (fun ppf flat_element ->
+                fprintf ppf "[%s]"
+                    (Types.flat_element_to_lowercase_string flat_element)))
              flat_fields)
 
 let tag_and_constructor_shape print_value_kind ppf (tag, shape) =
@@ -317,15 +316,13 @@ let block_shape ppf shape = match shape with
         t;
       Format.fprintf ppf ")"
 
-let flat_element ppf : flat_element -> unit = function
-  | Imm -> pp_print_string ppf "int"
-  | Float -> pp_print_string ppf "float"
-  | Float64 -> pp_print_string ppf "float64"
+let flat_element ppf : flat_element -> unit = fun x ->
+  pp_print_string ppf (Types.flat_element_to_lowercase_string x)
 
 let flat_element_read ppf : flat_element_read -> unit = function
-  | Flat_read_imm -> pp_print_string ppf "int"
+  | Flat_read flat ->
+      pp_print_string ppf (Types.flat_element_to_lowercase_string flat)
   | Flat_read_float m -> fprintf ppf "float[%a]" alloc_mode m
-  | Flat_read_float64 -> pp_print_string ppf "float64"
 
 let mixed_block_read ppf : mixed_block_read -> unit = function
   | Mread_value_prefix Immediate -> pp_print_string ppf "value_int"