Merge pull request #236

PR#6263: Add size_in_bytes and kind_size_in_bytes to Bigarray

Changes

@@ -285,6 +285,9 @@ Other libraries:
   Before, a handled signal could cause Unix.sleep to return early.
   Now, the sleep is restarted until the given time is elapsed.
   (Xavier Leroy)
+- PR#6263: add kind_size_in_bytes and size_in_bytes functions
+  to Bigarray module.
+  (Runhang Li, review by Mark Shinwell)
 - PR#6289: Unix.utimes uses the current time only if both arguments
     are exactly 0.0.  Also, use sub-second resolution if available.
   (Xavier Leroy, report by Christophe Troestler)

otherlibs/bigarray/bigarray.ml

@@ -62,6 +62,21 @@ let complex32 = Complex32
 let complex64 = Complex64
 let char = Char
 
+let kind_size_in_bytes : type a b. (a, b) kind -> int = function
+  | Float32 -> 4
+  | Float64 -> 8
+  | Int8_signed -> 1
+  | Int8_unsigned -> 1
+  | Int16_signed -> 2
+  | Int16_unsigned -> 2
+  | Int32 -> 4
+  | Int64 -> 8
+  | Int -> Sys.word_size / 8
+  | Nativeint -> Sys.word_size / 8
+  | Complex32 -> 8
+  | Complex64 -> 16
+  | Char -> 1
+
 type c_layout = C_layout_typ
 type fortran_layout = Fortran_layout_typ
 
@@ -90,9 +105,13 @@ module Genarray = struct
     let d = Array.make n 0 in
     for i = 0 to n-1 do d.(i) <- nth_dim a i done;
     d
