diff --git a/enum.ml b/enum.ml index f481c607..b26d03c4 100644 --- a/enum.ml +++ b/enum.ml @@ -23,7 +23,7 @@ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. *) -(** {1 Consumable generators} *) +(** {1 Restartable generators} *) exception EOG (** End of Generation *) @@ -474,9 +474,31 @@ module Heap = struct m end +(** Binary sorted merge of two sorted sequences *) +let sorted_merge ?(cmp=compare) e1 e2 = + fun () -> + let gen1, gen2 = e1 (), e2 () in + let next1 () = try Some (gen1 ()) with EOG -> None in + let next2 () = try Some (gen2 ()) with EOG -> None in + let x1 = ref (next1 ()) in + let x2 = ref (next2 ()) in + fun () -> + match !x1, !x2 with + | None, None -> raise EOG + | Some y1, Some y2 -> + if cmp y1 y2 <= 0 + then (x1 := next1 (); y1) + else (x2 := next2 (); y2) + | Some y1, None -> + x1 := next1 (); + y1 + | None, Some y2 -> + x2 := next2 (); + y2 + (** Assuming subsequences are sorted in increasing order, merge them into an increasing sequence *) -let merge_sorted ?(cmp=compare) enum = +let sorted_merge_n ?(cmp=compare) enum = fun () -> (* make a heap of (value, generator) *) let cmp (v1,_) (v2,_) = cmp v1 v2 in diff --git a/enum.mli b/enum.mli index 04d166c5..9f4c9dfd 100644 --- a/enum.mli +++ b/enum.mli @@ -23,7 +23,7 @@ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. *) -(** {1 Consumable generators} *) +(** {1 Restartable generators} *) (** This structure is inspired from Ocaml Batteries' BatEnum.t. It features restartable generators. *) @@ -195,9 +195,11 @@ module Heap : sig val pop : 'a t -> 'a end -val merge_sorted : ?cmp:('a -> 'a -> int) -> 'a t t -> 'a t - (** Assuming subsequences are sorted in increasing order, merge them - into an increasing sequence *) +val sorted_merge : ?cmp:('a -> 'a -> int) -> 'a t -> 'a t -> 'a t + (** Merge two sorted sequences into a sorted sequence *) + +val sorted_merge_n : ?cmp:('a -> 'a -> int) -> 'a t t -> 'a t + (** Sorted merge of multiple sorted sequences *) val persistent : 'a generator -> 'a t (** Store content of the generator in memory, to be able to iterate on it