From e83298ef24a83bd553f4e977ae46bb2a63fcbf7f Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 19 Mar 2013 11:44:48 +0100
Subject: [PATCH] moved generator functions into Enum.Gen

---
 enum.ml            | 51 ++++++++++++++++++++++++++++------------------
 enum.mli           | 19 +++++++++++++----
 tests/test_enum.ml | 12 +++++++++--
 3 files changed, 56 insertions(+), 26 deletions(-)

diff --git a/enum.ml b/enum.ml
index e0ce1ec4..544de648 100644
--- a/enum.ml
+++ b/enum.ml
@@ -34,6 +34,33 @@ and 'a generator = unit -> 'a
   (** A generator may be called several times, yielding the next value
       each time. It raises EOG when it reaches the end. *)
 
+(** {2 Generator functions} *)
+
+let start enum = enum ()
+
+module Gen = struct
+  let next gen = gen ()
+
+  let junk gen = ignore (gen ())
+
+  let rec fold f acc gen =
+    let acc', stop =
+      try f acc (gen ()), false
+      with EOG -> acc, true in
+    if stop then acc' else fold f acc' gen
+
+  let rec iter f gen =
+    let stop =
+      try f (gen ()); false
+      with EOG -> true in
+    if stop then () else iter f gen
+
+  let length gen =
+    fold (fun acc _ -> acc + 1) 0 gen
+end
+
+(** {2 Basic constructors} *)
+
 let empty () = fun () -> raise EOG
 
 let singleton x =
@@ -57,36 +84,20 @@ let iterate x f =
       acc := f cur;
       cur
 
-let start enum = enum ()
-
-let next gen = gen ()
-
-let junk gen = ignore (gen ())
+(** {2 Basic combinators} *)
 
 let is_empty enum =
   try ignore ((enum ()) ()); false
   with EOG -> true
 
 let fold f acc enum =
-  let rec fold acc gen =
-    let acc', stop =
-      try f acc (gen ()), false
-      with EOG -> acc, true in
-    if stop then acc' else fold acc' gen
-  in
-  fold acc (enum ())
+  Gen.fold f acc (enum ())
 
 let iter f enum =
-  let rec iter gen =
-    let stop =
-      try f (gen ()); false
-      with EOG -> true in
-    if stop then () else iter gen
-  in
-  iter (enum ())
+  Gen.iter f (enum ())
 
 let length enum =
-  fold (fun acc _ -> acc + 1) 0 enum
+  Gen.length (enum ())
               
 let map f enum =
   (* another enum *)
diff --git a/enum.mli b/enum.mli
index c446b916..85089402 100644
--- a/enum.mli
+++ b/enum.mli
@@ -42,11 +42,22 @@ and 'a generator = unit -> 'a
 val start : 'a t -> 'a generator
   (** Create a new generator *)
 
-val next : 'a generator -> 'a
-  (** Get next element, or raise EOG *)
+module Gen : sig
+  val next : 'a generator -> 'a
+    (** Get next element, or raise EOG *)
 
-val junk : 'a generator -> unit
-  (** Drop element *)
+  val junk : 'a generator -> unit
+    (** Drop element *)
+
+  val fold : ('b -> 'a -> 'b) -> 'b -> 'a generator -> 'b
+    (** Fold over the generator *)
+
+  val iter : ('a -> unit) -> 'a generator -> unit
+    (** Iterate on the generator *)
+
+  val length : 'a generator -> int
+    (** Consume generator to compute its length *)
+end
 
 (** {2 Basic constructors} *)
 
diff --git a/tests/test_enum.ml b/tests/test_enum.ml
index fe7174c5..cdedd461 100644
--- a/tests/test_enum.ml
+++ b/tests/test_enum.ml
@@ -11,8 +11,8 @@ let pstrlist l = Utils.sprintf "%a"
 let test_singleton () =
   let e = Enum.singleton 42 in
   let gen = Enum.start e in
-  OUnit.assert_equal 42 (Enum.next gen);
-  OUnit.assert_raises Enum.EOG (fun () -> Enum.next gen);
+  OUnit.assert_equal 42 (Enum.Gen.next gen);
+  OUnit.assert_raises Enum.EOG (fun () -> Enum.Gen.next gen);
   OUnit.assert_equal 1 (Enum.length e);
   ()
 
@@ -79,6 +79,13 @@ let test_tee () =
     OUnit.assert_equal [2;4;6;8;10] (Enum.to_list b)
   | _ -> OUnit.assert_failure "wrong list lenght"
 
+let test_strong_tee () =
+  let e = Enum.tee ~n:3 (1 -- 999) in
+  let l = Enum.to_list e in
+  let l' = List.map Enum.Gen.length l in
+  OUnit.assert_equal [333;333;333] l';
+  ()
+
 let test_interleave () =
   let e1 = Enum.of_list [1;3;5;7;9] in
   let e2 = Enum.of_list [2;4;6;8;10] in
@@ -109,6 +116,7 @@ let suite =
       "test_round_robin" >:: test_round_robin;
       "test_persistent" >:: test_persistent;
       "test_tee" >:: test_tee;
+      "test_strong_tee" >:: test_strong_tee;
       "test_interleave" >:: test_interleave;
       "test_intersperse" >:: test_intersperse;
       "test_product" >:: test_product;