Static allocation / simplification for unboxed product arrays

middle_end/flambda2/from_lambda/closure_conversion_aux.ml

@@ -504,9 +504,10 @@ module Acc = struct
       (* For immutable float blocks, we can statically allocate them in classic
          mode, but they are not currently provided with approximations. *)
       | Immutable_float_array _ | Immutable_float32_array _
-      | Immutable_value_array _ | Empty_array _ | Immutable_int32_array _
-      | Immutable_int64_array _ | Immutable_nativeint_array _
-      | Immutable_vec128_array _ | Mutable_string _ | Immutable_string _ ->
+      | Immutable_vec128_array _ | Immutable_value_array _
+      | Immutable_non_scannable_unboxed_product_array _ | Empty_array _
+      | Immutable_int32_array _ | Immutable_int64_array _
+      | Immutable_nativeint_array _ | Mutable_string _ | Immutable_string _ ->
         Value_unknown
     in
     let symbol_approximations =

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -759,6 +759,9 @@ let static_const env (sc : Static_const.t) : Fexpr.static_data =
     Misc.fatal_error
       "fexpr support for unboxed float32/int32/64/nativeint/vec128 arrays not \
        yet implemented"
+  | Immutable_non_scannable_unboxed_product_array _ ->
+    Misc.fatal_error
+      "fexpr support for unboxed product arrays not yet implemented"
   | Empty_array array_kind -> Empty_array array_kind
   | Mutable_string { initial_value } -> Mutable_string { initial_value }
   | Immutable_string s -> Immutable_string s

middle_end/flambda2/reaper/rebuild.ml

@@ -85,6 +85,10 @@ let rewrite_simple_with_debuginfo kinds env (simple : Simple.With_debuginfo.t) =
     (rewrite_simple kinds env (Simple.With_debuginfo.simple simple))
     (Simple.With_debuginfo.dbg simple)
 
+let rewrite_simple_with_debuginfo_and_full_kind kinds env
+    ((simple : Simple.With_debuginfo.t), full_kind) =
+  rewrite_simple_with_debuginfo kinds env simple, full_kind
+
 let rewrite_static_const kinds (env : env) (sc : Static_const.t) =
   match sc with
   | Set_of_closures sc ->
@@ -168,6 +172,11 @@ let rewrite_static_const kinds (env : env) (sc : Static_const.t) =
         fields
     in
     Static_const.immutable_vec128_array fields
+  | Immutable_non_scannable_unboxed_product_array fields ->
+    let fields =
+      List.map (rewrite_simple_with_debuginfo_and_full_kind kinds env) fields
+    in
+    Static_const.immutable_non_scannable_unboxed_product_array fields
   | Empty_array _ | Mutable_string _ | Immutable_string _ -> sc
 
 let rewrite_static_const_or_code kinds env (sc : Static_const_or_code.t) =

middle_end/flambda2/simplify/lifting/reification.ml

@@ -65,6 +65,11 @@ let create_static_const dacc dbg (to_lift : T.to_lift) : RSC.t =
   | Immutable_vec128_array { fields } ->
     let fields = List.map (fun f -> Or_variable.Const f) fields in
     RSC.create_immutable_vec128_array art fields
+  | Immutable_non_scannable_unboxed_product_array { fields } ->
+    let fields' = fields in
+    let fields = convert_fields (List.map fst fields) in
+    let fields = List.combine fields (List.map snd fields') in
+    RSC.create_immutable_non_scannable_unboxed_product_array art fields
   | Immutable_value_array { fields } ->
     let fields = convert_fields fields in
     RSC.create_immutable_value_array art fields

middle_end/flambda2/simplify/rebuilt_static_const.ml

@@ -185,6 +185,15 @@ let create_immutable_nativeint_array =
 let create_immutable_vec128_array =
   create_immutable_naked_number_array SC.immutable_vec128_array
 
+let create_immutable_non_scannable_unboxed_product_array are_rebuilding fields =
+  if ART.do_not_rebuild_terms are_rebuilding
+  then
+    let free_names = free_names_of_fields (List.map fst fields) in
+    Block_not_rebuilt { free_names }
+  else
+    create_normal_non_code
+      (SC.immutable_non_scannable_unboxed_product_array fields)
+
 let create_immutable_value_array are_rebuilding fields =
   if ART.do_not_rebuild_terms are_rebuilding
   then
@@ -227,7 +236,8 @@ let map_set_of_closures t ~f =
       | Immutable_float_block _ | Immutable_float_array _
       | Immutable_float32_array _ | Immutable_int32_array _
       | Immutable_int64_array _ | Immutable_nativeint_array _
-      | Immutable_vec128_array _ | Immutable_value_array _ | Empty_array _
+      | Immutable_vec128_array _ | Immutable_value_array _
+      | Immutable_non_scannable_unboxed_product_array _ | Empty_array _
       | Mutable_string _ | Immutable_string _ ->
         t))
   | Block_not_rebuilt _ | Set_of_closures_not_rebuilt _ | Code_not_rebuilt _ ->

