Implement missing specializations on the PutBack adaptor and on the MergeJoinBy Iterator

src/adaptors/mod.rs

@@ -243,6 +243,28 @@ impl<I> Iterator for PutBack<I>
         size_hint::add_scalar(self.iter.size_hint(), self.top.is_some() as usize)
     }
 
+    fn count(self) -> usize {
+        self.iter.count() + (self.top.is_some() as usize)
+    }
+
+    fn last(self) -> Option<Self::Item> {
+        self.iter.last().or(self.top)
+    }
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        match self.top {
+            None => self.iter.nth(n),
+            ref mut some => {
+                if n == 0 {
+                    some.take()
+                } else {
+                    *some = None;
+                    self.iter.nth(n - 1)
+                }
+            }
+        }
+    }
+
     fn all<G>(&mut self, mut f: G) -> bool
         where G: FnMut(Self::Item) -> bool
     {

src/merge_join.rs

@@ -32,10 +32,10 @@ pub struct MergeJoinBy<I: Iterator, J: Iterator, F> {
 }
 
 impl<I, J, F> Clone for MergeJoinBy<I, J, F>
-    where I: Clone + Iterator,
-          I::Item: Clone,
-          J: Clone + Iterator,
-          J::Item: Clone,
+    where I: Iterator,
+          J: Iterator,
+          PutBack<Fuse<I>>: Clone,
+          PutBack<Fuse<J>>: Clone,
           F: Clone,
 {
     fn clone(&self) -> Self {
@@ -100,4 +100,74 @@ impl<I, J, F> Iterator for MergeJoinBy<I, J, F>
 
         (lower, upper)
     }
+
+    fn count(mut self) -> usize {
+        let mut count = 0;
+        loop {
+            match (self.left.next(), self.right.next()) {
+                (None, None) => break count,
+                (Some(_left), None) => break count + 1 + self.left.into_parts().1.count(),
+                (None, Some(_right)) => break count + 1 + self.right.into_parts().1.count(),
+                (Some(left), Some(right)) => {
+                    count += 1;
+                    match (self.cmp_fn)(&left, &right) {
+                        Ordering::Equal => {}
+                        Ordering::Less => self.right.put_back(right),
+                        Ordering::Greater => self.left.put_back(left),
+                    }
+                }
+            }
+        }
+    }
+
+    fn last(mut self) -> Option<Self::Item> {
+        let mut previous_element = None;
+        loop {
+            match (self.left.next(), self.right.next()) {
+                (None, None) => break previous_element,
+                (Some(left), None) => {
+                    break Some(EitherOrBoth::Left(
+                        self.left.into_parts().1.last().unwrap_or(left),
+                    ))
+                }
+                (None, Some(right)) => {
+                    break Some(EitherOrBoth::Right(
+                        self.right.into_parts().1.last().unwrap_or(right),
+                    ))
+                }
+                (Some(left), Some(right)) => {
+                    previous_element = match (self.cmp_fn)(&left, &right) {
+                        Ordering::Equal => Some(EitherOrBoth::Both(left, right)),
+                        Ordering::Less => {
+                            self.right.put_back(right);
+                            Some(EitherOrBoth::Left(left))
+                        }
+                        Ordering::Greater => {
+                            self.left.put_back(left);
+                            Some(EitherOrBoth::Right(right))
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
+        loop {
+            if n == 0 {
+                break self.next();
+            }
+            n -= 1;
+            match (self.left.next(), self.right.next()) {
+                (None, None) => break None,
+                (Some(_left), None) => break self.left.nth(n).map(EitherOrBoth::Left),
+                (None, Some(_right)) => break self.right.nth(n).map(EitherOrBoth::Right),
+                (Some(left), Some(right)) => match (self.cmp_fn)(&left, &right) {
+                    Ordering::Equal => {}
+                    Ordering::Less => self.right.put_back(right),
+                    Ordering::Greater => self.left.put_back(left),
+                },
+            }
+        }
+    }
 }

tests/specializations.rs

@@ -0,0 +1,154 @@
+extern crate itertools;
+
+#[macro_use]
+extern crate quickcheck;
+
+use itertools::{EitherOrBoth, Itertools};
+
+use std::fmt::Debug;
+use std::ops::BitXor;
+
+struct Unspecialized<I>(I);
+impl<I> Iterator for Unspecialized<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    #[inline(always)]
+    fn next(&mut self) -> Option<I::Item> {
+        self.0.next()
+    }
+
+    #[inline(always)]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.0.size_hint()
+    }
+}
+
+fn check_specialized<'a, V, IterItem, Iter, F>(iterator: &Iter, mapper: F)
+where
+    V: Eq + Debug,
+    IterItem: 'a,
+    Iter: Iterator<Item = IterItem> + Clone + 'a,
+    F: Fn(Box<Iterator<Item = IterItem> + 'a>) -> V,
+{
+    assert_eq!(
+        mapper(Box::new(Unspecialized(iterator.clone()))),
+        mapper(Box::new(iterator.clone()))
+    )
+}
+
+fn check_specialized_count_last_nth_sizeh<'a, IterItem, Iter>(
+    it: &Iter,
+    known_expected_size: Option<usize>,
+) where
+    IterItem: 'a + Eq + Debug,
+    Iter: Iterator<Item = IterItem> + Clone + 'a,
+{
+    let size = it.clone().count();
+    if let Some(expected_size) = known_expected_size {
+        assert_eq!(size, expected_size);
+    }
+    check_specialized(it, |i| i.count());
+    check_specialized(it, |i| i.last());
+    for n in 0..size + 2 {
+        check_specialized(it, |mut i| i.nth(n));
+    }
+    let mut it_sh = it.clone();
+    for n in 0..size + 2 {
+        let len = it_sh.clone().count();
+        let (min, max) = it_sh.size_hint();
+        assert_eq!((size - n.min(size)), len);
+        assert!(min <= len);
+        if let Some(max) = max {
+            assert!(len <= max);
+        }
+        it_sh.next();
+    }
+}
+
+fn check_specialized_fold_xor<'a, IterItem, Iter>(it: &Iter)
+where
+    IterItem: 'a
+        + BitXor
+        + Eq
+        + Debug
+        + BitXor<<IterItem as BitXor>::Output, Output = <IterItem as BitXor>::Output>
+        + Clone,
+    <IterItem as BitXor>::Output:
+        BitXor<Output = <IterItem as BitXor>::Output> + Eq + Debug + Clone,
+    Iter: Iterator<Item = IterItem> + Clone + 'a,
+{
+    check_specialized(it, |mut i| {
+        let first = i.next().map(|f| f.clone() ^ (f.clone() ^ f));
+        i.fold(first, |acc, v: IterItem| acc.map(move |a| v ^ a))
+    });
+}
+
+fn put_back_test(test_vec: Vec<i32>, known_expected_size: Option<usize>) {
+    {
+        // Lexical lifetimes support
+        let pb = itertools::put_back(test_vec.iter());
+        check_specialized_count_last_nth_sizeh(&pb, known_expected_size);
+        check_specialized_fold_xor(&pb);
+    }
+
+    let mut pb = itertools::put_back(test_vec.into_iter());
+    pb.put_back(1);
+    check_specialized_count_last_nth_sizeh(&pb, known_expected_size.map(|x| x + 1));
+    check_specialized_fold_xor(&pb)
+}
+
+#[test]
+fn put_back() {
+    put_back_test(vec![7, 4, 1], Some(3));
+}
+
+quickcheck! {
+    fn put_back_qc(test_vec: Vec<i32>) -> () {
+        put_back_test(test_vec, None)
+    }
+}
+
+fn merge_join_by_test(i1: Vec<usize>, i2: Vec<usize>, known_expected_size: Option<usize>) {
+    let i1 = i1.into_iter();
+    let i2 = i2.into_iter();
+    let mjb = i1.clone().merge_join_by(i2.clone(), std::cmp::Ord::cmp);
+    check_specialized_count_last_nth_sizeh(&mjb, known_expected_size);
+    // Rust 1.24 compatibility:
+    fn eob_left_z(eob: EitherOrBoth<usize, usize>) -> usize {
+        eob.left().unwrap_or(0)
+    }
+    fn eob_right_z(eob: EitherOrBoth<usize, usize>) -> usize {
+        eob.left().unwrap_or(0)
+    }
+    fn eob_both_z(eob: EitherOrBoth<usize, usize>) -> usize {
+        let (a, b) = eob.both().unwrap_or((0, 0));
+        assert_eq!(a, b);
+        a
+    }
+    check_specialized_fold_xor(&mjb.clone().map(eob_left_z));
+    check_specialized_fold_xor(&mjb.clone().map(eob_right_z));
+    check_specialized_fold_xor(&mjb.clone().map(eob_both_z));
+
+    // And the other way around
+    let mjb = i2.merge_join_by(i1, std::cmp::Ord::cmp);
+    check_specialized_count_last_nth_sizeh(&mjb, known_expected_size);
+    check_specialized_fold_xor(&mjb.clone().map(eob_left_z));
+    check_specialized_fold_xor(&mjb.clone().map(eob_right_z));
+    check_specialized_fold_xor(&mjb.clone().map(eob_both_z));
+}
+
+#[test]
+fn merge_join_by() {
+    let i1 = vec![1, 3, 5, 7, 8, 9];
+    let i2 = vec![0, 3, 4, 5];
+    merge_join_by_test(i1, i2, Some(8));
+}
+
+quickcheck! {
+    fn merge_join_by_qc(i1: Vec<usize>, i2: Vec<usize>) -> () {
+        merge_join_by_test(i1, i2, None)
+    }
+}