From 96d3c7e8b79868f3fe1e1e56384f69b6e8e5c26a Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Tue, 19 Mar 2013 11:25:26 +0100
Subject: [PATCH] added a bunch of combinators for Enum, some of them quite
 advanced. A few are not yet implemented.

---
 enum.ml  | 288 ++++++++++++++++++++++++++++++++++++++++++++++++++++---
 enum.mli |  79 ++++++++++++++-
 2 files changed, 347 insertions(+), 20 deletions(-)

diff --git a/enum.ml b/enum.ml
index 1d00b3f4..e0ce1ec4 100644
--- a/enum.ml
+++ b/enum.ml
@@ -44,6 +44,19 @@ let singleton x =
         then raise EOG
         else begin stop := true; x end
 
+let repeat x =
+  let f () = x in
+  fun () -> f
+
+(** [iterate x f] is [[x; f x; f (f x); f (f (f x)); ...]] *)
+let iterate x f =
+  fun () ->
+    let acc = ref x in
+    fun () ->
+      let cur = !acc in
+      acc := f cur;
+      cur
+
 let start enum = enum ()
 
 let next gen = gen ()
@@ -99,6 +112,17 @@ let append e1 e2 =
         end else raise EOG  (* done *)
     in next
 
+let cycle enum =
+  assert (not (is_empty enum));
+  fun () ->
+    let gen = ref (enum ()) in
+    let rec next () =
+      try !gen ()
+      with EOG ->
+        gen := enum ();
+        next ()
+    in next
+
 let flatten enum =
   fun () ->
     let next_gen = enum () in
@@ -150,6 +174,254 @@ let drop n enum =
         then begin incr count; ignore (gen ()); next () end
         else gen ()
     in next
