diff --git a/src/adaptors/mod.rs b/src/adaptors/mod.rs index 3bc419945..b69eb1335 100644 --- a/src/adaptors/mod.rs +++ b/src/adaptors/mod.rs @@ -233,6 +233,28 @@ impl Iterator for PutBack 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.iter.last().or(self.top) + } + + fn nth(&mut self, n: usize) -> Option { + 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(&mut self, mut f: G) -> bool where G: FnMut(Self::Item) -> bool { diff --git a/src/merge_join.rs b/src/merge_join.rs index 46c5d0c84..41832798f 100644 --- a/src/merge_join.rs +++ b/src/merge_join.rs @@ -32,10 +32,10 @@ pub struct MergeJoinBy { } impl Clone for MergeJoinBy - where I: Clone + Iterator, - I::Item: Clone, - J: Clone + Iterator, - J::Item: Clone, + where I: Iterator, + J: Iterator, + PutBack>: Clone, + PutBack>: Clone, F: Clone, { clone_fields!(left, right, cmp_fn); @@ -94,4 +94,74 @@ impl Iterator for MergeJoinBy (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 { + 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 { + 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), + }, + } + } + } } diff --git a/tests/specializations.rs b/tests/specializations.rs new file mode 100644 index 000000000..87e6b2b26 --- /dev/null +++ b/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); +impl Iterator for Unspecialized +where + I: Iterator, +{ + type Item = I::Item; + + #[inline(always)] + fn next(&mut self) -> Option { + self.0.next() + } + + #[inline(always)] + fn size_hint(&self) -> (usize, Option) { + self.0.size_hint() + } +} + +fn check_specialized<'a, V, IterItem, Iter, F>(iterator: &Iter, mapper: F) +where + V: Eq + Debug, + IterItem: 'a, + Iter: Iterator + Clone + 'a, + F: Fn(Box + '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, +) where + IterItem: 'a + Eq + Debug, + Iter: Iterator + 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<::Output, Output = ::Output> + + Clone, + ::Output: + BitXor::Output> + Eq + Debug + Clone, + Iter: Iterator + 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, known_expected_size: Option) { + { + // 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) -> () { + put_back_test(test_vec, None) + } +} + +fn merge_join_by_test(i1: Vec, i2: Vec, known_expected_size: Option) { + 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 { + eob.left().unwrap_or(0) + } + fn eob_right_z(eob: EitherOrBoth) -> usize { + eob.left().unwrap_or(0) + } + fn eob_both_z(eob: EitherOrBoth) -> 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, i2: Vec) -> () { + merge_join_by_test(i1, i2, None) + } +}