Allow lifting of immutable arrays (ocaml-flambda#721)

middle_end/flambda2/from_lambda/lambda_to_flambda_primitives.ml

@@ -601,9 +601,9 @@ let convert_lprim ~big_endian (prim : L.primitive) (args : Simple.t list)
         If_then_else
           ( Unary (Is_boxed_float, elt),
             Variadic
-              ( Make_array (Naked_floats, Immutable, mode),
+              ( Make_array (Naked_floats, mutability, mode),
                 List.map unbox_float args ),
-            Variadic (Make_array (Values, Immutable, mode), args) )))
+            Variadic (Make_array (Values, mutability, mode), args) )))
   | Popaque, [arg] -> Unary (Opaque_identity, arg)
   | Pduprecord (repr, num_fields), [arg] ->
     let kind : P.Duplicate_block_kind.t =

middle_end/flambda2/parser/fexpr.ml

@@ -87,6 +87,7 @@ type static_data =
   | Boxed_nativeint of targetint or_variable
   | Immutable_float_block of float or_variable list
   | Immutable_float_array of float or_variable list
+  | Immutable_value_array of field_of_block list
   | Empty_array
   | Mutable_string of { initial_value : string }
   | Immutable_string of string

middle_end/flambda2/parser/fexpr_to_flambda.ml

@@ -662,6 +662,9 @@ let rec expr env (e : Fexpr.expr) : Flambda.Expr.t =
         | Immutable_float_array elements ->
           static_const
             (Immutable_float_array (List.map (or_variable float env) elements))
+        | Immutable_value_array elements ->
+          static_const
+            (Immutable_value_array (List.map (field_of_block env) elements))
         | Empty_array -> static_const Empty_array
         | Mutable_string { initial_value = s } ->
           static_const (Mutable_string { initial_value = s })

middle_end/flambda2/parser/flambda_to_fexpr.ml

@@ -579,6 +579,8 @@ let static_const env (sc : Static_const.t) : Fexpr.static_data =
     Immutable_float_block (List.map (or_variable float env) elements)
   | Immutable_float_array elements ->
     Immutable_float_array (List.map (or_variable float env) elements)
+  | Immutable_value_array elements ->
+    Immutable_value_array (List.map (field_of_block env) elements)
   | Empty_array -> Empty_array
   | Mutable_string { initial_value } -> Mutable_string { initial_value }
   | Immutable_string s -> Immutable_string s

middle_end/flambda2/parser/print_fexpr.ml

@@ -295,6 +295,10 @@ let static_data ppf : static_data -> unit = function
     Format.fprintf ppf "Float_array [|%a|]"
       (pp_semi_list float_or_variable)
       elements
+  | Immutable_value_array elements ->
+    Format.fprintf ppf "Value_array [|%a|]"
+      (pp_semi_list field_of_block)
+      elements
   | Empty_array -> Format.fprintf ppf "Empty_array"
   | Mutable_string { initial_value = s } ->
     Format.fprintf ppf "mutable \"%s\"" (s |> String.escaped)

middle_end/flambda2/simplify/lifting/reification.ml

@@ -17,34 +17,35 @@
 open! Simplify_import
 
 let create_static_const dacc dbg (to_lift : T.to_lift) : RSC.t =
+  let[@inline always] convert_fields fields =
+    ListLabels.map fields ~f:(fun field ->
+        let module F = Field_of_static_block in
+        Simple.pattern_match' field
+          ~var:(fun var ~coercion ->
+            if not (Coercion.is_id coercion)
+            then
+              Misc.fatal_errorf "Expected identity coercion on variable:@ %a"
+                Simple.print field;
+            F.Dynamically_computed (var, dbg))
+          ~symbol:(fun sym ~coercion ->
+            if not (Coercion.is_id coercion)
+            then
+              Misc.fatal_errorf "Expected identity coercion on symbol:@ %a"
+                Simple.print field;
+            F.Symbol sym)
+          ~const:(fun const ->
+            match Reg_width_const.descr const with
+            | Tagged_immediate imm -> F.Tagged_immediate imm
+            | Naked_immediate _ | Naked_float _ | Naked_int32 _ | Naked_int64 _
+            | Naked_nativeint _ ->
+              Misc.fatal_errorf
+                "Expected a constant of kind [Value] but got %a (dbg %a)"
+                Reg_width_const.print const Debuginfo.print_compact dbg))
+  in
   let art = DA.are_rebuilding_terms dacc in
   match to_lift with
   | Immutable_block { tag; is_unique; fields } ->