middle_end/flambda2/simplify/rebuilt_static_const.mli

@@ -97,6 +97,11 @@ val create_immutable_vec128_array :
   Vector_types.Vec128.Bit_pattern.t Or_variable.t list ->
   t
 
+val create_immutable_non_scannable_unboxed_product_array :
+  Are_rebuilding_terms.t ->
+  (Simple.With_debuginfo.t * Flambda_kind.With_subkind.t) list ->
+  t
+
 val create_immutable_value_array :
   Are_rebuilding_terms.t -> Simple.With_debuginfo.t list -> t
 

middle_end/flambda2/simplify/simplify_binary_primitive.ml

@@ -927,20 +927,12 @@ let simplify_phys_equal (op : P.equality_comparison) dacc ~original_term _dbg
     SPR.create original_term ~try_reify:false dacc
 
 let simplify_array_load (array_kind : P.Array_kind.t)
-    (array_load_kind : P.Array_load_kind.t) mutability dacc ~original_term:_ dbg
-    ~arg1:array ~arg1_ty:array_ty ~arg2:index ~arg2_ty:index_ty ~result_var =
-  let result_kind =
-    match array_load_kind with
-    | Immediates -> (* CR mshinwell: use the subkind *) K.value
-    | Values -> K.value
-    | Naked_floats -> K.naked_float
-    | Naked_float32s -> K.naked_float32
-    | Naked_int32s -> K.naked_int32
-    | Naked_int64s -> K.naked_int64
-    | Naked_nativeints -> K.naked_nativeint
-    | Naked_vec128s -> K.naked_vec128
-  in
+    (array_load_kind : P.Array_load_kind.t) mutability dacc ~original_prim
+    ~original_term:_ dbg ~arg1:array ~arg1_ty:array_ty ~arg2:index
+    ~arg2_ty:index_ty ~result_var =
+  let result_kind = P.result_kind' original_prim in
   let array_kind =
+    (* XXX mshinwell: should the array_load_kind be specialised too? *)
     Simplify_common.specialise_array_kind dacc array_kind ~array_ty
   in
   (* CR-someday mshinwell: should do a meet on the new value too *)
@@ -959,36 +951,38 @@ let simplify_array_load (array_kind : P.Array_kind.t)
       SPR.create named ~try_reify dacc
     in
     let[@inline] contents_unknown () =
