From 5b4e2055912da52f955999b4bfddda42bbf2b804 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 19 Mar 2013 23:59:07 +0100
Subject: [PATCH] Enum.dup combinator

---
 enum.ml  | 34 ++++++++++++++++++++++++++++++++++
 enum.mli |  5 +++++
 2 files changed, 39 insertions(+)

diff --git a/enum.ml b/enum.ml
index 31067fa0..1637d4bb 100644
--- a/enum.ml
+++ b/enum.ml
@@ -368,6 +368,40 @@ let tee ?(n=2) enum =
       let j = !i in
       if j = n then raise EOG else (incr i; fun () -> next j)
 
+(** Duplicate the enum into [n] generators (default 2). The generators
+    share the same underlying instance of the enum, so the optimal case is
+    when they are consumed evenly *)
+let dup ?(n=2) enum =
+  fun () ->
+    (* array of queues, together with their index *)
+    let qs = Array.init n (fun i -> Queue.create ()) in
+    let gen = enum () in  (* unique generator! *)
+    let finished = ref false in (* is [gen] exhausted? *)
+    (* get next element for the i-th queue *)
+    let rec next i =
+      if Queue.is_empty qs.(i)
+        then
+          if !finished then raise EOG
+          else get_next i  (* consume generator *)
+        else Queue.pop qs.(i)
+    (* consume one more element *)
+    and get_next i =
+      try
+        let x = gen () in
+        for j = 0 to n-1 do
+          if j <> i then Queue.push x qs.(j)
+        done;
+        x
+      with EOG ->
+        finished := true;
+        raise EOG
+    in
+    (* generator of generators *)
+    let i = ref 0 in
+    fun () ->
+      let j = !i in
+      if j = n then raise EOG else (incr i; fun () -> next j)
+
 (** Yield elements from a and b alternatively *)
 let interleave a b =
   fun () ->
diff --git a/enum.mli b/enum.mli
index 4406b45f..b089a0d4 100644
--- a/enum.mli
+++ b/enum.mli
@@ -144,6 +144,11 @@ val tee : ?n:int -> 'a t -> 'a generator t
       [index = k mod n] with go to the k-th enum. [n] defaults value
       is 2. *)
 
+val dup : ?n:int -> 'a t -> 'a generator t
+  (** Duplicate the enum into [n] generators (default 2). The generators
+      share the same underlying instance of the enum, so the optimal case is
+      when they are consumed evenly *)
+
 val interleave : 'a t -> 'a t -> 'a t
   (** [interleave a b] yields an element of [a], then an element of [b],
       and so on until the end of [a] or [b] is reached. *)