From 13fd7da1429a95e02a758cfeff17d695356a1ed3 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 21 Mar 2013 16:48:46 +0100
Subject: [PATCH] added a lot of operations on Enum, including sort, uniq,
 for_all, exists, mem, min, max, group, unzip, unfold, scan

---
 enum.ml  | 156 +++++++++++++++++++++++++++++++++++++++++++++++++++++++
 enum.mli |  57 ++++++++++++++++++++
 2 files changed, 213 insertions(+)

diff --git a/enum.ml b/enum.ml
index 96fdd31f..638d7f82 100644
--- a/enum.ml
+++ b/enum.ml
@@ -116,6 +116,17 @@ let iterate x f =
       acc := f cur;
       cur
 
+(** Dual of {!fold}, with a deconstructing operation *)
+let unfold f acc =
+  fun () ->
+    let acc = ref acc in
+    fun () ->
+      match f !acc with
+      | None -> raise EOG
+      | Some (x, acc') ->
+        acc := acc';
+        x
+
 (** {2 Basic combinators} *)
 
 let is_empty enum =
@@ -125,9 +136,37 @@ let is_empty enum =
 let fold f acc enum =
   Gen.fold f acc (enum ())
 
+let fold2 f acc e1 e2 =
+  let acc = ref acc in
+  let gen1 = e1 () and gen2 = e2 () in
+  (try
+    while true do acc := f !acc (gen1 ()) (gen2 ()) done
+  with EOG -> ());
+  !acc
+
+(** Successive values of the accumulator *)
+let scan f acc e =
+  fun () ->
+    let acc = ref acc in
+    let first = ref true in
+    let gen = e () in
+    fun () ->
+      if !first
+        then (first := false; !acc)
+        else begin
+          acc := f !acc (gen ());
+          !acc
+        end
+
 let iter f enum =
   Gen.iter f (enum ())
 
+let iter2 f e1 e2 =
+  let gen1 = e1 () and gen2 = e2 () in
+  try
+    while true do f (gen1 ()) (gen2 ()) done;
+  with EOG -> ()
+
 let length enum =
   Gen.length (enum ())
               
@@ -192,6 +231,13 @@ let flatMap f enum =
         next ()  (* try again, with [gen = f x] *)
     in next
 
+let mem ?(eq=(=)) x enum =
+  try
+    iter (fun y -> if eq x y then raise Exit) enum;
+    false
+  with Exit -> 
+    true
+
 let take n enum =
   assert (n >= 0);
   fun () ->
@@ -212,6 +258,16 @@ let drop n enum =
         else gen ()
     in next
 
+let nth n enum =
+  assert (n>=0);
+  let gen = enum () in
+  let rec iter i =
+    let x = gen () in
+    if n = i then x else iter (i+1)
+  in
+  try iter 0
+  with EOG -> raise Not_found
+
 let filter p enum =
   fun () ->
     let gen = enum () in
@@ -276,6 +332,52 @@ let zipIndex enum =
       incr r;
       n, x
 
+let unzip e =
+  map fst e, map snd e
+
+(** [partition p l] returns the elements that satisfy [p],
+    and the elements that do not satisfy [p] *)
+let partition p enum =
+  filter p enum, filter (fun x -> not (p x)) enum
+
+let for_all p enum =
+  try
+    iter (fun x -> if not (p x) then raise Exit) enum;
+    true
+  with Exit ->
+    false
+
+let exists p enum =
+  try
+    iter (fun x -> if p x then raise Exit) enum;
+    false
+  with Exit ->
+    true
+
+let for_all2 p e1 e2 =
+  try
+    iter2 (fun x y -> if not (p x y) then raise Exit) e1 e2;
+    true
+  with Exit ->
+    false
+
+let exists2 p e1 e2 =
+  try
+    iter2 (fun x y -> if p x y then raise Exit) e1 e2;
+    false
+  with Exit ->
+    true
+
+let min ?(lt=fun x y -> x < y) enum =
+  let gen = enum () in
+  let first = try gen () with EOG -> raise Not_found in
+  Gen.fold (fun min x -> if lt x min then x else min) first gen
+
+let max ?(lt=fun x y -> x < y) enum =
+  let gen = enum () in
+  let first = try gen () with EOG -> raise Not_found in
+  Gen.fold (fun max x -> if lt max x then x else max) first gen
+
 (** {2 Complex combinators} *)
 
 (** Pick elements fairly in each sub-enum *)
@@ -569,6 +671,59 @@ let product a b =
       in
       next
 
+(** Group equal consecutive elements together. *)
+let group ?(eq=(=)) enum =
+  fun () ->
+    let gen = enum () in
+    try
+      let cur = ref [gen ()] in
+      let rec next () =
+        (* try to get an element *)
+        let next_x =
+          if !cur = []
+            then None
+            else try Some (gen ()) with EOG -> None in
+        match next_x, !cur with
+        | None, [] -> raise EOG
+        | None, l ->
+          cur := [];
+          l
+        | Some x, y::_ when eq x y ->
+          cur := x::!cur;
+          next ()  (* same group *)
+        | Some x, l ->
+          cur := [x];
+          l
+      in next
+    with EOG ->
+      fun () -> raise EOG
+
+let uniq ?(eq=(=)) enum =
+  fun () ->
+    let gen = enum () in
+    let prev = ref (Obj.magic 0) in
+    let first = ref true in
+    let rec next () =
+      let x = gen () in
+      if !first then (first := false; prev := x; x)
+      else if eq x !prev then next ()
+      else (prev := x; x)
+    in next
+
+let sort ?(cmp=compare) enum =
+  fun () ->
+    (* build heap *)
+    let h = Heap.empty ~cmp in
+    iter (Heap.insert h) enum;
+    fun () ->
+      if Heap.is_empty h
+        then raise EOG
+        else Heap.pop h
+
+let sort_uniq ?(cmp=compare) enum =
+  uniq ~eq:(fun x y -> cmp x y = 0) (sort ~cmp enum)
+
+(*
 let permutations enum =
   failwith "not implemented" (* TODO *)
 
@@ -578,6 +733,7 @@ let combinations n enum =
 
 let powerSet enum =
   failwith "not implemented"
+*)
 
 (** {2 Basic conversion functions} *)
 
diff --git a/enum.mli b/enum.mli
index 70b555d2..f20f1211 100644
--- a/enum.mli
+++ b/enum.mli
@@ -83,6 +83,9 @@ val repeat : 'a -> 'a t
 val iterate : 'a -> ('a -> 'a) -> 'a t
   (** [iterate x f] is [[x; f x; f (f x); f (f (f x)); ...]] *)
 
+val unfold : ('b -> ('a * 'b) option) -> 'b -> 'a t
+  (** Dual of {!fold}, with a deconstructing operation *)
+
 (** {2 Basic combinators} *)
 
 val is_empty : _ t -> bool
@@ -91,9 +94,18 @@ val is_empty : _ t -> bool
 val fold : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b
   (** Fold on the generator *)
 
+val fold2 : ('c -> 'a -> 'b -> 'c) -> 'c -> 'a t -> 'b t -> 'c
+  (** Fold on the two enums in parallel *)
+
+val scan : ('b -> 'a -> 'b) -> 'b -> 'a t -> 'b t
+  (** Successive values of the accumulator *)
+
 val iter : ('a -> unit) -> 'a t -> unit
   (** Iterate on the enum *)
 
+val iter2 : ('a -> 'b -> unit) -> 'a t -> 'b t -> unit
+  (** Iterate on the two sequences *)
+
 val length : _ t -> int
   (** Length of an enum (linear time) *)
 
@@ -112,12 +124,18 @@ val flatten : 'a t t -> 'a t
 val flatMap : ('a -> 'b t) -> 'a t -> 'b t
   (** Monadic bind *)
 
+val mem : ?eq:('a -> 'a -> bool) -> 'a -> 'a t -> bool
+  (** Is the given element, member of the enum? *)
+
 val take : int -> 'a t -> 'a t
   (** Take at most n elements *)
 
 val drop : int -> 'a t -> 'a t
   (** Drop n elements *)
 
+val nth : int -> 'a t -> 'a
+  (** n-th element, or Not_found *)
+
 val filter : ('a -> bool) -> 'a t -> 'a t
   (** Filter out elements that do not satisfy the predicate.  *)
 
@@ -139,6 +157,29 @@ val zip : 'a t -> 'b t -> ('a * 'b) t
 val zipIndex : 'a t -> (int * 'a) t
   (** Zip elements with their index in the enum *)
 
+val unzip : ('a * 'b) t -> 'a t * 'b t
+  (** Unzip into two sequences *)
+
+val partition : ('a -> bool) -> 'a t -> 'a t * 'a t
+  (** [partition p l] returns the elements that satisfy [p],
+      and the elements that do not satisfy [p] *)
+
+val for_all : ('a -> bool) -> 'a t -> bool
+  (** Predicate true for all elements? *)
+
+val exists : ('a -> bool) -> 'a t -> bool
+  (** Predicate true for at least one element? *)
+
+val for_all2 : ('a -> 'b -> bool) -> 'a t -> 'b t -> bool
+
+val exists2 : ('a -> 'b -> bool) -> 'a t -> 'b t -> bool
+
+val min : ?lt:('a -> 'a -> bool) -> 'a t -> 'a
+  (** Minimum element *)
+
+val max : ?lt:('a -> 'a -> bool) -> 'a t -> 'a
+  (** Maximum element *)
+
 (** {2 Complex combinators} *)
 
 val merge : 'a t t -> 'a t
@@ -182,6 +223,20 @@ val intersperse : 'a -> 'a t -> 'a t
 val product : 'a t -> 'b t -> ('a * 'b) t
   (** Cartesian product *)
 
+val group : ?eq:('a -> 'a -> bool) -> 'a t -> 'a list t
+  (** Group equal consecutive elements together. *)
+
+val uniq : ?eq:('a -> 'a -> bool) -> 'a t -> 'a t
+  (** Remove consecutive duplicate elements. Basically this is
+      like [fun e -> map List.hd (group e)]. *)
+
+val sort : ?cmp:('a -> 'a -> int) -> 'a t -> 'a t
+  (** Sort according to the given comparison function *)
+
+val sort_uniq : ?cmp:('a -> 'a -> int) -> 'a t -> 'a t
+  (** Sort and remove duplicates *)
+
+(* TODO later
 val permutations : 'a t -> 'a t t
   (** Permutations of the enum. Each permutation becomes unavailable once
       the next one is produced. *)
@@ -191,6 +246,7 @@ val combinations : int -> 'a t -> 'a t t
 
 val powerSet : 'a t -> 'a t t
   (** All subsets of the enum (in no particular order) *)
+*)
 
 (** {2 Basic conversion functions} *)
 
@@ -215,3 +271,4 @@ module Infix : sig
   val (--) : int -> int -> int t
   val (|>) : 'a -> ('a -> 'b) -> 'b
 end
+