-    let fields =
-      ListLabels.map fields ~f:(fun field ->
-          let module F = Field_of_static_block in
-          Simple.pattern_match' field
-            ~var:(fun var ~coercion ->
-              if not (Coercion.is_id coercion)
-              then
-                Misc.fatal_errorf "Expected identity coercion on variable:@ %a"
-                  Simple.print field;
-              F.Dynamically_computed (var, dbg))
-            ~symbol:(fun sym ~coercion ->
-              if not (Coercion.is_id coercion)
-              then
-                Misc.fatal_errorf "Expected identity coercion on symbol:@ %a"
-                  Simple.print field;
-              F.Symbol sym)
-            ~const:(fun const ->
-              match Reg_width_const.descr const with
-              | Tagged_immediate imm -> F.Tagged_immediate imm
-              | Naked_immediate _ | Naked_float _ | Naked_int32 _
-              | Naked_int64 _ | Naked_nativeint _ ->
-                Misc.fatal_errorf
-                  "Expected a constant of kind [Value] but got %a (dbg %a)"
-                  Reg_width_const.print const Debuginfo.print_compact dbg))
-    in
+    let fields = convert_fields fields in
     let mut : Mutability.t =
       if is_unique then Immutable_unique else Immutable
     in
@@ -53,6 +54,12 @@ let create_static_const dacc dbg (to_lift : T.to_lift) : RSC.t =
   | Boxed_int32 i -> RSC.create_boxed_int32 art (Const i)
   | Boxed_int64 i -> RSC.create_boxed_int64 art (Const i)
   | Boxed_nativeint i -> RSC.create_boxed_nativeint art (Const i)
+  | Immutable_float_array { fields } ->
+    let fields = List.map (fun f -> Or_variable.Const f) fields in
+    RSC.create_immutable_float_array art fields
+  | Immutable_value_array { fields } ->
+    let fields = convert_fields fields in
+    RSC.create_immutable_value_array art fields
   | Empty_array -> RSC.create_empty_array art
 
 let lift dacc ty ~bound_to static_const : _ Or_invalid.t * DA.t =

middle_end/flambda2/simplify/rebuilt_static_const.ml

@@ -157,6 +157,18 @@ let create_immutable_float_array are_rebuilding fields =
     Block_not_rebuilt { free_names }
   else create_normal_non_code (Immutable_float_array fields)
 
+let create_immutable_value_array are_rebuilding fields =
+  if ART.do_not_rebuild_terms are_rebuilding
+  then
+    let free_names =
+      ListLabels.fold_left fields ~init:Name_occurrences.empty
+        ~f:(fun free_names field ->
+          Name_occurrences.union free_names
+            (Field_of_static_block.free_names field))
+    in
+    Block_not_rebuilt { free_names }
+  else create_normal_non_code (Immutable_value_array fields)
+
 let create_empty_array are_rebuilding =
   if ART.do_not_rebuild_terms are_rebuilding
   then Block_not_rebuilt { free_names = Name_occurrences.empty }
@@ -189,7 +201,8 @@ let map_set_of_closures t ~f =
           }
       | Block _ | Boxed_float _ | Boxed_int32 _ | Boxed_int64 _
       | Boxed_nativeint _ | Immutable_float_block _ | Immutable_float_array _
-      | Empty_array | Mutable_string _ | Immutable_string _ ->
+      | Immutable_value_array _ | Empty_array | Mutable_string _
+      | Immutable_string _ ->
         t))
   | Block_not_rebuilt _ | Set_of_closures_not_rebuilt _ | Code_not_rebuilt _ ->
     t

middle_end/flambda2/simplify/rebuilt_static_const.mli

