mergesorted

src/Iterators.jl

@@ -17,6 +17,7 @@ export
     subsets,
     iterate,
     takenth,
+    mergesorted,
     @itr
 
 # iteratorsize is new in 0.5, declare it here for older versions. However,
@@ -714,6 +715,109 @@ start(it::Iterate) = it.seed
 next(it::Iterate, state) = (state, it.f(state))
 @compat done(it::Iterate, state) = (state==Union{})
 
+# PrefetchIterator: helper class for preliminary fetching one element in advance
+
+immutable PrefetchIterator
+    inner      # underlying iterator
+end
+
+immutable PrefetchState
+    hd         # current head of iterator (next to emit)
+    hd_prev    # previous head (just emitted)
+    st         # curent state of iterator (next to emit)
+    st_prev    # previous state (just emitted)
+end
+
+function head(state::PrefetchState)
+    return state.hd
+end
+
+function start(it::PrefetchIterator)
+    st0 = start(it.inner)
+    hd, st = next(it.inner, st0)
+    return PrefetchState(hd, hd, st, st0)
+end
+
+function next(it::PrefetchIterator, state::PrefetchState)
+    if done(it.inner, state.st)
+        # can't read any more elements, repeating current one forever
+        return state.hd, PrefetchState(state.hd, state.hd, state.st, state.st)
+    else
+        new_hd, new_st = next(it.inner, state.st)
+        return state.hd, PrefetchState(new_hd, state.hd, new_st, state.st)
+    end
+end
+
+function done(it::PrefetchIterator, state::PrefetchState)
+    return done(it.inner, state.st_prev)
+end
+
+
+# mergesorted: merge 2 sorted iterators
+
+type MergedIter
+    it1::PrefetchIterator   # iterator 1
+    it2::PrefetchIterator   # iterator 2
+    lt                      # custom `<` operator to compare elements
+                            # of 2 iterators
+end
+
+type MergedState
+    s1::PrefetchState
+    s2::PrefetchState
+end
+
+
+function start(it::MergedIter)
+    s1 = start(it.it1)
+    s2 = start(it.it2)
+    return MergedState(s1, s2)
+end
+
+function next(it::MergedIter, s::MergedState)
+    lt = it.lt
+    it1 = it.it1; it2 = it.it2; s1 = s.s1; s2 = s.s2
+    while head(s1) == head(s2) && !done(it2, s2)
+        # move second iterator till heads are different or iterator is done
+        _, s2 = next(it2, s2)
+    end
+    hd1 = head(s1); hd2 = head(s2)
+    if done(it1, s1)
+        _, next_s2 = next(it2, s2)
+        return hd2, MergedState(s1, next_s2)
+    elseif done(it2, s2)
+        _, next_s1 = next(it1, s1)
+        return hd1, MergedState(next_s1, s2)
+    elseif lt(hd1, hd2) # hd1 < hd2
+        _, next_s1 = next(it1, s1)
+        return hd1, MergedState(next_s1, s2)
+    else # hd2 < hd1
+        _, next_s2 = next(it2, s2)
+        return hd2, MergedState(s1, next_s2)
+    end
+end
+
+function done(it::MergedIter, s::MergedState)
+    return done(it.it1, s.s1) && done(it.it2, s.s2)
+end
+
+"""
+Merge 2 sorted iterators. Let `hd1` be a head of 1st iterator and `hd2` -
+a head of 2nd iterator. Then merged iterator works as follows:
+
+ * if hd1 < hd2, hd1 is emitted
+ * if hd1 > hd2, hd2 is emitted
+ * if hd1 == hd2, hd1 is emitted and hd2 is discarded (so no duplicates
+                  are produced)
+ * if one of iterators is done, tail of another is taken
+
+Comparison is done using `lt` option that defaults to `isless` function.
+"""
+function mergesorted(it1, it2; lt = isless)
+    return MergedIter(PrefetchIterator(it1), PrefetchIterator(it2), lt)
+end
+
+
 using Base.Meta
 
 ## @itr macro for auto-inlining in for loops

test/runtests.jl

@@ -224,6 +224,27 @@ test_groupby(
 @test collect(takenth(10:20, 1)) == collect(10:20)
 
 
+# mergesorted
+
+macro test_mergesorted(it1, it2, expected)
+    x = gensym()
+    w = :(mergesorted($it1, $it2))
+    quote
+	actual = Any[]
+	for $x in $w
+	    push!(actual, $x)
+	end
+	@test actual == $expected
+    end
+end
+
+
+@test_mergesorted 1:10 5:15 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+@test_mergesorted 5:15 1:10 [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+@test_mergesorted [:a, :c, :d, :f] [:b, :c, :d, :e] [:a, :b, :c, :d, :e, :f]
+
+
+
 ## @itr
 ## ====