From fed7dd41e063454883cd6eaee83bf894de00796a Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 13 Oct 2016 10:04:58 +0200
Subject: [PATCH] add functions in `CCArray`: fold2,iter2,map2

---
 src/core/CCArray.ml  | 52 ++++++++++++++++++++++++++++++++++++++++----
 src/core/CCArray.mli | 15 +++++++++++++
 2 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/core/CCArray.ml b/src/core/CCArray.ml
index 4890a9a3..6daf0f08 100644
--- a/src/core/CCArray.ml
+++ b/src/core/CCArray.ml
@@ -106,8 +106,18 @@ module type S = sig
       @raise Invalid_argument if they have distinct lengths
       allow different types @since NEXT_RELEASE *)
 
+  val fold2 : ('acc -> 'a -> 'b -> 'acc) -> 'acc -> 'a t -> 'b t -> 'acc
+  (** Fold on two arrays stepwise.
+      @raise Invalid_argument if they have distinct lengths
+      @since NEXT_RELEASE *)
+
+  val iter2 : ('a -> 'b -> unit) -> 'a t -> 'b t -> unit
+  (** Iterate on two arrays stepwise.
+      @raise Invalid_argument if they have distinct lengths
+      @since NEXT_RELEASE *)
+
   val shuffle : 'a t -> unit
-  (** shuffle randomly the array, in place *)
+  (** Shuffle randomly the array, in place *)
 
   val shuffle_with : Random.State.t -> 'a t -> unit
   (** Like shuffle but using a specialized random state *)
@@ -124,15 +134,15 @@ module type S = sig
 
   val pp: ?sep:string -> (Buffer.t -> 'a -> unit) ->
           Buffer.t -> 'a t -> unit
-  (** print an array of items with printing function *)
+  (** Print an array of items with printing function *)
 
   val pp_i: ?sep:string -> (Buffer.t -> int -> 'a -> unit) ->
             Buffer.t -> 'a t -> unit
-  (** print an array, giving the printing function both index and item *)
+  (** Print an array, giving the printing function both index and item *)
 
   val print : ?sep:string -> (Format.formatter -> 'a -> unit) ->
               Format.formatter -> 'a t -> unit
-  (** print an array of items with printing function *)
+  (** Print an array of items with printing function *)
 end
 
 (** {2 General Implementation}
@@ -293,6 +303,10 @@ let empty = [| |]
 
 let map = Array.map
 
+let map2 f a b =
+  if Array.length a <> Array.length b then invalid_arg "map2";
+  Array.init (Array.length a) (fun i -> f (Array.unsafe_get a i) (Array.unsafe_get b i))
+
 let length = Array.length
 
 let get = Array.get
@@ -472,6 +486,28 @@ let for_all2 p a b =
 let exists2 p a b =
   _exists2 p a b 0 0 ~len:(min (Array.length a) (Array.length b))
 
+let _iter2 f a b i j ~len =
+  for o = 0 to len-1 do
+    f (Array.get a (i+o)) (Array.get b (j+o))
+  done
+
+let _fold2 f acc a b i j ~len =
+  let rec aux acc o =
+    if o=len then acc
+    else
+      let acc = f acc (Array.get a (i+o)) (Array.get b (j+o)) in
+      aux acc (o+1)
+  in
+  aux acc 0
+
+let iter2 f a b =
+  if length a <> length b then invalid_arg "iter2";
+  _iter2 f a b 0 0 ~len:(Array.length a)
+
+let fold2 f acc a b =
+  if length a <> length b then invalid_arg "fold2";
+  _fold2 f acc a b 0 0 ~len:(Array.length a)
+
 let (--) i j =
   if i<=j
   then
@@ -675,6 +711,14 @@ module Sub = struct
   Sub.exists2 (=) (Sub.make [| 1;2;3;4 |] 1 ~len:2) (Sub.make [| 0;1;3;4 |] 1 ~len:3)
   *)
 
+  let iter2 f a b =
+    if length a <> length b then invalid_arg "iter2";
+    _iter2 f a.arr b.arr a.i b.i ~len:(length a)
+
+  let fold2 f acc a b =
+    if length a <> length b then invalid_arg "fold2";
+    _fold2 f acc a.arr b.arr a.i b.i ~len:(length a)
+
   let shuffle a =
     _shuffle Random.int a.arr a.i a.j
 
diff --git a/src/core/CCArray.mli b/src/core/CCArray.mli
index 6b417c51..24be5897 100644
--- a/src/core/CCArray.mli
+++ b/src/core/CCArray.mli
@@ -108,6 +108,16 @@ module type S = sig
       @raise Invalid_argument if they have distinct lengths
       allow different types @since NEXT_RELEASE *)
 
+  val fold2 : ('acc -> 'a -> 'b -> 'acc) -> 'acc -> 'a t -> 'b t -> 'acc
+  (** Fold on two arrays stepwise.
+      @raise Invalid_argument if they have distinct lengths
+      @since NEXT_RELEASE *)
+
+  val iter2 : ('a -> 'b -> unit) -> 'a t -> 'b t -> unit
+  (** Iterate on two arrays stepwise.
+      @raise Invalid_argument if they have distinct lengths
+      @since NEXT_RELEASE *)
+
   val shuffle : 'a t -> unit
   (** Shuffle randomly the array, in place *)
 
@@ -145,6 +155,11 @@ include S with type 'a t := 'a t
 
 val map : ('a -> 'b) -> 'a t -> 'b t
 
+val map2 : ('a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
+(** Map on two arrays stepwise.
+      @raise Invalid_argument if they have distinct lengths
+      @since NEXT_RELEASE *)
+
 val rev : 'a t -> 'a t
 (** Copy + reverse in place
     @since NEXT_RELEASE *)