From a378af8a8e90543918133721cbb32c340e1d8c99 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 18 Mar 2013 23:42:29 +0100
Subject: [PATCH] append, flatMap and flatten in Enum

---
 enum.ml  | 32 ++++++++++++++++++++++++++++++++
 enum.mli |  9 +++++++++
 2 files changed, 41 insertions(+)

diff --git a/enum.ml b/enum.ml
index 10c62714..f73fe542 100644
--- a/enum.ml
+++ b/enum.ml
@@ -83,6 +83,38 @@ let map f enum =
     fun () ->
       try f (gen ())
       with EOG -> raise EOG
+
+let append e1 e2 =
+  fun () ->
+    let gen = ref (e1 ()) in
+    let first = ref true in
+    (* get next element *)
+    let rec next () =
+      try !gen ()
+      with EOG ->
+        if !first then begin
+          first := false;
+          gen := e2 ();  (* switch to the second generator *)
+          next ()
+        end else raise EOG  (* done *)
+    in next
+
+let flatten enum =
+  fun () ->
+    let next_gen = enum () in
+    let gen = ref (fun () -> raise EOG) in
+    (* get next element *)
+    let rec next () =
+      try !gen ()
+      with EOG ->
+        (* jump to next sub-enum *)
+        let stop =
+          try gen := !next_gen (); false
+          with EOG -> true in
+        if stop then raise EOG else next ()
+    in next
+      
+let flatMap f enum = flatten (map f enum)
     
 let of_list l =
   fun () ->
diff --git a/enum.mli b/enum.mli
index 4b2229c3..344f9fd2 100644
--- a/enum.mli
+++ b/enum.mli
@@ -67,6 +67,15 @@ val length : _ t -> int
 val map : ('a -> 'b) -> 'a t -> 'b t
   (** Lazy map *)
 
+val append : 'a t -> 'a t -> 'a t
+  (** Append the two enums *)
+
+val flatten : 'a t t -> 'a t
+  (** Flatten the enum of enum *)
+
+val flatMap : ('a -> 'b t) -> 'a t -> 'b t
+  (** Monadic bind *)
+
 val of_list : 'a list -> 'a t
   (** Enumerate the list *)