Fix bigarray kinds + prim translation

Add boxing/unboxing when translating bigarray loads/sets since
the rpimitives take unboxed/untagged arguments, and returns
untagged/unboxed results.
This uncovered another bug, which was that the kinds for bigarray
elements in flambda_primitives.ml were wrong.

Author: Gbury

middle_end/flambda2.0/from_lambda/lambda_to_flambda_primitives.ml

@@ -52,38 +52,127 @@ let bint_shift bi prim arg1 arg2 =
     (Binary (Int_shift (C.standard_int_of_boxed_integer bi, prim),
              unbox_bint bi arg1, untag_int arg2))
 
-let string_or_bytes_access_validity_condition str kind index : H.expr_primitive =
+(* offset to substract to a string length depending
+   on the size of the read/write *)
+let length_offset_of_size size =
+  let offset =
+    match (size : Flambda_primitive.string_accessor_width) with
+    | Eight -> 0
+    | Sixteen -> 1
+    | Thirty_two -> 3
+    | Sixty_four -> 7
+  in
+  Immediate.int (Targetint.OCaml.of_int offset)
+
+(* equivalent to the one in `cmm_helpers.ml` *)
+let max_or_zero size_int x =
+  let register_bitsize_minus_one =
+    H.Simple (Simple.const (Naked_immediate (
+        Immediate.int (Targetint.OCaml.of_int (size_int * 8 - 1)))))
+  in
+  let sign =
+    H.Prim (Binary (Int_shift (Naked_nativeint, Asr),
+                  x, register_bitsize_minus_one)) in
+  let minus_one =
+    H.Simple (Simple.const (Naked_nativeint (Targetint.of_int (-1))))
+  in
+  let sign_negation =
+    H.Prim (Binary (Int_arith (Naked_nativeint, Xor),
+                  sign, minus_one)) in
+  let ret =
+    H.Prim (Binary (Int_arith (Naked_nativeint, And),
+                  sign_negation, x)) in
+  ret
+
+(* actual (strict) upper bound for an index in a string read/write *)
+let actual_max_length size_int size length =
+  if size = (Eight : Flambda_primitive.string_accessor_width) then
+    length (* micro-optimization *)
+  else begin
+    let offset = length_offset_of_size size in
+    let reduced_length =
+      H.Prim (Binary (Int_arith (Naked_immediate, Sub),
+                    length, Simple (Simple.const (Naked_immediate offset))))
+    in
+    let reduced_length_nativeint =
+      H.Prim (Unary (Num_conv { src = Naked_immediate; dst = Naked_nativeint },
+                   reduced_length))
+    in
+    let nativeint_res = max_or_zero size_int reduced_length_nativeint in
+    H.Prim (Unary (Num_conv { src = Naked_nativeint; dst = Naked_immediate },
+                   nativeint_res))
+  end
+
+let string_or_bytes_access_validity_condition
+    size_int str kind size index : H.expr_primitive =
   Binary (Int_comp (I.Naked_immediate, Unsigned, Lt),
           untag_int index,
-          (Prim (Unary (String_length kind, str))))
+          actual_max_length size_int size (Prim (Unary (String_length kind, str))))
 
-let string_or_bytes_ref kind arg1 arg2 dbg : H.expr_primitive =
+let string_or_bytes_ref size_int kind arg1 arg2 dbg : H.expr_primitive =
   Checked {
     primitive =
       tag_int (Binary (String_or_bigstring_load (kind, Eight), arg1, arg2));
     validity_conditions = [
-      string_or_bytes_access_validity_condition arg1 String arg2;
+      string_or_bytes_access_validity_condition size_int arg1 String Eight arg2;
     ];
     failure = Index_out_of_bounds;
     dbg;
   }
 
-let bigstring_access_validity_condition bstr index : H.expr_primitive =
+let bigstring_access_validity_condition size_int bstr size index : H.expr_primitive =
   Binary (Int_comp (I.Naked_immediate, Unsigned, Lt),
           untag_int index,
-          (Prim (Unary (Bigarray_length { dimension = 1; }, bstr))))
+          actual_max_length size_int size
+            (Prim (Unary (Bigarray_length { dimension = 1; }, bstr))))
 
 (* CR mshinwell: Same problems as previous function *)
-let bigstring_ref size arg1 arg2 dbg : H.expr_primitive =
+let bigstring_ref size_int size arg1 arg2 dbg : H.expr_primitive =
+  let wrap =
+    match (size : Flambda_primitive.string_accessor_width) with
+    | Eight | Sixteen -> tag_int
+    | Thirty_two -> box_bint Pint32
+    | Sixty_four -> box_bint Pint64
+  in
   Checked {
-    primitive = Binary (String_or_bigstring_load (Bigstring, size), arg1, arg2);
+    primitive =
+      wrap (Binary (String_or_bigstring_load (Bigstring, size), arg1, arg2));
     validity_conditions = [
-      bigstring_access_validity_condition arg1 arg2;
+      bigstring_access_validity_condition size_int arg1 size arg2;
     ];
     failure = Index_out_of_bounds;
     dbg;
   }
 
