append, flatMap and flatten in Enum

enum.ml

@@ -84,6 +84,38 @@ let map f enum =
       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 () ->
     let l = ref l in

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