diff --git a/src/libcore/benches/iter.rs b/src/libcore/benches/iter.rs index b284d855c4515..6c597301ac204 100644 --- a/src/libcore/benches/iter.rs +++ b/src/libcore/benches/iter.rs @@ -281,3 +281,32 @@ bench_sums! { bench_take_while_chain_ref_sum, (0i64..1000000).chain(1000000..).take_while(|&x| x < 1111111) } + +// Checks whether Skip> is as fast as Zip, Skip>, from +// https://users.rust-lang.org/t/performance-difference-between-iterator-zip-and-skip-order/15743 +#[bench] +fn bench_zip_then_skip(b: &mut Bencher) { + let v: Vec<_> = (0..100_000).collect(); + let t: Vec<_> = (0..100_000).collect(); + + b.iter(|| { + let s = v.iter().zip(t.iter()).skip(10000) + .take_while(|t| *t.0 < 10100) + .map(|(a, b)| *a + *b) + .sum::(); + assert_eq!(s, 2009900); + }); +} +#[bench] +fn bench_skip_then_zip(b: &mut Bencher) { + let v: Vec<_> = (0..100_000).collect(); + let t: Vec<_> = (0..100_000).collect(); + + b.iter(|| { + let s = v.iter().skip(10000).zip(t.iter().skip(10000)) + .take_while(|t| *t.0 < 10100) + .map(|(a, b)| *a + *b) + .sum::(); + assert_eq!(s, 2009900); + }); +} diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs index 623cad754dd72..257d7d6caaaf8 100644 --- a/src/libcore/iter/mod.rs +++ b/src/libcore/iter/mod.rs @@ -1045,6 +1045,11 @@ impl Iterator for Zip where A: Iterator, B: Iterator fn size_hint(&self) -> (usize, Option) { ZipImpl::size_hint(self) } + + #[inline] + fn nth(&mut self, n: usize) -> Option { + ZipImpl::nth(self, n) + } } #[stable(feature = "rust1", since = "1.0.0")] @@ -1065,6 +1070,14 @@ trait ZipImpl { fn new(a: A, b: B) -> Self; fn next(&mut self) -> Option; fn size_hint(&self) -> (usize, Option); + fn nth(&mut self, n: usize) -> Option; + fn super_nth(&mut self, mut n: usize) -> Option { + while let Some(x) = self.next() { + if n == 0 { return Some(x) } + n -= 1; + } + None + } fn next_back(&mut self) -> Option where A: DoubleEndedIterator + ExactSizeIterator, B: DoubleEndedIterator + ExactSizeIterator; @@ -1094,6 +1107,11 @@ impl ZipImpl for Zip }) } + #[inline] + default fn nth(&mut self, n: usize) -> Option { + self.super_nth(n) + } + #[inline] default fn next_back(&mut self) -> Option<(A::Item, B::Item)> where A: DoubleEndedIterator + ExactSizeIterator, @@ -1174,6 +1192,24 @@ impl ZipImpl for Zip (len, Some(len)) } + #[inline] + fn nth(&mut self, n: usize) -> Option { + let delta = cmp::min(n, self.len - self.index); + let end = self.index + delta; + while self.index < end { + let i = self.index; + self.index += 1; + if A::may_have_side_effect() { + unsafe { self.a.get_unchecked(i); } + } + if B::may_have_side_effect() { + unsafe { self.b.get_unchecked(i); } + } + } + + self.super_nth(n - delta) + } + #[inline] fn next_back(&mut self) -> Option<(A::Item, B::Item)> where A: DoubleEndedIterator + ExactSizeIterator, diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs index edd75f7795ed7..a962efadd64e9 100644 --- a/src/libcore/tests/iter.rs +++ b/src/libcore/tests/iter.rs @@ -144,6 +144,43 @@ fn test_iterator_chain_find() { assert_eq!(iter.next(), None); } +#[test] +fn test_zip_nth() { + let xs = [0, 1, 2, 4, 5]; + let ys = [10, 11, 12]; + + let mut it = xs.iter().zip(&ys); + assert_eq!(it.nth(0), Some((&0, &10))); + assert_eq!(it.nth(1), Some((&2, &12))); + assert_eq!(it.nth(0), None); + + let mut it = xs.iter().zip(&ys); + assert_eq!(it.nth(3), None); + + let mut it = ys.iter().zip(&xs); + assert_eq!(it.nth(3), None); +} + +#[test] +fn test_zip_nth_side_effects() { + let mut a = Vec::new(); + let mut b = Vec::new(); + let value = [1, 2, 3, 4, 5, 6].iter().cloned() + .map(|n| { + a.push(n); + n * 10 + }) + .zip([2, 3, 4, 5, 6, 7, 8].iter().cloned().map(|n| { + b.push(n * 100); + n * 1000 + })) + .skip(1) + .nth(3); + assert_eq!(value, Some((50, 6000))); + assert_eq!(a, vec![1, 2, 3, 4, 5]); + assert_eq!(b, vec![200, 300, 400, 500, 600]); +} + #[test] fn test_iterator_step_by() { // Identity