Allow %arrayblit at all layouts

Author: mshinwell

lambda/lambda.ml

@@ -2399,8 +2399,8 @@ let array_set_kind mode = function
   | Pgcscannableproductarray kinds -> Pgcscannableproductarray_set (mode, kinds)
   | Pgcignorableproductarray kinds -> Pgcignorableproductarray_set kinds
 
-let array_ref_kind_of_array_set_kind_for_unboxed_types_and_int
-      (kind : array_set_kind) ~print_array_set_kind : array_ref_kind =
+let array_ref_kind_of_array_set_kind (kind : array_set_kind) mode
+      : array_ref_kind =
   match kind with
   | Pintarray_set -> Pintarray_ref
   | Punboxedfloatarray_set uf -> Punboxedfloatarray_ref uf
@@ -2410,12 +2410,9 @@ let array_ref_kind_of_array_set_kind_for_unboxed_types_and_int
     Pgcscannableproductarray_ref scannables
   | Pgcignorableproductarray_set ignorables ->
     Pgcignorableproductarray_ref ignorables
-  | Pgenarray_set _
-  | Paddrarray_set _
-  | Pfloatarray_set ->
-    Misc.fatal_errorf "Array set kind %a cannot be converted via \
-        array_ref_kind_of_array_set_kind_for_unboxed_types_and_int"
-      print_array_set_kind kind
+  | Pgenarray_set _ -> Pgenarray_ref mode
+  | Paddrarray_set _ -> Paddrarray_ref
+  | Pfloatarray_set -> Pfloatarray_ref mode
 
 let may_allocate_in_region lam =
   (* loop_region raises, if the lambda might allocate in parent region *)

lambda/lambda.mli

@@ -1169,11 +1169,10 @@ val array_ref_kind : locality_mode -> array_kind -> array_ref_kind
 (** The mode will be discarded if unnecessary for the given [array_kind] *)
 val array_set_kind : modify_mode -> array_kind -> array_set_kind
 
-(** Note that this fails on [Pfloatarray_set] *)
-val array_ref_kind_of_array_set_kind_for_unboxed_types_and_int
-  : array_set_kind
-  -> print_array_set_kind:(Format.formatter -> array_set_kind -> unit)
-  -> array_ref_kind
+(** Any mode information in the given [array_set_kind] is ignored.  Any mode
+    in the return value always comes from the [locality_mode] parameter. *)
+val array_ref_kind_of_array_set_kind
+  : array_set_kind -> locality_mode -> array_ref_kind
 
 (* Returns true if the given lambda can allocate on the local stack *)
 val may_allocate_in_region : lambda -> bool

middle_end/flambda2/from_lambda/lambda_to_lambda_transforms.ml

@@ -462,8 +462,11 @@ let makearray_dynamic env (lambda_array_kind : L.array_kind)
 let arrayblit env ~(src_mutability : L.mutable_flag)
     ~(dst_array_set_kind : L.array_set_kind) args loc =
   let src_array_ref_kind =
-    L.array_ref_kind_of_array_set_kind_for_unboxed_types_and_int
-      dst_array_set_kind ~print_array_set_kind:Printlambda.array_set_kind
+    (* We don't expect any allocation (e.g. occurring from the reading of a
+       [float array]) to persist after simplification. We use [alloc_local] just
+       in case that simplification doesn't happen for some reason (this seems
+       unlikely). *)
+    L.array_ref_kind_of_array_set_kind dst_array_set_kind L.alloc_local
   in
   match args with
   | [src_expr; src_start_pos_expr; dst_expr; dst_start_pos_expr; length_expr] ->
@@ -521,8 +524,12 @@ let arrayblit env ~(src_mutability : L.mutable_flag)
     let env, copy_backwards = make_loop env Downto in
     let env, copy_forwards = make_loop env Upto in
     let body =
-      L.Lifthenelse
-        (must_copy_backwards, copy_backwards, copy_forwards, L.layout_unit)
+      (* The region is expected to be redundant (see comment above about
+         modes). *)
+      L.Lregion
+        ( L.Lifthenelse
+            (must_copy_backwards, copy_backwards, copy_forwards, L.layout_unit),
+          L.layout_unit )
     in
     let expr =
       (* Preserve right-to-left evaluation order. *)