Enum.take and Enum.drop

enum.ml

@@ -130,6 +130,26 @@ let flatMap f enum =
           with EOG -> true in
         if stop then raise EOG else next ()
     in next
+
+let take n enum =
+  assert (n >= 0);
+  fun () ->
+    let gen = enum () in
+    let count = ref 0 in  (* how many yielded elements *)
+    fun () ->
+      if !count = n then raise EOG
+      else begin incr count; gen () end
+
+let drop n enum =
+  assert (n >= 0);
+  fun () ->
+    let gen = enum () in
+    let count = ref 0 in  (* how many droped elements? *)
+    let rec next () =
+      if !count < n
+        then begin incr count; ignore (gen ()); next () end
+        else gen ()
+    in next
     
 let of_list l =
   fun () ->
@@ -141,7 +161,9 @@ let of_list l =
 
 let to_list enum =
   let rec fold gen =
-    try (gen ()) :: fold gen
+    try
+      let x = gen () in
+      x :: fold gen
     with EOG -> []
   in fold (enum ())
 
@@ -158,9 +180,9 @@ let int_range i j =
   fun () ->
     let r = ref i in
     fun () ->
-      if !r > j then raise EOG
+      let x = !r in
+      if x > j then raise EOG
         else begin
-          let x = !r in
           incr r;
           x
         end

enum.mli

@@ -76,6 +76,12 @@ val flatten : 'a t t -> 'a t
 val flatMap : ('a -> 'b t) -> 'a t -> 'b t
   (** Monadic bind *)
 
+val take : int -> 'a t -> 'a t
+  (** Take at most n elements *)
+
+val drop : int -> 'a t -> 'a t
+  (** Drop n elements *)
+
 val of_list : 'a list -> 'a t
   (** Enumerate the list *)
 

tests/test_enum.ml

@@ -3,6 +3,10 @@ open OUnit
 open Enum.Infix
 
 let pint i = string_of_int i
+let plist l = Utils.sprintf "%a"
+  (Sequence.pp_seq Format.pp_print_int) (Sequence.of_list l)
+let pstrlist l = Utils.sprintf "%a"
+  (Sequence.pp_seq Format.pp_print_string) (Sequence.of_list l)
 
 let test_singleton () =
   let e = Enum.singleton 42 in
@@ -15,11 +19,19 @@ let test_singleton () =
 let test_iter () =
   let e = 1 -- 10 in
   OUnit.assert_equal ~printer:pint 10 (Enum.length e);
-  (* TODO *)
+  OUnit.assert_equal [1;2] (Enum.to_list (1 -- 2));
+  OUnit.assert_equal [1;2;3;4;5] (Enum.to_list (Enum.take 5 e));
+  ()
+
+let test_map () =
+  let e = 1 -- 10 in
+  let e' = Enum.map string_of_int e in
+  OUnit.assert_equal ~printer:pstrlist ["9"; "10"] (Enum.to_list (Enum.drop 8 e'));
   ()
 
 let suite =
   "test_enum" >:::
     [ "test_singleton" >:: test_singleton;
       "test_iter" >:: test_iter;
+      "test_map" >:: test_map;
     ]