flambda-backend: More precise layout for array patterns (#1968)

* more precise layout

* define array_ref_kind_result_layout

Author: alanechang

lambda/lambda.ml

@@ -1530,6 +1530,11 @@ let layout_of_native_repr : Primitive.native_repr -> _ = function
     | Void -> assert false
     end
 
+let array_ref_kind_result_layout = function
+  | Pintarray_ref -> layout_int
+  | Pfloatarray_ref _ -> layout_boxed_float
+  | Pgenarray_ref _ | Paddrarray_ref -> layout_field
+
 let primitive_result_layout (p : primitive) =
   assert !Clflags.native_code;
   match p with
@@ -1570,10 +1575,7 @@ let primitive_result_layout (p : primitive) =
   | Pprobe_is_enabled _ | Pbswap16
     -> layout_int
   | Parrayrefu array_ref_kind | Parrayrefs array_ref_kind ->
-      (match array_ref_kind with
-       | Pintarray_ref -> layout_int
-       | Pfloatarray_ref _ -> layout_boxed_float
-       | Pgenarray_ref _ | Paddrarray_ref -> layout_field)
+    array_ref_kind_result_layout array_ref_kind
   | Pbintofint (bi, _) | Pcvtbint (_,bi,_)
   | Pnegbint (bi, _) | Paddbint (bi, _) | Psubbint (bi, _)
   | Pmulbint (bi, _) | Pdivbint {size = bi} | Pmodbint {size = bi}

lambda/lambda.mli

@@ -764,6 +764,8 @@ val structured_constant_layout : structured_constant -> layout
 
 val primitive_result_layout : primitive -> layout
 
+val array_ref_kind_result_layout: array_ref_kind -> layout
+
 val compute_expr_layout : (Ident.t -> layout option) -> lambda -> layout
 
 (** The mode will be discarded if unnecessary for the given [array_kind] *)

lambda/matching.ml

@@ -2197,18 +2197,19 @@ let get_expr_args_array ~scopes kind head (arg, _mut, _sort, _layout) rem =
     if pos >= len then
       rem
     else
-      (* CR ncourant: could do better than layout_field using kind *)
+      (* TODO: The resulting float should be allocated to at the mode of the
+         array pattern, once that's available *)
+      let ref_kind = Lambda.(array_ref_kind alloc_heap kind) in
+      let result_layout = array_ref_kind_result_layout ref_kind in
       ( Lprim
-          (* TODO: The resulting float should be allocated to at the mode of the
-             array pattern, once that's available *)
-          (Parrayrefu Lambda.(array_ref_kind alloc_heap kind),
+          (Parrayrefu ref_kind,
            [ arg; Lconst (Const_base (Const_int pos)) ],
            loc),
         (match am with
         | Mutable   -> StrictOpt
         | Immutable -> Alias),
         Jkind.Sort.for_array_get_result,
-        layout_field)
+        result_layout)
       :: make_args (pos + 1)
   in
   make_args 0