Repository: bluss/rust-itertools
Pull request #27 (2015-04-16 19:05:51)

Add merge_by
Doesn't yet compile.

bluss: Great, this is a useful addition to itertools. Small comments in the patch. Please update the PR title & message & commit message since that's the information that's associated with the change when it ends up in git.

Patch:
diff --git a/src/adaptors.rs b/src/adaptors.rs
index d3f2527..e5bd40e 100644
--- a/src/adaptors.rs
+++ b/src/adaptors.rs
@@ -9,6 +9,7 @@ use std::mem;
 use std::num::One;
 #[cfg(feature = "unstable")]
 use std::ops::Add;
+use std::cmp::Ordering;
 use std::usize;
 use std::iter::{Fuse, Peekable};
 use super::Itertools;
@@ -465,49 +466,62 @@ impl<I> Iterator for Step<I>
 /// If both base iterators are sorted (ascending), the result is sorted.
 ///
 /// Iterator element type is **I::Item**.
-pub struct Merge<I, J> where
+pub struct Merge<I, J, F> where
     I: Iterator,
     J: Iterator<Item=I::Item>,
+    F: Fn(&I::Item, &I::Item) -> Ordering
 {
     a: Peekable<I>,
     b: Peekable<J>,
+    cmp: F,
 }
 
-impl<I, J> Merge<I, J> where
+impl<I, J, F> Merge<I, J, F> where
     I: Iterator,
     J: Iterator<Item=I::Item>,
+    F: Fn(&I::Item, &I::Item) -> Ordering
 {
     /// Create a **Merge** iterator.
-    pub fn new(a: I, b: J) -> Self
+    pub fn new(a: I, b: J, cmp: F) -> Self
     {
         Merge {
             a: a.peekable(),
             b: b.peekable(),
+            cmp: cmp,
         }
     }
 }
 
-impl<I, J> Clone for Merge<I, J> where
+impl<I, J, F> Clone for Merge<I, J, F> where
     I: Iterator,
     J: Iterator<Item=I::Item>,
     Peekable<I>: Clone,
     Peekable<J>: Clone,
+    F: Clone,
+    F: Fn(&I::Item, &I::Item) -> Ordering
 {
     fn clone(&self) -> Self {
-        clone_fields!(Merge, self, a, b)
+        clone_fields!(Merge, self, a, b, cmp)
     }
 }
 
-impl<I, J> Iterator for Merge<I, J> where
+impl<I, J, F> Iterator for Merge<I, J, F> where
     I: Iterator,
     I::Item: PartialOrd,
     J: Iterator<Item=I::Item>,
+    F: Fn(&I::Item, &I::Item) -> Ordering
 {
     type Item = I::Item;
 
     fn next(&mut self) -> Option<I::Item> {
         if match (self.a.peek(), self.b.peek()) {
-            (Some(a), Some(b)) => a <= b,
+            (Some(a), Some(b)) => {
+                match (self.cmp)(a, b) {
+                    Ordering::Less => true,
+                    Ordering::Equal => true,
+                    Ordering::Greater => false,
+                }
+            }
             (Some(_), None) => true,
             (None, Some(_)) => false,
             (None, None) => return None,
diff --git a/src/lib.rs b/src/lib.rs
index d6a907a..2b47e1e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -31,6 +31,7 @@
 //!
 
 use std::fmt::Write;
+use std::cmp::Ordering;
 
 pub use adaptors::{
     Interleave,
@@ -424,12 +425,45 @@ pub trait Itertools : Iterator {
     /// assert_eq!(it.next(), Some(4));
     /// assert_eq!(it.next(), Some(6));
     /// ```
-    fn merge<J>(self, other: J) -> Merge<Self, J> where
+    fn merge<J>(self, other: J)
+        -> Merge<Self, J, fn(&Self::Item, &Self::Item) -> Ordering> where
         Self: Sized,
         Self::Item: PartialOrd,
         J: Iterator<Item=Self::Item>,
     {
-        Merge::new(self, other)
+        fn wrapper<A: PartialOrd>(a: &A, b: &A) -> Ordering { 
+            a.partial_cmp(b).unwrap_or(Ordering::Less) 
+        };
+        self.merge_by(other, wrapper)
+    }
+
+    /// Return an iterator adaptor that merges the two base iterators in an order.
+    /// This is much like merge_by but allows for descending orders or sorting tuples.
+    /// This can be especially useful walking a BTreeMap structure.
+    ///
+    /// Iterator element type is **Self::Item**.
+    ///
+    /// ## Example
+    /// ```
+    /// use itertools::Itertools;
+    ///
+    /// let a = (0..10).step(2);
+    /// let b = (1..10).step(3);
+    /// let mut it = a.merge_by(b, |x, y|{x.cmp(&y)});
+    /// assert_eq!(it.next(), Some(0));
+    /// assert_eq!(it.next(), Some(1));
+    /// assert_eq!(it.next(), Some(2));
+    /// assert_eq!(it.next(), Some(4));
+    /// assert_eq!(it.next(), Some(4));
+    /// assert_eq!(it.next(), Some(6));
+    /// ```
+
+    fn merge_by<J, F>(self, other: J, cmp: F) -> Merge<Self, J, F> where
+        Self: Sized,
+        J: Iterator<Item=Self::Item>,
+        F: Fn(&Self::Item, &Self::Item) -> Ordering
+    {
+        Merge::new(self, other, cmp)
     }
 
     /// Return an iterator adaptor that iterates over the cartesian product of


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 7a5b180..7c1dba0 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -321,6 +321,29 @@ fn merge() {
     assert_iters_equal((0..10).step(2).merge((1..10).step(2)), (0..10));
 }
 
+#[test]
+fn merge_by() {
+    let odd : Vec<(u32, &str)> = vec![(1, "hello"), (3, "world"), (5, "!")];
+    let even = vec![(2, "foo"), (4, "bar"), (6, "baz")];
+    let expected = vec![(1, "hello"), (2, "foo"), (3, "world"), (4, "bar"), (5, "!"), (6, "baz")];
+    let results = odd.iter().merge_by(even.iter(), |a, b|{ a.0.cmp(&b.0)});
+    assert_iters_equal(results, expected.iter());
+}
+
+#[test]
+fn merge_by_btree() {
+    use std::collections::BTreeMap;
+    let mut bt1 = BTreeMap::new();
+    bt1.insert("hello", 1);
+    bt1.insert("world", 3);
+    let mut bt2 = BTreeMap::new();
+    bt2.insert("foo", 2);
+    bt2.insert("bar", 4);
+    let results = bt1.into_iter().merge_by(bt2.into_iter(), |a, b|{a.0.cmp(&b.0)});
+    let expected = vec![("bar", 4), ("foo", 2), ("hello", 1), ("world", 3)];
+    assert_iters_equal(results, expected.into_iter());
+}
+
 #[test]
 fn to_string_join() {
     let many = [1, 2, 3];