-      return_given_type (T.unknown (P.result_kind' prim)) ~try_reify:false
+      return_given_type (T.unknown result_kind) ~try_reify:false
     in
     (* CR mshinwell/vlaviron: if immutable array accesses were consistently
        setting [mutability] to [Immutable], we could restrict the following code
        to immutable loads only and use [T.meet_is_immutable_array] instead. *)
     match T.prove_is_immutable_array (DA.typing_env dacc) array_ty with
     | Unknown -> contents_unknown ()
-    | Proved (elt_kind, fields, _mode) -> (
-      match elt_kind with
+    | Proved (elt_kinds, fields, _mode) -> (
+      match elt_kinds with
       | Unknown | Bottom -> contents_unknown ()
-      | Ok elt_kind -> (
-        if not (K.equal (K.With_subkind.kind elt_kind) result_kind)
-        then contents_unknown ()
-        else
-          match
-            T.prove_equals_tagged_immediates (DA.typing_env dacc) index_ty
-          with
-          | Unknown -> contents_unknown ()
-          | Proved imms -> (
-            match Targetint_31_63.Set.get_singleton imms with
-            | None -> contents_unknown ()
-            | Some imm ->
-              if Targetint_31_63.( < ) imm Targetint_31_63.zero
-                 || Targetint_31_63.( >= ) imm
-                      (Array.length fields |> Targetint_31_63.of_int)
-              then SPR.create_invalid dacc
-              else
-                return_given_type
-                  fields.(Targetint_31_63.to_int imm)
-                  ~try_reify:true))))
+      | Ok elt_kinds -> (
+        match
+          T.prove_equals_tagged_immediates (DA.typing_env dacc) index_ty
+        with
+        | Unknown -> contents_unknown ()
+        | Proved imms -> (
+          match Targetint_31_63.Set.get_singleton imms with
+          | None -> contents_unknown ()
+          | Some imm ->
+            if Targetint_31_63.( < ) imm Targetint_31_63.zero
+               || Targetint_31_63.( >= ) imm
+                    (Array.length fields |> Targetint_31_63.of_int)
+            then SPR.create_invalid dacc
+            else
+              let index = Targetint_31_63.to_int imm in
+              let num_elt_kinds = List.length elt_kinds in
+              let elt_kind = List.nth elt_kinds (index mod num_elt_kinds) in
+              if not (K.equal (K.With_subkind.kind elt_kind) result_kind)
+              then
+                (* CR mshinwell: I think this should try to reinterpret *)
+                contents_unknown ()
+              else return_given_type fields.(index) ~try_reify:true))))
 
 let simplify_string_or_bigstring_load _string_like_value _string_accessor_width
     ~original_prim dacc ~original_term _dbg ~arg1:_ ~arg1_ty:_ ~arg2:_
@@ -1034,7 +1028,7 @@ let simplify_binary_primitive0 dacc original_prim (prim : P.binary_primitive)
     match prim with
     | Block_set { kind; init; field } -> simplify_block_set kind init ~field
     | Array_load (array_kind, width, mutability) ->