+
+let filter p enum =
+  fun () ->
+    let gen = enum () in
+    let rec next () =
+      match (try Some (gen ()) with EOG -> None) with
+      | None -> raise EOG
+      | Some x ->
+        if p x
+          then x (* yield element *)
+          else next ()  (* discard element *)
+    in next
+
+let takeWhile p enum =
+  fun () ->
+    let gen = enum () in
+    let rec next () =
+      match (try Some (gen ()) with EOG -> None) with
+      | None -> raise EOG
+      | Some x ->
+        if p x
+          then x (* yield element *)
+          else raise EOG (* stop *)
+    in next
+
+let dropWhile p enum =
+  fun () ->
+    let gen = enum () in
+    let stop_drop = ref false in
+    let rec next () =
+      match (try Some (gen ()) with EOG -> None) with
+      | None -> raise EOG
+      | Some x when !stop_drop -> x (* yield *)
+      | Some x ->
+        if p x
+          then next ()  (* drop *)
+          else (stop_drop := true; x) (* stop dropping, and yield *)
+    in next
+
+let filterMap f enum =
+  fun () ->
+    let gen = enum () in
+    (* tailrec *)
+    let rec next () =
+      match (try Some (gen ()) with EOG -> None) with
+      | None -> raise EOG
+      | Some x ->
+        begin
+          match f x with
+          | None -> next ()  (* drop element *)
+          | Some y -> y  (* return [f x] *)
+        end
+    in next
+
+let zipWith f a b =
+  fun () ->
+    let gen_a = a () in
+    let gen_b = b () in
+    fun () ->
+      f (gen_a ()) (gen_b ())
+
+let zip a b = zipWith (fun x y -> x,y) a b
+
+let zipIndex enum =
+  fun () ->
+    let r = ref 0 in
+    let gen = enum () in
+    fun () ->
+      let x = gen () in
+      let n = !r in
+      incr r;
+      n, x
+
+(** {2 Complex combinators} *)
+
+(** Pick elements fairly in each sub-enum *)
+let round_robin enum =
+  (* list of sub-enums *)
+  let l = fold (fun acc x -> x::acc) [] enum in
+  let l = List.rev l in
+  fun () ->
+    let q = Queue.create () in
+    List.iter (fun enum' -> Queue.push (enum' ()) q) l;
+    (* recursive function to get next element *)
+    let rec next () =
+      if Queue.is_empty q
+        then raise EOG
+        else
+          let gen = Queue.pop q in
+          match (try Some (gen ()) with EOG -> None) with
+          | None -> next ()  (* exhausted generator, drop it *)
+          | Some x ->
+            Queue.push gen q; (* put generator at the end, return x *)
+            x
+    in next
+
+(** {3 Mutable double-linked list, similar to {! Deque.t} *)
+module MList = struct
+  type 'a t = 'a node option ref
+  and 'a node = {
+    content : 'a;
+    mutable prev : 'a node;
+    mutable next : 'a node;
+  }
+
+  let create () = ref None
+
+  let is_empty d =
+    match !d with
+    | None -> true
+    | Some _ -> false
+
+  let push_back d x =
+    match !d with
+    | None ->
+      let rec elt = {
+        content = x; prev = elt; next = elt; } in
+      d := Some elt
+    | Some first ->
+      let elt = { content = x; next=first; prev=first.prev; } in
+      first.prev.next <- elt;
+      first.prev <- elt
+
+  (* conversion to enum *)
+  let to_enum d =
+    fun () ->
+      match !d with
+      | None -> (fun () -> raise EOG)
+      | Some first ->
+        let cur = ref first in (* current elemnt of the list *)
+        let stop = ref false in (* are we done yet? *)
+        (fun () ->
+          (if !stop then raise EOG);
+          let x = (!cur).content in
+          cur := (!cur).next;
+          (if !cur == first then stop := true); (* EOG, we made a full cycle *)
+          x)
+end
+
+(** Store content of the generator in an enum *)
+let persistent gen =
+  let l = MList.create () in
+  (try
+    while true do MList.push_back l (gen ()); done
+  with EOG ->
+    ());
+  (* done recursing through the generator *)
+  MList.to_enum l
+
+let tee ?(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 cur = ref 0 in
+    (* get next element for the i-th queue *)
+    let rec next i =
+      let q = qs.(i) in
+      if Queue.is_empty q
+        then update_to_i i  (* consume generator *)
+        else Queue.pop q
+    (* consume [gen] until some element for [i]-th generator is
+       available. It raises EOG if [gen] is exhausted before *)
+    and update_to_i i =
+      let x = gen () in
+      let j = !cur in
+      cur := (j+1) mod n;  (* move cursor to next generator *)
+      let q = qs.(j) in
+      if j = i
+        then begin
+          assert (Queue.is_empty q);
+          x  (* return the element *)
+        end else begin
+          Queue.push x q;
+          update_to_i i  (* continue consuming [gen] *)
+        end
+    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 () ->
+    let gen_a = a () in
+    let gen_b = b () in
+    let left = ref true in  (* left or right? *)
+    fun () ->
+      if !left
+        then (left := false; gen_a ())
+        else (left := true; gen_b ())
+
+(** Put [x] between elements of [enum] *)
+let intersperse x enum =
+  fun () ->
+    let next_elem = ref None in
+    let gen = enum () in
+    (* must see whether the gen is empty (first element must be from enum) *)
+    try
+      next_elem := Some (gen ());
+      (* get next element *)
+      let rec next () =
+        match !next_elem with
+        | None -> next_elem := Some (gen ()); x  (* yield x, gen is not exhausted *)
+        | Some y -> next_elem := None; y (* yield element of gen *)
+      in next
+    with EOG ->
+      fun () -> raise EOG
+
+(** Cartesian product *)
+let product a b =
+  fun () ->
+    (* [a] is the outer relation *)
+    let gen_a = a () in
+    try
+      (* current element of [a] *)
+      let cur_a = ref (gen_a ()) in
+      let gen_b = ref (b ()) in
+      let rec next () =
+        try !cur_a, !gen_b ()
+        with EOG ->
+          (* gen_b exhausted, get next elem of [a] *)
+          cur_a := gen_a ();
+          gen_b := b ();
+          next ()
+      in
+      next
+    with EOG ->
+      raise EOG  (* [a] is empty *)
+
+let permutations enum =
+  failwith "not implemented" (* TODO *)
+
+let combinations n enum =
+  assert (n >= 0);
+  failwith "not implemented" (* TODO *)
+
+(** {2 Basic conversion functions} *)
+
+let to_list enum =
+  let rec fold gen =
+    try
+      let x = gen () in
+      x :: fold gen
+    with EOG -> []
+  in fold (enum ())
     
 let of_list l =
   fun () ->
@@ -159,22 +431,8 @@ let of_list l =
       | [] -> raise EOG
       | x::l' -> l := l'; x
 
-let to_list enum =
-  let rec fold gen =
-    try
-      let x = gen () in
-      x :: fold gen
-    with EOG -> []
-  in fold (enum ())
-
 let to_rev_list enum =
-  let rec fold acc gen =
-    let acc', stop =
-      try let x = gen () in x :: acc, false
-      with EOG -> acc, true
-    in if stop then acc' else fold acc' gen
-  in
-  fold [] (enum ())
+  fold (fun acc x -> x :: acc) [] enum
 
 let int_range i j =
   fun () ->
diff --git a/enum.mli b/enum.mli
index 5abc9f30..c446b916 100644
--- a/enum.mli
+++ b/enum.mli
@@ -37,11 +37,7 @@ and 'a generator = unit -> 'a
   (** A generator may be called several times, yielding the next value
       each time. It raises EOG when it reaches the end. *)
 
-val empty : 'a t
-  (** Enmpty enum *)
-
-val singleton : 'a -> 'a t
-  (** One-element enum *)
+(** {2 Generator functions} *)
 
 val start : 'a t -> 'a generator
   (** Create a new generator *)
@@ -52,6 +48,22 @@ val next : 'a generator -> 'a
 val junk : 'a generator -> unit
   (** Drop element *)
 
+(** {2 Basic constructors} *)
+
+val empty : 'a t
+  (** Enmpty enum *)
+
+val singleton : 'a -> 'a t
+  (** One-element enum *)
+
+val repeat : 'a -> 'a t
+  (** Repeat same element endlessly *)
+
+val iterate : 'a -> ('a -> 'a) -> 'a t
+  (** [iterate x f] is [[x; f x; f (f x); f (f (f x)); ...]] *)
+
+(** {2 Basic combinators} *)
+
 val is_empty : _ t -> bool
   (** Check whether the enum is empty *)
 
@@ -70,6 +82,9 @@ val map : ('a -> 'b) -> 'a t -> 'b t
 val append : 'a t -> 'a t -> 'a t
   (** Append the two enums *)
 
+val cycle : 'a t -> 'a t
+  (** Cycle through the enum, endlessly. The enum must not be empty. *)
+
 val flatten : 'a t t -> 'a t
   (** Flatten the enum of enum *)
 
@@ -82,6 +97,60 @@ val take : int -> 'a t -> 'a t
 val drop : int -> 'a t -> 'a t
   (** Drop n elements *)
 
+val filter : ('a -> bool) -> 'a t -> 'a t
+  (** Filter out elements that do not satisfy the predicate. The outer
+      enum must be finite. *)
+
+val takeWhile : ('a -> bool) -> 'a t -> 'a t
+  (** Take elements while they satisfy the predicate *)
+
+val dropWhile : ('a -> bool) -> 'a t -> 'a t
+  (** Drop elements while they satisfy the predicate *)
+
+val filterMap : ('a -> 'b option) -> 'a t -> 'b t
+  (** Maps some elements to 'b, drop the other ones *)
+
+val zipWith : ('a -> 'b -> 'c) -> 'a t -> 'b t -> 'c t
+  (** Combine common part of the enums (stops when one is exhausted) *)
+
+val zip : 'a t -> 'b t -> ('a * 'b) t
+  (** Zip together the common part of the enums *)
+
+val zipIndex : 'a t -> (int * 'a) t
+  (** Zip elements with their index in the enum *)
+
+(** {2 Complex combinators} *)
+
+val round_robin : 'a t t -> 'a t
+  (** Pick elements fairly in each sub-enum *)
+
+val persistent : 'a generator -> 'a t
+  (** Store content of the generator in memory, to be able to iterate on it
+      several times later *)
+
+val tee : ?n:int -> 'a t -> 'a generator t
+  (** Split the enum into [n] generators in a fair way. Elements with
+      [index = k mod n] with go to the k-th enum. [n] defaults value
+      is 2. *)
+
+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. *)
+
+val intersperse : 'a -> 'a t -> 'a t
+  (** Put the separator element between all elements of the given enum *)
+
+val product : 'a t -> 'b t -> ('a * 'b) t
+  (** Cartesian product *)
+
+val permutations : 'a t -> 'a t t
+  (** Permutations of the enum *)
+
+val combinations : int -> 'a t -> 'a t t
+  (** Combinations of given length *)
+
+(** {2 Basic conversion functions} *)
+
 val of_list : 'a list -> 'a t
   (** Enumerate the list *)