Support for mixed blocks in Flambda2 types (#2533)

Author: lthls

middle_end/flambda2/classic_mode_types/value_approximation.ml

@@ -28,7 +28,8 @@ type 'code t =
         code : 'code;
         symbol : Symbol.t option
       }
-  | Block_approximation of Tag.t * 'code t array * Alloc_mode.For_types.t
+  | Block_approximation of
+      Tag.Scannable.t * 'code t array * Alloc_mode.For_types.t
 
 let rec print fmt = function
   | Value_unknown -> Format.fprintf fmt "?"
@@ -41,7 +42,8 @@ let rec print fmt = function
     if len < 1
     then Format.fprintf fmt "{}"
     else (
-      Format.fprintf fmt "@[<hov 2>{%a:%a" Tag.print tag print fields.(0);
+      Format.fprintf fmt "@[<hov 2>{%a:%a" Tag.Scannable.print tag print
+        fields.(0);
       for i = 1 to len - 1 do
         Format.fprintf fmt "@ %a" print fields.(i)
       done;

middle_end/flambda2/classic_mode_types/value_approximation.mli

@@ -28,7 +28,8 @@ type 'code t =
         code : 'code;
         symbol : Symbol.t option
       }
-  | Block_approximation of Tag.t * 'code t array * Alloc_mode.For_types.t
+  | Block_approximation of
+      Tag.Scannable.t * 'code t array * Alloc_mode.For_types.t
 
 val print : Format.formatter -> 'a t -> unit
 

middle_end/flambda2/from_lambda/closure_conversion.ml

@@ -1053,7 +1053,6 @@ let close_let acc env let_bound_ids_with_kinds user_visible defining_expr
             ( Variadic
                 (Make_block (Values (tag, _), Immutable, alloc_mode), fields),
               _ ) -> (
-          let tag' = Tag.Scannable.to_tag tag in
           let approxs =
             List.map (find_value_approximation body_env) fields |> Array.of_list
           in
@@ -1071,7 +1070,7 @@ let close_let acc env let_bound_ids_with_kinds user_visible defining_expr
             let acc =
               Acc.add_symbol_approximation acc sym
                 (Value_approximation.Block_approximation
-                   (tag', approxs, Alloc_mode.For_allocations.as_type alloc_mode))
+                   (tag, approxs, Alloc_mode.For_allocations.as_type alloc_mode))
             in
             body acc body_env
           | Computed_static static_fields ->
@@ -1095,7 +1094,7 @@ let close_let acc env let_bound_ids_with_kinds user_visible defining_expr
             in
             let approx =
               Value_approximation.Block_approximation
-                (tag', approxs, Alloc_mode.For_allocations.as_type alloc_mode)
+                (tag, approxs, Alloc_mode.For_allocations.as_type alloc_mode)
             in
             let acc = Acc.add_symbol_approximation acc symbol approx in
             let acc, body = body acc body_env in
@@ -1105,7 +1104,7 @@ let close_let acc env let_bound_ids_with_kinds user_visible defining_expr
               defining_expr ~body
           | Dynamic_block ->
             let body_env =
-              Env.add_block_approximation body_env var tag' approxs
+              Env.add_block_approximation body_env var tag approxs
                 (Alloc_mode.For_allocations.as_type alloc_mode)
             in
             bind acc body_env)

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -494,7 +494,6 @@ module Acc = struct
             | Tagged_immediate i -> Value_int i
             | Dynamically_computed _ -> Value_unknown
           in
-          let tag = Tag.Scannable.to_tag tag in
           let fields = List.map approx_of_field fields |> Array.of_list in
           Block_approximation (tag, fields, Alloc_mode.For_types.unknown ())
         else Value_unknown

middle_end/flambda2/from_lambda/closure_conversion_aux.mli

@@ -169,7 +169,7 @@ module Env : sig
   val add_block_approximation :
     t ->
     Variable.t ->
-    Tag.t ->
+    Tag.Scannable.t ->
     value_approximation array ->
     Alloc_mode.For_types.t ->
     t

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -999,8 +999,8 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     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)]
+    let shape = K.Mixed_block_shape.from_lambda shape in
+    [Variadic (Make_block (Mixed (tag, shape), mutability, mode), args)]
   | Pmakearray (lambda_array_kind, mutability, mode), _ -> (
     let args = List.flatten args in
     let mode = Alloc_mode.For_allocations.from_lambda mode ~current_region in
@@ -1447,7 +1447,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
       Naked_floats { size = Unknown }
     in
     [Binary (Block_load (block_access, mutability), arg, Simple field)]
-  | Pmixedfield (field, read, sem), [[arg]] -> (
+  | Pmixedfield (field, read, shape, sem), [[arg]] -> (
     let imm = Targetint_31_63.of_int field in
     check_non_negative_imm imm "Pmixedfield";
     let field = Simple.const (Reg_width_const.tagged_immediate imm) in
@@ -1460,11 +1460,11 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
         | Mread_flat_suffix read ->
           Flat_suffix
             (match read with
-            | Flat_read flat_element ->
-              P.Mixed_block_flat_element.from_lambda flat_element
-            | Flat_read_float_boxed _ -> Float_boxed)
+            | Flat_read flat_element -> K.from_lambda_flat_element flat_element
+            | Flat_read_float_boxed _ -> K.naked_float)
       in
-      Mixed { tag = Unknown; field_kind; size = Unknown }
+      let shape = K.Mixed_block_shape.from_lambda shape in
+      Mixed { tag = Unknown; field_kind; shape; size = Unknown }
     in
     let block_access : H.expr_primitive =
       Binary (Block_load (block_access, mutability), arg, Simple field)
@@ -1510,7 +1510,7 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
     [ Ternary
         (Block_set (block_access, init_or_assign), block, Simple field, value)
     ]
-  | ( Psetmixedfield (field, write, initialization_or_assignment),
+  | ( Psetmixedfield (field, write, shape, initialization_or_assignment),
       [[block]; [value]] ) ->
     let imm = Targetint_31_63.of_int field in
     check_non_negative_imm imm "Psetmixedfield";
@@ -1523,9 +1523,10 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list list)
               Value_prefix
                 (convert_block_access_field_kind immediate_or_pointer)
             | Mwrite_flat_suffix flat ->
-              Flat_suffix (P.Mixed_block_flat_element.from_lambda flat));
-          size = Unknown;
-          tag = Unknown
+              Flat_suffix (K.from_lambda_flat_element flat));
+          shape = K.Mixed_block_shape.from_lambda shape;
+          tag = Unknown;
+          size = Unknown
         }
     in
     let init_or_assign = convert_init_or_assign initialization_or_assignment in