-      simplify_array_load array_kind width mutability
+      simplify_array_load array_kind width mutability ~original_prim
     | Int_arith (kind, op) -> (
       match kind with
       | Tagged_immediate -> Binary_int_arith_tagged_immediate.simplify op

middle_end/flambda2/simplify/simplify_extcall.ml

@@ -128,17 +128,17 @@ let simplify_comparison ~dbg ~dacc ~cont ~tagged_prim ~float_prim
 
 let simplify_caml_make_vect dacc ~len_ty ~init_value_ty : t =
   let typing_env = DA.typing_env dacc in
-  let element_kind : _ Or_unknown_or_bottom.t =
+  let element_kinds : _ Or_unknown_or_bottom.t =
     (* We can't deduce subkind information, e.g. an array is all-immediates
        rather than arbitrary values, but we can deduce kind information. *)
     if not (Flambda_features.flat_float_array ())
-    then Ok (Flambda_kind.With_subkind.anything (T.kind init_value_ty))
+    then Ok [Flambda_kind.With_subkind.anything (T.kind init_value_ty)]
     else
       match T.prove_is_or_is_not_a_boxed_float typing_env init_value_ty with
       | Proved true ->
         (* A boxed float provided to [caml_make_vect] with the float array
            optimisation on will always yield a flat array of naked floats. *)
-        Ok Flambda_kind.With_subkind.naked_float
+        Ok [Flambda_kind.With_subkind.naked_float]
       | Proved false | Unknown -> Unknown
   in
   (* CR-someday mshinwell: We should really adjust the kind of the parameter of
@@ -151,7 +151,7 @@ let simplify_caml_make_vect dacc ~len_ty ~init_value_ty : t =
      Also maybe we should allow static allocation of these arrays for reasonable
      sizes. *)
   let type_of_returned_array =
-    T.mutable_array ~element_kind ~length:len_ty Alloc_mode.For_types.heap
+    T.mutable_array ~element_kinds ~length:len_ty Alloc_mode.For_types.heap
   in
   Unchanged { return_types = Known [type_of_returned_array] }
 

middle_end/flambda2/simplify/simplify_static_const.ml

@@ -53,7 +53,7 @@ let rebuild_naked_number_array dacc ~bind_result_sym kind type_creator creator
   in
   let dacc =
     bind_result_sym
-      (T.immutable_array ~element_kind:(Ok kind) ~fields:field_tys
+      (T.immutable_array ~element_kinds:(Ok [kind]) ~fields:field_tys
          Alloc_mode.For_types.heap)
   in
   creator (DA.are_rebuilding_terms dacc) fields, dacc
@@ -199,6 +199,24 @@ let simplify_static_const_of_kind_value dacc (static_const : Static_const.t)
   | Immutable_vec128_array fields ->
     rebuild_naked_number_array dacc ~bind_result_sym KS.naked_vec128
       T.this_naked_vec128 RSC.create_immutable_vec128_array ~fields
+  | Immutable_non_scannable_unboxed_product_array fields ->
+    let kinds = List.map snd fields in
+    let fields_with_tys =
+      List.map
+        (fun (field, kind) ->
+          simplify_field_of_block dacc (field, K.With_subkind.kind kind))
+        fields
+    in
+    let fields, field_tys = List.split fields_with_tys in
+    let dacc =
+      bind_result_sym
+        (T.immutable_array ~element_kinds:(Ok kinds) ~fields:field_tys
+           Alloc_mode.For_types.heap)
+    in
+    ( Rebuilt_static_const.create_immutable_non_scannable_unboxed_product_array
+        (DA.are_rebuilding_terms dacc)
+        (List.combine fields kinds),
+      dacc )
   | Immutable_value_array fields ->
     let fields_with_tys =
       List.map
@@ -208,7 +226,7 @@ let simplify_static_const_of_kind_value dacc (static_const : Static_const.t)
     let fields, field_tys = List.split fields_with_tys in
     let dacc =
       bind_result_sym
-        (T.immutable_array ~element_kind:(Ok KS.any_value) ~fields:field_tys
+        (T.immutable_array ~element_kinds:(Ok [KS.any_value]) ~fields:field_tys
            Alloc_mode.For_types.heap)
     in
     ( Rebuilt_static_const.create_immutable_value_array
@@ -218,7 +236,7 @@ let simplify_static_const_of_kind_value dacc (static_const : Static_const.t)
   | Empty_array array_kind ->
     let dacc =
       bind_result_sym
-        (T.array_of_length ~element_kind:Bottom
+        (T.array_of_length ~element_kinds:Bottom
            ~length:(T.this_tagged_immediate Targetint_31_63.zero)
            Alloc_mode.For_types.heap)
     in

middle_end/flambda2/simplify/simplify_switch_expr.ml

@@ -293,7 +293,7 @@ let rebuild_switch_with_single_arg_to_same_destination uacc ~dacc_before_switch
         consts
     in
     let block_type =
-      T.immutable_array ~element_kind:(Ok KS.tagged_immediate)
+      T.immutable_array ~element_kinds:(Ok [KS.tagged_immediate])
         ~fields:
           (List.map
              (fun const ->

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -233,7 +233,7 @@ let simplify_array_length _array_kind dacc ~original_term ~arg:_
   Simplify_common.simplify_projection dacc ~original_term
     ~deconstructing:array_ty
     ~shape:
-      (T.array_of_length ~element_kind:Unknown
+      (T.array_of_length ~element_kinds:Unknown
          ~length:(T.alias_type_of K.value result)
          (Alloc_mode.For_types.unknown ()))
     ~result_var ~result_kind:K.value
@@ -660,11 +660,11 @@ let simplify_duplicate_array ~kind:_ ~(source_mutability : Mutability.t)
     | Need_meet ->
       let dacc = DA.add_variable dacc result_var T.any_value in
       SPR.create original_term ~try_reify:false dacc
-    | Known_result (element_kind, fields, alloc_mode) ->
+    | Known_result (element_kinds, fields, alloc_mode) ->
       let length =
         T.this_tagged_immediate (Array.length fields |> Targetint_31_63.of_int)
       in
-      let ty = T.mutable_array ~element_kind ~length alloc_mode in
+      let ty = T.mutable_array ~element_kinds ~length alloc_mode in
       let dacc = DA.add_variable dacc result_var ty in
       SPR.create original_term ~try_reify:false dacc)
   | ( (Immutable | Immutable_unique | Mutable),

middle_end/flambda2/simplify/simplify_variadic_primitive.ml

@@ -94,19 +94,15 @@ let simplify_make_array (array_kind : P.Array_kind.t)
     | None -> T.unknown K.value
   in
   let element_kinds = P.Array_kind.element_kinds array_kind in
-  let element_kind =
-    (* CR mshinwell: support unboxed product arrays in the type system *)
-    (* Remember that the element subkinds cannot in general be deduced from the
-       types of the array members, it must be obtained from the array kind
-       annotations that came via [Lambda]. *)
-    match P.Array_kind.element_kinds array_kind with
-    | [kind] -> Some kind
-    | _ :: _ -> None
-    | [] ->
-      Misc.fatal_errorf
-        "Empty list of element kinds given for array kind:@ %a@ %a"
-        P.Array_kind.print array_kind Debuginfo.print_compact dbg
-  in
+  (* Remember that the element subkinds cannot in general be deduced from the
+     types of the array members, it must be obtained from the array kind
+     annotations that came via [Lambda]. *)
+  (match element_kinds with
+  | _ :: _ -> ()
+  | [] ->
+    Misc.fatal_errorf
+      "Empty list of element kinds given for array kind:@ %a@ %a"
+      P.Array_kind.print array_kind Debuginfo.print_compact dbg);
   let num_element_kinds = List.length element_kinds in
   if List.length args mod num_element_kinds <> 0
   then
@@ -115,34 +111,35 @@ let simplify_make_array (array_kind : P.Array_kind.t)
        list:@ array_kind=%a@ num args=%d@ %a"
       P.Array_kind.print array_kind (List.length args) Named.print original_term;
   let env_extension =
-    match element_kind with
-    | None -> Or_bottom.Ok TEE.empty
-    | Some element_kind ->
-      let initial_element_type = T.unknown_with_subkind element_kind in
-      let typing_env = DA.typing_env dacc in
-      List.fold_left
-        (fun env_extension element_type ->
+    let element_kinds = Array.of_list element_kinds in
+    let typing_env = DA.typing_env dacc in
+    List.fold_left
+      (fun (env_extension, i) element_type ->
+        let initial_element_type =
+          T.unknown_with_subkind element_kinds.(i mod num_element_kinds)
+        in
+        let env_extension =
           let open Or_bottom.Let_syntax in
           let<* env_extension = env_extension in
           let<* _, env_extension' =
             T.meet typing_env initial_element_type element_type
           in
-          TEE.meet typing_env env_extension env_extension')
-        (Or_bottom.Ok TEE.empty) tys
+          TEE.meet typing_env env_extension env_extension'
+        in
+        env_extension, i + 1)
+      (Or_bottom.Ok TEE.empty, 0)
+      tys
+    |> fst
   in
   match env_extension with
   | Bottom -> SPR.create_invalid dacc
   | Ok env_extension ->
     let ty =
       let alloc_mode = Alloc_mode.For_allocations.as_type alloc_mode in
-      let element_kind : _ Or_unknown_or_bottom.t =
-        match element_kind with
-        | None -> (* Array of unboxed products *) Unknown
-        | Some element_kind -> Ok element_kind
-      in
+      let element_kinds : _ Or_unknown_or_bottom.t = Ok element_kinds in
       match mutable_or_immutable with
-      | Mutable -> T.mutable_array ~element_kind ~length alloc_mode
-      | Immutable -> T.immutable_array ~element_kind ~fields:tys alloc_mode
+      | Mutable -> T.mutable_array ~element_kinds ~length alloc_mode
+      | Immutable -> T.immutable_array ~element_kinds ~fields:tys alloc_mode
       | Immutable_unique ->
         Misc.fatal_errorf "Immutable_unique is not expected for arrays:@ %a"
           Named.print original_term