+
   external kind: ('a, 'b, 'c) t -> ('a, 'b) kind = "caml_ba_kind"
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
 
+  let size_in_bytes arr =
+    (kind_size_in_bytes (kind arr)) * (Array.fold_left ( * ) 1 (dims arr))
+
   external sub_left: ('a, 'b, c_layout) t -> int -> int -> ('a, 'b, c_layout) t
      = "caml_ba_sub"
   external sub_right: ('a, 'b, fortran_layout) t -> int -> int ->
@@ -126,6 +145,10 @@ module Array1 = struct
   external dim: ('a, 'b, 'c) t -> int = "%caml_ba_dim_1"
   external kind: ('a, 'b, 'c) t -> ('a, 'b) kind = "caml_ba_kind"
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
+
+  let size_in_bytes arr =
+    (kind_size_in_bytes (kind arr)) * (dim arr)
+
   external sub: ('a, 'b, 'c) t -> int -> int -> ('a, 'b, 'c) t = "caml_ba_sub"
   external blit: ('a, 'b, 'c) t -> ('a, 'b, 'c) t -> unit = "caml_ba_blit"
   external fill: ('a, 'b, 'c) t -> 'a -> unit = "caml_ba_fill"
@@ -156,6 +179,10 @@ module Array2 = struct
   external dim2: ('a, 'b, 'c) t -> int = "%caml_ba_dim_2"
   external kind: ('a, 'b, 'c) t -> ('a, 'b) kind = "caml_ba_kind"
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
+
+  let size_in_bytes arr =
+    (kind_size_in_bytes (kind arr)) * (dim1 arr) * (dim2 arr)
+
   external sub_left: ('a, 'b, c_layout) t -> int -> int -> ('a, 'b, c_layout) t
      = "caml_ba_sub"
   external sub_right:
@@ -203,6 +230,10 @@ module Array3 = struct
   external dim3: ('a, 'b, 'c) t -> int = "%caml_ba_dim_3"
   external kind: ('a, 'b, 'c) t -> ('a, 'b) kind = "caml_ba_kind"
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
+
+  let size_in_bytes arr =
+    (kind_size_in_bytes (kind arr)) * (dim1 arr) * (dim2 arr) * (dim3 arr)
+
   external sub_left: ('a, 'b, c_layout) t -> int -> int -> ('a, 'b, c_layout) t
      = "caml_ba_sub"
   external sub_right:

otherlibs/bigarray/bigarray.mli

@@ -168,6 +168,10 @@ val char : (char, int8_unsigned_elt) kind
    characters instead of arrays of small integers, by using
    the kind value [char] instead of [int8_unsigned]. *)
 
+val kind_size_in_bytes : ('a, 'b) kind -> int
+(** [kind_size_in_bytes k] is the number of bytes used to store
+   an element of type [k]. *)
+
 (** {6 Array layouts} *)
 
 type c_layout = C_layout_typ (**)
@@ -280,6 +284,10 @@ module Genarray :
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
   (** Return the layout of the given big array. *)
 
+  val size_in_bytes : ('a, 'b, 'c) t -> int
+  (** [size_in_bytes a] is the number of elements in [a] multiplied
+    by [a]'s {!kind_size_in_bytes}.*)
+
   external get: ('a, 'b, 'c) t -> int array -> 'a = "caml_ba_get_generic"
   (** Read an element of a generic big array.
      [Genarray.get a [|i1; ...; iN|]] returns the element of [a]
@@ -490,6 +498,10 @@ module Array1 : sig
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
   (** Return the layout of the given big array. *)
 
+  val size_in_bytes : ('a, 'b, 'c) t -> int
+  (** [size_in_bytes a] is the number of elements in [a]
+    multiplied by [a]'s {!kind_size_in_bytes}. *)
+
   external get: ('a, 'b, 'c) t -> int -> 'a = "%caml_ba_ref_1"
   (** [Array1.get a x], or alternatively [a.{x}],
      returns the element of [a] at index [x].
@@ -572,6 +584,10 @@ module Array2 :
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
   (** Return the layout of the given big array. *)
 
+  val size_in_bytes : ('a, 'b, 'c) t -> int
+  (** [[size_in_bytes a] is the number of elements in [a]
+    multiplied by [a]'s {!kind_size_in_bytes}. *)
+
   external get: ('a, 'b, 'c) t -> int -> int -> 'a = "%caml_ba_ref_2"
   (** [Array2.get a x y], also written [a.{x,y}],
      returns the element of [a] at coordinates ([x], [y]).
@@ -678,6 +694,10 @@ module Array3 :
   external layout: ('a, 'b, 'c) t -> 'c layout = "caml_ba_layout"
   (** Return the layout of the given big array. *)
 
+  val size_in_bytes : ('a, 'b, 'c) t -> int
+  (** [size_in_bytes a] is the number of elements in [a]
+    multiplied by [a]'s {!kind_size_in_bytes}. *)
+
   external get: ('a, 'b, 'c) t -> int -> int -> int -> 'a = "%caml_ba_ref_3"
   (** [Array3.get a x y z], also written [a.{x,y,z}],
      returns the element of [a] at coordinates ([x], [y], [z]).

testsuite/tests/lib-bigarray/bigarrays.ml

@@ -411,6 +411,14 @@ let _ =
   test 1 (Array1.dim (from_list int [1;2;3;4;5])) 5;
   test 2 (Array1.dim (from_list_fortran int [1;2;3])) 3;
 
+  testing_function "size_in_bytes_one";
+  test 1 (Array1.size_in_bytes (from_list int [1;2;3;4;5]))
+    (5 * (kind_size_in_bytes int));
+  test 2 (Array1.size_in_bytes (from_list int [])) 0;
+  let int64list = (from_list int64 (List.map Int64.of_int [1;2;3;4;5])) in
+  test 3 (Array1.size_in_bytes int64list) (5 * (kind_size_in_bytes int64));
+  test 4 (Array1.size_in_bytes (from_list int64 (List.map Int64.of_int []))) 0;
+
   testing_function "kind & layout";
   let a = from_list int [1;2;3] in
   test 1 (Array1.kind a) int;
@@ -595,6 +603,10 @@ let _ =
   test 3 (Array2.dim1 b) 4;
   test 4 (Array2.dim2 b) 6;
 
+  testing_function "size_in_bytes_two";
+  let a = Array2.create int c_layout 4 6 in
+  test 1 (Array2.size_in_bytes a) (24 * (kind_size_in_bytes int));
+
   testing_function "sub";
   let a = make_array2 int c_layout 0 5 3 id in
   let b = Array2.sub_left a 2 2 in
@@ -746,6 +758,10 @@ let _ =
   test 5 (Array3.dim2 b) 5;
   test 6 (Array3.dim3 b) 6;
 
+  testing_function "size_in_bytes_three";
+  let a = Array3.create int c_layout 4 5 6 in
+  test 1 (Array3.size_in_bytes a) (120 * (kind_size_in_bytes int));
+
   testing_function "slice1";
   let a = make_array3 int c_layout 0 3 3 3 id in
   test 1 (Array3.slice_left_1 a 0 0) (from_list int [0;1;2]);
@@ -757,6 +773,39 @@ let _ =
   test 6 (Array3.slice_right_1 a 1 2) (from_list_fortran int [112;212;312]);
   test 7 (Array3.slice_right_1 a 3 1) (from_list_fortran int [131;231;331]);
 
+  testing_function "size_in_bytes_general";
+  let a = Genarray.create int c_layout [|2;2;2;2;2|] in
+  test 1 (Genarray.size_in_bytes a) (32 * (kind_size_in_bytes int));
+
+(* Kind size *)
+  testing_function "kind_size_in_bytes";
+  let arr1 = Array1.create Float32 c_layout 1 in
+  test 1 (kind_size_in_bytes Float32) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Float64 c_layout 1 in
+  test 2 (kind_size_in_bytes Float64) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int8_signed c_layout 1 in
+  test 3 (kind_size_in_bytes Int8_signed) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int8_unsigned c_layout 1 in
+  test 4 (kind_size_in_bytes Int8_unsigned) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int16_signed c_layout 1 in
+  test 5 (kind_size_in_bytes Int16_signed) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int16_unsigned c_layout 1 in
+  test 6 (kind_size_in_bytes Int16_unsigned) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int32 c_layout 1 in
+  test 7 (kind_size_in_bytes Int32) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int64 c_layout 1 in
+  test 8 (kind_size_in_bytes Int64) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Int c_layout 1 in
+  test 9 (kind_size_in_bytes Int) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Nativeint c_layout 1 in
+  test 10 (kind_size_in_bytes Nativeint) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Complex32 c_layout 1 in
+  test 11 (kind_size_in_bytes Complex32) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Complex64 c_layout 1 in
+  test 12 (kind_size_in_bytes Complex64) (Array1.size_in_bytes arr1);
+  let arr1 = Array1.create Char c_layout 1 in
+  test 13 (kind_size_in_bytes Char) (Array1.size_in_bytes arr1);
+
 (* Reshaping *)
   print_newline();
   testing_function "------ Reshaping --------";

testsuite/tests/lib-bigarray/bigarrays.reference

@@ -11,6 +11,8 @@ comparisons
  1... 2... 3... 4... 5... 6... 7... 8... 9... 10... 11... 12... 13... 14... 15... 16... 17... 18... 19... 20... 21... 22... 23... 24... 25... 26... 27... 28... 29... 30... 31... 32... 44... 45... 46... 47... 48... 49...
 dim
  1... 2...
+size_in_bytes_one
+ 1... 2... 3... 4...
 kind & layout
  1... 2... 1... 2...
 sub
@@ -28,6 +30,8 @@ set/get (unsafe, specialized)
  1... 2...
 dim
  1... 2... 3... 4...
+size_in_bytes_two
+ 1...
 sub
  1... 2...
 slice
@@ -43,8 +47,14 @@ set/get (unsafe, specialized)
  1...
 dim
  1... 2... 3... 4... 5... 6...
+size_in_bytes_three
+ 1...
 slice1
  1... 2... 3... 4... 5... 6... 7...
+size_in_bytes_general
+ 1...
+kind_size_in_bytes
+ 1... 2... 3... 4... 5... 6... 7... 8... 9... 10... 11... 12... 13...
 
 ------ Reshaping --------