middle_end/flambda2/identifiers/reg_width_const.ml

@@ -40,3 +40,20 @@ let is_tagged_immediate t =
   | Naked_immediate _ | Naked_float _ | Naked_float32 _ | Naked_int32 _
   | Naked_int64 _ | Naked_nativeint _ | Naked_vec128 _ ->
     None
+
+let of_int_of_kind (kind : Flambda_kind.t) i =
+  match kind with
+  | Value -> tagged_immediate (Targetint_31_63.of_int i)
+  | Naked_number Naked_float ->
+    naked_float (Numeric_types.Float_by_bit_pattern.create (float_of_int i))
+  | Naked_number Naked_float32 ->
+    naked_float32 (Numeric_types.Float32_by_bit_pattern.create (float_of_int i))
+  | Naked_number Naked_immediate -> naked_immediate (Targetint_31_63.of_int i)
+  | Naked_number Naked_int32 -> naked_int32 (Int32.of_int i)
+  | Naked_number Naked_int64 -> naked_int64 (Int64.of_int i)
+  | Naked_number Naked_nativeint -> naked_nativeint (Targetint_32_64.of_int i)
+  | Naked_number Naked_vec128 ->
+    let i = Int64.of_int i in
+    naked_vec128 (Vector_types.Vec128.Bit_pattern.of_bits { high = i; low = i })
+  | Region | Rec_info ->
+    Misc.fatal_errorf "Invalid kind %a" Flambda_kind.print kind

middle_end/flambda2/identifiers/reg_width_const.mli

@@ -25,3 +25,7 @@ val of_descr : Descr.t -> t
 val is_naked_immediate : t -> Targetint_31_63.t option
 
 val is_tagged_immediate : t -> Targetint_31_63.t option
+
+(** Create a numeric constant of the given kind ([Region] and [Rec_info] are
+    forbidden). *)
+val of_int_of_kind : Flambda_kind.t -> int -> t