@@ -88,6 +88,9 @@ val create_immutable_float_array :
   Numeric_types.Float_by_bit_pattern.t Or_variable.t list ->
   t
 
+val create_immutable_value_array :
+  Are_rebuilding_terms.t -> Field_of_static_block.t list -> t
+
 val create_empty_array : Are_rebuilding_terms.t -> t
 
 val create_mutable_string : Are_rebuilding_terms.t -> initial_value:string -> t

middle_end/flambda2/simplify/simplify_binary_primitive.ml

@@ -1006,6 +1006,9 @@ let simplify_array_load (array_kind : P.Array_kind.t) mutability dacc
     let dacc = DA.add_variable dacc result_var ty in
     SPR.create_invalid dacc
   | Ok array_kind ->
+    (* CR mshinwell: Add proper support for immutable arrays here (probably not
+       required at present since they only go into [Duplicate_array]
+       operations). *)
     let result_kind' =
       P.Array_kind.element_kind array_kind |> K.With_subkind.kind
     in

middle_end/flambda2/simplify/simplify_extcall.ml

@@ -147,7 +147,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.array_of_length ~element_kind ~length:len_ty
+      T.mutable_array ~element_kind ~length:len_ty (Known Heap)
     in
     Unchanged { return_types = Known [type_of_returned_array] }
 

middle_end/flambda2/simplify/simplify_static_const.ml

@@ -162,12 +162,30 @@ let simplify_static_const_of_kind_value dacc (static_const : Static_const.t)
             field K.naked_float)
         fields
     in
-    let fields, _field_tys = List.split fields_with_tys in
-    let dacc = bind_result_sym T.any_value in
+    let fields, field_tys = List.split fields_with_tys in
+    let dacc =
+      bind_result_sym
+        (T.immutable_array ~element_kind:(Known K.With_subkind.naked_float)
+           ~fields:field_tys (Known Heap))
+    in
     ( Rebuilt_static_const.create_immutable_float_array
         (DA.are_rebuilding_terms dacc)
         fields,
       dacc )
+  | Immutable_value_array fields ->
+    let fields_with_tys =
+      List.map (fun field -> simplify_field_of_block dacc field) fields
+    in
+    let fields, field_tys = List.split fields_with_tys in
+    let dacc =
+      bind_result_sym
+        (T.immutable_array ~element_kind:(Known K.With_subkind.any_value)
+           ~fields:field_tys (Known Heap))
+    in
+    ( Rebuilt_static_const.create_immutable_value_array
+        (DA.are_rebuilding_terms dacc)
+        fields,
+      dacc )
   | Empty_array ->
     (* CR-someday lmaurer: Comment from lthls:
 
@@ -181,7 +199,8 @@ let simplify_static_const_of_kind_value dacc (static_const : Static_const.t)
     let dacc =
       bind_result_sym
         (T.array_of_length ~element_kind:Unknown
-           ~length:(T.this_tagged_immediate Targetint_31_63.zero))
+           ~length:(T.this_tagged_immediate Targetint_31_63.zero)
+           (Known Heap))
     in
     Rebuilt_static_const.create_empty_array (DA.are_rebuilding_terms dacc), dacc
   | Mutable_string { initial_value } ->

middle_end/flambda2/simplify/simplify_unary_primitive.ml

