From 135989d55f0aa5664935810f2526c7e4b3a5f7b9 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 19 Mar 2013 00:22:41 +0100
Subject: [PATCH] Enum.take and Enum.drop

---
 enum.ml            | 28 +++++++++++++++++++++++++---
 enum.mli           |  6 ++++++
 tests/test_enum.ml | 14 +++++++++++++-
 3 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/enum.ml b/enum.ml
index a12d87a4..8b19620d 100644
--- a/enum.ml
+++ b/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
diff --git a/enum.mli b/enum.mli
index 4610da6f..33db8e33 100644
--- a/enum.mli
+++ b/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 *)
 
diff --git a/tests/test_enum.ml b/tests/test_enum.ml
index 3f2a835e..9037a6cd 100644
--- a/tests/test_enum.ml
+++ b/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;
     ]