+let boxable_number_of_naked_number_kind k : K.Boxable_number.t =
+  match (k : K.Naked_number_kind.t) with
+  | Naked_immediate -> Untagged_immediate
+  | Naked_float -> Naked_float
+  | Naked_int32 -> Naked_int32
+  | Naked_int64 -> Naked_int64
+  | Naked_nativeint -> Naked_nativeint
+
+let bigarray_wrap_of_kind kind =
+  match P.element_kind_of_bigarray_kind kind with
+  | Value -> Fun.id
+  | Fabricated -> assert false
+  | Naked_number k ->
+    let bi = boxable_number_of_naked_number_kind k in
+    (fun arg -> H.Unary (Box_number bi, Prim arg))
+
+let bigarray_unwrap_of_kind kind =
+  match P.element_kind_of_bigarray_kind kind with
+  | Value -> Fun.id
+  | Fabricated -> assert false
+  | Naked_number k ->
+    let bi = boxable_number_of_naked_number_kind k in
+    (fun arg -> H.Prim (Unary (Unbox_number bi, arg)))
+
+let bigarray_map_last f l =
+  match List.rev l with
+  | [] -> assert false
+  | h :: r -> List.rev (f h :: r)
+
 let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
       (dbg : Debuginfo.t) : H.expr_primitive =
   let args = List.map (fun arg : H.simple_or_prim -> Simple arg) args in
@@ -215,9 +304,11 @@ let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
   | Pbytesrefu, [arg1; arg2] ->
     tag_int (Binary (String_or_bigstring_load (Bytes, Eight), arg1, arg2))
   | Pbytesrefs, [arg1; arg2] ->
-    string_or_bytes_ref Bytes arg1 arg2 dbg
+    let module B = (val backend : Flambda2_backend_intf.S) in
+    string_or_bytes_ref B.size_int Bytes arg1 arg2 dbg
   | Pstringrefs, [arg1; arg2] ->
-    string_or_bytes_ref String arg1 arg2 dbg
+    let module B = (val backend : Flambda2_backend_intf.S) in
+    string_or_bytes_ref B.size_int String arg1 arg2 dbg
   | Pstring_load_16 true (* unsafe *), [arg1; arg2]
   | Pbytes_load_16 true (* unsafe *), [arg1; arg2] ->
     tag_int (Binary (String_or_bigstring_load (String, Sixteen), arg1, arg2))