@@ -197,7 +197,8 @@ let simplify_array_length dacc ~original_term ~arg:_ ~arg_ty:array_ty
     ~deconstructing:array_ty
     ~shape:
       (T.array_of_length ~element_kind:Unknown
-         ~length:(T.alias_type_of K.value result))
+         ~length:(T.alias_type_of K.value result)
+         Unknown)
     ~result_var ~result_kind:K.value
 
 (* CR-someday mshinwell: Consider whether "string length" should be treated like
@@ -477,20 +478,31 @@ let simplify_bigarray_length ~dimension:_ dacc ~original_term ~arg:_ ~arg_ty:_
   let dacc = DA.add_variable dacc result_var ty in
   SPR.create original_term ~try_reify:false dacc
 
-let simplify_duplicate_array ~kind:_ ~source_mutability:_
-    ~destination_mutability:_ dacc ~original_term ~arg:_ ~arg_ty ~result_var =
-  (* This simplification should eliminate bounds checks on array literals.
-
-     Any alias in the type to the whole array will be dropped, but aliases
-     inside the type (in this case for the length) can remain. Similarly for
-     blocks in [simplify_duplicate_block] below, aliases to the fields can
-     remain. *)
-  let ty = T.remove_outermost_alias (DA.typing_env dacc) arg_ty in
-  let dacc = DA.add_variable dacc result_var ty in
-  SPR.create original_term ~try_reify:false dacc
+let simplify_duplicate_array ~kind:_ ~(source_mutability : Mutability.t)
+    ~(destination_mutability : Mutability.t) dacc ~original_term ~arg:_ ~arg_ty
+    ~result_var =
+  (* This simplification should eliminate bounds checks on array literals. *)
+  match source_mutability, destination_mutability with
+  | Immutable, Mutable -> (
+    match T.meet_is_immutable_array (DA.typing_env dacc) arg_ty with
+    | Invalid -> SPR.create_invalid dacc
+    | 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, length, alloc_mode) ->
+      let ty = T.mutable_array ~element_kind ~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),
+      (Immutable | Immutable_unique | Mutable) ) ->
+    Misc.fatal_errorf
+      "Combination of mutabilities not supported for [Duplicate_array]:@ %a"
+      Named.print original_term
 
 let simplify_duplicate_block ~kind:_ dacc ~original_term ~arg:_ ~arg_ty
     ~result_var =
+  (* Any alias in the type to the whole block will be dropped, but aliases
+     inside the type (e.g. in fields) can remain. *)
   let ty = T.remove_outermost_alias (DA.typing_env dacc) arg_ty in
   let dacc = DA.add_variable dacc result_var ty in
   SPR.create original_term ~try_reify:false dacc

middle_end/flambda2/simplify/simplify_variadic_primitive.ml

@@ -91,8 +91,9 @@ let simplify_make_block_of_floats ~original_prim ~mutable_or_immutable
     Tag.double_array_tag ~shape ~mutable_or_immutable alloc_mode dacc
     ~original_term dbg ~args_with_tys ~result_var
 
-let simplify_make_array (array_kind : P.Array_kind.t) ~mutable_or_immutable
-    alloc_mode dacc ~original_term:_ dbg ~args_with_tys ~result_var =
+let simplify_make_array (array_kind : P.Array_kind.t)
+    ~(mutable_or_immutable : Mutability.t) alloc_mode dacc ~original_term dbg
+    ~args_with_tys ~result_var =
   let args, tys = List.split args_with_tys in
   let length =
     match Targetint_31_63.of_int_option (List.length args) with
@@ -121,7 +122,18 @@ let simplify_make_array (array_kind : P.Array_kind.t) ~mutable_or_immutable
   match env_extension with
   | Bottom -> SPR.create_invalid dacc
   | Ok env_extension ->
-    let ty = T.array_of_length ~element_kind:(Known element_kind) ~length in
+    let ty =
+      match mutable_or_immutable with
+      | Mutable ->
+        T.mutable_array ~element_kind:(Known element_kind) ~length
+          (Known alloc_mode)
+      | Immutable ->
+        T.immutable_array ~element_kind:(Known element_kind) ~fields:tys
+          (Known alloc_mode)
+      | Immutable_unique ->
+        Misc.fatal_errorf "Immutable_unique is not expected for arrays:@ %a"
+          Named.print original_term
+    in
     let named =
       Named.create_prim
         (Variadic

middle_end/flambda2/terms/flambda.ml

@@ -1414,7 +1414,8 @@ module Named = struct
                  ( Block _ | Boxed_float _ | Boxed_int32 _ | Boxed_int64 _
                  | Boxed_nativeint _ | Immutable_float_block _
                  | Immutable_float_array _ | Mutable_string _
-                 | Immutable_string _ | Empty_array ) ->
+                 | Immutable_string _ | Empty_array | Immutable_value_array _ )
+               ->
                acc)
            init
 end