Rollup merge of #48635 - scottmcm:faster-zip-nth, r=kennytm

Fixes #47311.
r? @nrc

3 files changed, 102 insertions, 0 deletions

diff --git a/src/libcore/benches/iter.rs b/src/libcore/benches/iter.rs
index b284d855c45..6c597301ac2 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<Zip<A,B>> is as fast as Zip<Skip<A>, Skip<B>>, 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::<u64>();
+        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::<u64>();
+        assert_eq!(s, 2009900);
+    });
+}
diff --git a/src/libcore/iter/mod.rs b/src/libcore/iter/mod.rs
index 623cad754dd..257d7d6caaa 100644
--- a/src/libcore/iter/mod.rs
+++ b/src/libcore/iter/mod.rs
@@ -1045,6 +1045,11 @@ impl<A, B> Iterator for Zip<A, B> where A: Iterator, B: Iterator
     fn size_hint(&self) -> (usize, Option<usize>) {
         ZipImpl::size_hint(self)
     }
+
+    #[inline]
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        ZipImpl::nth(self, n)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
@@ -1065,6 +1070,14 @@ trait ZipImpl<A, B> {
     fn new(a: A, b: B) -> Self;
     fn next(&mut self) -> Option<Self::Item>;
     fn size_hint(&self) -> (usize, Option<usize>);
+    fn nth(&mut self, n: usize) -> Option<Self::Item>;
+    fn super_nth(&mut self, mut n: usize) -> Option<Self::Item> {
+        while let Some(x) = self.next() {
+            if n == 0 { return Some(x) }
+            n -= 1;
+        }
+        None
+    }
     fn next_back(&mut self) -> Option<Self::Item>
         where A: DoubleEndedIterator + ExactSizeIterator,
               B: DoubleEndedIterator + ExactSizeIterator;
@@ -1095,6 +1108,11 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
     }
 
     #[inline]
+    default fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.super_nth(n)
+    }
+
+    #[inline]
     default fn next_back(&mut self) -> Option<(A::Item, B::Item)>
         where A: DoubleEndedIterator + ExactSizeIterator,
               B: DoubleEndedIterator + ExactSizeIterator
@@ -1175,6 +1193,24 @@ impl<A, B> ZipImpl<A, B> for Zip<A, B>
     }
 
     #[inline]
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        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,
               B: DoubleEndedIterator + ExactSizeIterator
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index edd75f7795e..a962efadd64 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -145,6 +145,43 @@ fn test_iterator_chain_find() {
 }
 
 #[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
     let mut it = (0..).step_by(1).take(3);