@@ -232,72 +323,78 @@ let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
       Prim (Binary (String_or_bigstring_load (String, Sixty_four),
         arg1, arg2)))
   | Pstring_load_16 false, [str; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         tag_int
           (Binary (String_or_bigstring_load (String, Sixteen), str, index));
       validity_conditions = [
-        string_or_bytes_access_validity_condition str String index;
+        string_or_bytes_access_validity_condition B.size_int str String Sixteen index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pstring_load_32 false, [str; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Unary (Box_number Naked_int32,
           Prim (Binary (String_or_bigstring_load (String, Thirty_two),
             str, index)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition str String index;
+        string_or_bytes_access_validity_condition B.size_int str String Thirty_two index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pstring_load_64 false, [str; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Unary (Box_number Naked_int64,
           Prim (Binary (String_or_bigstring_load (String, Sixty_four),
             str, index)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition str String index;
+        string_or_bytes_access_validity_condition B.size_int str String Sixty_four index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   (* CR mshinwell: factor out *)
   | Pbytes_load_16 false, [bytes; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         tag_int
           (Binary (String_or_bigstring_load (Bytes, Sixteen), bytes, index));
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Sixteen index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbytes_load_32 false, [bytes; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Unary (Box_number Naked_int32,
           Prim (Binary (String_or_bigstring_load (Bytes, Thirty_two),
             bytes, index)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Thirty_two index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbytes_load_64 false, [bytes; index] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Unary (Box_number Naked_int64,
           Prim (Binary (String_or_bigstring_load (Bytes, Sixty_four),
             bytes, index)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Sixty_four index;
       ];
       failure = Index_out_of_bounds;
       dbg;
@@ -313,34 +410,37 @@ let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
     Ternary (Bytes_or_bigstring_set (Bytes, Sixty_four),
       bytes, index, Prim (Unary (Unbox_number Naked_int64, new_value)))
   | Pbytes_set_16 false, [bytes; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bytes, Sixteen),
           bytes, index, untag_int new_value);
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Sixteen index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbytes_set_32 false, [bytes; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bytes, Thirty_two),
           bytes, index, Prim (Unary (Unbox_number Naked_int32, new_value)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Thirty_two index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbytes_set_64 false, [bytes; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bytes, Sixty_four),
           bytes, index, Prim (Unary (Unbox_number Naked_int64, new_value)));
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Sixty_four index;
       ];
       failure = Index_out_of_bounds;
       dbg;
@@ -631,12 +731,13 @@ let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
     Ternary (Bytes_or_bigstring_set (Bytes, Eight), bytes, index,
              untag_int new_value)
   | Pbytessets, [bytes; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bytes, Eight),
           bytes, index, untag_int new_value);
       validity_conditions = [
-        string_or_bytes_access_validity_condition bytes Bytes index;
+        string_or_bytes_access_validity_condition B.size_int bytes Bytes Eight index;
       ];
       failure = Index_out_of_bounds;
       dbg;
@@ -694,64 +795,73 @@ let convert_lprim ~backend (prim : L.primitive) (args : Simple.t list)
     let is_safe : P.is_safe = if unsafe then Unsafe else Safe in
     let kind = C.convert_bigarray_kind kind in
     let layout = C.convert_bigarray_layout layout in
-    Variadic (Bigarray_load (is_safe, num_dimensions, kind, layout), args)
+    let wrap = bigarray_wrap_of_kind kind in
+    wrap (Variadic (Bigarray_load (is_safe, num_dimensions, kind, layout), args))
   | Pbigarrayset (unsafe, num_dimensions, kind, layout), args ->
     let is_safe : P.is_safe = if unsafe then Unsafe else Safe in
     let kind = C.convert_bigarray_kind kind in
     let layout = C.convert_bigarray_layout layout in
-    Variadic (Bigarray_set (is_safe, num_dimensions, kind, layout), args)
+    let unwrap = bigarray_unwrap_of_kind kind in
+    let new_args = bigarray_map_last unwrap args in
+    Variadic (Bigarray_set (is_safe, num_dimensions, kind, layout), new_args)
   | Pbigarraydim dimension, [arg] ->
     tag_int (Unary (Bigarray_length { dimension; }, arg))
   | Pbigstring_load_16 true, [arg1; arg2] ->
     Binary (String_or_bigstring_load (Bigstring, Sixteen), arg1, arg2)
   | Pbigstring_load_16 false, [arg1; arg2] ->
-    bigstring_ref Sixteen arg1 arg2 dbg
+    let module B = (val backend : Flambda2_backend_intf.S) in
+    bigstring_ref B.size_int Sixteen arg1 arg2 dbg
   | Pbigstring_load_32 true, [arg1; arg2] ->
     Binary (String_or_bigstring_load (Bigstring, Thirty_two), arg1, arg2)
   | Pbigstring_load_32 false, [arg1; arg2] ->
-    bigstring_ref Thirty_two arg1 arg2 dbg
+    let module B = (val backend : Flambda2_backend_intf.S) in
+    bigstring_ref B.size_int Thirty_two arg1 arg2 dbg
   | Pbigstring_load_64 true, [arg1; arg2] ->
     Binary (String_or_bigstring_load (Bigstring, Sixty_four), arg1, arg2)
   | Pbigstring_load_64 false, [arg1; arg2] ->
-    bigstring_ref Sixty_four arg1 arg2 dbg
+    let module B = (val backend : Flambda2_backend_intf.S) in
+    bigstring_ref B.size_int Sixty_four arg1 arg2 dbg
   | Pbigstring_set_16 true, [bigstring; index; new_value] ->
     Ternary (Bytes_or_bigstring_set (Bigstring, Sixteen),
-      bigstring, index, new_value)
+      bigstring, index, untag_int new_value)
   | Pbigstring_set_32 true, [bigstring; index; new_value] ->
     Ternary (Bytes_or_bigstring_set (Bigstring, Thirty_two),
-      bigstring, index, new_value)
+      bigstring, index, Prim (Unary (Unbox_number Naked_int32, new_value)))
   | Pbigstring_set_64 true, [bigstring; index; new_value] ->
     Ternary (Bytes_or_bigstring_set (Bigstring, Sixty_four),
-      bigstring, index, new_value)
+      bigstring, index, Prim (Unary (Unbox_number Naked_int64, new_value)))
   | Pbigstring_set_16 false, [bigstring; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bigstring, Sixteen),
-          bigstring, index, new_value);
+          bigstring, index, untag_int new_value);
       validity_conditions = [
-        bigstring_access_validity_condition bigstring index;
+        bigstring_access_validity_condition B.size_int bigstring Sixteen index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbigstring_set_32 false, [bigstring; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bigstring, Thirty_two),
-          bigstring, index, new_value);
+          bigstring, index, Prim (Unary (Unbox_number Naked_int32, new_value)));
       validity_conditions = [
-        bigstring_access_validity_condition bigstring index;
+        bigstring_access_validity_condition B.size_int bigstring Thirty_two index;
       ];
       failure = Index_out_of_bounds;
       dbg;
     }
   | Pbigstring_set_64 false, [bigstring; index; new_value] ->
+    let module B = (val backend : Flambda2_backend_intf.S) in
     Checked {
       primitive =
         Ternary (Bytes_or_bigstring_set (Bigstring, Sixty_four),
-          bigstring, index, new_value);
+          bigstring, index, Prim (Unary (Unbox_number Naked_int64, new_value)));
       validity_conditions = [
-        bigstring_access_validity_condition bigstring index;
+        bigstring_access_validity_condition B.size_int bigstring Sixty_four index;
       ];
       failure = Index_out_of_bounds;
       dbg;