From d033b4621c67e9e056898cf19996696bf0e9c86f Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 14 Sep 2015 19:42:57 +0200
Subject: [PATCH] add fair functions to `CCKList`

---
 src/iter/CCKList.ml  | 38 ++++++++++++++++++++++++++++++++++++--
 src/iter/CCKList.mli | 22 ++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/iter/CCKList.ml b/src/iter/CCKList.ml
index adf6421e..6adf9d1d 100644
--- a/src/iter/CCKList.ml
+++ b/src/iter/CCKList.ml
@@ -108,8 +108,12 @@ let rec take n (l:'a t) () = match l () with
 
 let rec take_while p l () = match l () with
   | `Nil -> `Nil
-  | `Cons (x,l') when p x -> `Cons (x, take_while p l')
-  | `Cons (_,l') -> take_while p l' ()
+  | `Cons (x,l') ->
+      if p x then `Cons (x, take_while p l') else `Nil
+
+(*$T
+  of_list [1;2;3;4] |> take_while (fun x->x < 4) |> to_list = [1;2;3]
+*)
 
 let rec drop n (l:'a t) () = match l () with
   | l' when n=0 -> l'
@@ -229,6 +233,11 @@ let rec group eq l () = match l() with
   | `Cons (x, l') ->
       `Cons (cons x (take_while (eq x) l'), group eq (drop_while (eq x) l'))
 
+(*$T
+  of_list [1;1;1;2;2;3;3;1] |> group (=) |> map to_list |> to_list = \
+    [[1;1;1]; [2;2]; [3;3]; [1]]
+*)
+
 let rec _uniq eq prev l () = match prev, l() with
   | _, `Nil -> `Nil
   | None, `Cons (x, l') ->
@@ -431,6 +440,31 @@ let sort_uniq ?(cmp=Pervasives.compare) l =
   let l = to_list l in
   uniq (fun x y -> cmp x y = 0) (of_list (List.sort cmp l))
 
+(** {2 Fair Combinations} *)
+
+let rec interleave a b () = match a() with
+  | `Nil -> b ()
+  | `Cons (x, tail) -> `Cons (x, interleave b tail)
+
+let rec fair_flat_map f a () = match a() with
+  | `Nil -> `Nil
+  | `Cons (x, tail) ->
+      let y = f x in
+      interleave y (fair_flat_map f tail) ()
+
+let rec fair_app f a () = match f() with
+  | `Nil -> `Nil
+  | `Cons (f1, fs) ->
+      interleave (map f1 a) (fair_app fs a) ()
+
+let (>>-) a f = fair_flat_map f a
+let (<.>) f a = fair_app f a
+
+(*$T
+  interleave (of_list [1;3;5]) (of_list [2;4;6]) |> to_list = [1;2;3;4;5;6]
+  fair_app (of_list [(+)1; ( * ) 3]) (of_list [1; 10]) \
+    |> to_list |> List.sort Pervasives.compare = [2; 3; 11; 30]
+*)
 
 (** {2 Monadic Operations} *)
 module type MONAD = sig
diff --git a/src/iter/CCKList.mli b/src/iter/CCKList.mli
index 268fd39c..e2998296 100644
--- a/src/iter/CCKList.mli
+++ b/src/iter/CCKList.mli
@@ -191,6 +191,20 @@ val sort_uniq : ?cmp:'a ord -> 'a t -> 'a t
     finite. O(n ln(n)) time and space.
     @since 0.3.3 *)
 
+(** {2 Fair Combinations} *)
+
+val interleave : 'a t -> 'a t -> 'a t
+(** Fair interleaving of both streams.
+    @since NEXT_RELEASE *)
+
+val fair_flat_map : ('a -> 'b t) -> 'a t -> 'b t
+(** Fair version of {!flat_map}.
+    @since NEXT_RELEASE *)
+
+val fair_app : ('a -> 'b) t -> 'a t -> 'b t
+(** Fair version of {!(<*>)}
+    @since NEXT_RELEASE *)
+
 (** {2 Implementations}
     @since 0.3.3 *)
 
@@ -200,6 +214,14 @@ val (>>=) : 'a t -> ('a -> 'b t) -> 'b t
 val (>|=) : 'a t -> ('a -> 'b) -> 'b t
 val (<*>) : ('a -> 'b) t -> 'a t -> 'b t
 
+val (>>-) : 'a t -> ('a -> 'b t) -> 'b t
+(** Infix version of {! fair_flat_map}
+    @since NEXT_RELEASE *)
+
+val (<.>) : ('a -> 'b) t -> 'a t -> 'b t
+(** Infix version of {!fair_app}
+    @since NEXT_RELEASE *)
+
 (** {2 Monadic Operations} *)
 module type MONAD = sig
   type 'a t