moved generator functions into Enum.Gen

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 *)

enum.mli

@@ -42,12 +42,23 @@ and 'a generator = unit -> 'a
 val start : 'a t -> 'a generator
   (** Create a new generator *)
 
+module Gen : sig
   val next : 'a generator -> 'a
     (** Get next element, or raise EOG *)
 
   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} *)
 
 val empty : 'a t

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;