Implement nth_back for Zip iterator

library/core/src/iter/adapters/zip.rs

@@ -32,6 +32,19 @@ impl<A: Iterator, B: Iterator> Zip<A, B> {
         }
         None
     }
+    fn super_nth_back(&mut self, mut n: usize) -> Option<(A::Item, B::Item)>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator,
+    {
+        while let Some(x) = <Self as DoubleEndedIterator>::next_back(self) {
+            if n == 0 {
+                return Some(x);
+            }
+            n -= 1;
+        }
+        None
+    }
 }
 
 /// Converts the arguments to iterators and zips them.
@@ -124,6 +137,11 @@ where
     fn next_back(&mut self) -> Option<(A::Item, B::Item)> {
         ZipImpl::next_back(self)
     }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<(A::Item, B::Item)> {
+        ZipImpl::nth_back(self, n)
+    }
 }
 
 // Zip specialization trait
@@ -135,6 +153,10 @@ trait ZipImpl<A, B> {
     fn size_hint(&self) -> (usize, Option<usize>);
     fn nth(&mut self, n: usize) -> Option<Self::Item>;
     fn next_back(&mut self) -> Option<Self::Item>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator;
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item>
     where
         A: DoubleEndedIterator + ExactSizeIterator,
         B: DoubleEndedIterator + ExactSizeIterator;
@@ -202,6 +224,15 @@ macro_rules! zip_impl_general_defaults {
                 _ => unreachable!(),
             }
         }
+
+        #[inline]
+        default fn nth_back(&mut self, n: usize) -> Option<(A::Item, B::Item)>
+        where
+            A: DoubleEndedIterator + ExactSizeIterator,
+            B: DoubleEndedIterator + ExactSizeIterator,
+        {
+            self.super_nth_back(n)
+        }
     };
 }
 
@@ -403,6 +434,80 @@ where
             None
         }
     }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<(A::Item, B::Item)>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator,
+    {
+        // No effects when the iterator is exhausted.
+        if self.index >= self.len {
+            return None;
+        }
+
+        let old_len = self.len;
+        let remaining = old_len - self.index;
+        // How many elements to skip from the back (capped at remaining).
+        let delta = cmp::min(n, remaining);
+        // The new back boundary after skipping `delta` elements.
+        let new_len = old_len - delta;
+
+        // Adjust a, b to equal length if needed on first backward iteration.
+        if A::MAY_HAVE_SIDE_EFFECT || B::MAY_HAVE_SIDE_EFFECT {
+            let sz_a = self.a.size();
+            let sz_b = self.b.size();
+            // This condition can and must only be true on the first backward
+            // iteration call, otherwise we would break the restriction on calls
+            // to `self.next_back()` after calling `get_unchecked()`.
+            if sz_a != sz_b && (old_len == sz_a || old_len == sz_b) {
+                if A::MAY_HAVE_SIDE_EFFECT && sz_a > old_len {
+                    for _ in 0..sz_a - old_len {
+                        self.a.next_back();
+                    }
+                }
+                if B::MAY_HAVE_SIDE_EFFECT && sz_b > old_len {
+                    for _ in 0..sz_b - old_len {
+                        self.b.next_back();
+                    }
+                }
+            }
+
+            // Execute side effects for the `delta` skipped elements from the back.
+            // Go from high to low indices since we're iterating backwards.
+            for i in (new_len..old_len).rev() {
+                // Update len before each access for panic safety, so that
+                // the same index won't be accessed twice, as required by
+                // TrustedRandomAccess.
+                self.len = i;
+                if A::MAY_HAVE_SIDE_EFFECT {
+                    // SAFETY: `i` < `old_len` <= `self.a.size()` and `self.b.size()`
+                    unsafe {
+                        self.a.__iterator_get_unchecked(i);
+                    }
+                }
+                if B::MAY_HAVE_SIDE_EFFECT {
+                    // SAFETY: same as above.
+                    unsafe {
+                        self.b.__iterator_get_unchecked(i);
+                    }
+                }
+            }
+        }
+
+        // If we couldn't skip enough, the iterator is exhausted.
+        if n >= remaining {
+            self.len = self.index;
+            return None;
+        }
+
+        // Get the target element at index `new_len - 1`.
+        let i = new_len - 1;
+        self.len = i;
+        // SAFETY: `i` is `old_len - n - 1`, which is >= `self.index` (since
+        // `n < remaining`) and < `old_len` <= `self.a.size()` and `self.b.size()`.
+        unsafe { Some((self.a.__iterator_get_unchecked(i), self.b.__iterator_get_unchecked(i))) }
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]

library/coretests/tests/iter/adapters/zip.rs

@@ -370,3 +370,45 @@ fn test_issue_82291() {
     zip.next();
     assert_eq!(called.get(), 1);
 }
+
+#[test]
+fn test_zip_nth_back() {
+    let xs = [0, 1, 2, 4, 5];
+    let ys = [10, 11, 12];
+
+    let mut it = xs.iter().zip(&ys);
+    assert_eq!(it.nth_back(0), Some((&2, &12)));
+    assert_eq!(it.nth_back(1), Some((&0, &10)));
+    assert_eq!(it.nth_back(0), None);
+
+    let mut it = xs.iter().zip(&ys);
+    assert_eq!(it.nth_back(3), None);
+
+    let mut it = ys.iter().zip(&xs);
+    assert_eq!(it.nth_back(3), None);
+}
+
+#[test]
+fn test_zip_nth_back_after_next() {
+    let xs = [0, 1, 2, 3, 4];
+    let ys = [10, 11, 12, 13, 14];
+
+    let mut it = xs.iter().zip(&ys);
+    assert_eq!(it.next(), Some((&0, &10)));
+    assert_eq!(it.nth_back(1), Some((&3, &13)));
+    assert_eq!(it.nth_back(0), Some((&2, &12)));
+    assert_eq!(it.next(), Some((&1, &11)));
+    assert_eq!(it.next(), None);
+}
+
+#[test]
+fn test_zip_rev_nth_uses_nth_back() {
+    // This is the motivating use-case from issue #54054:
+    // (0..N).rev().step_by(K) should be fast because step_by uses nth,
+    // and rev delegates nth to nth_back.
+    let xs = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9];
+    let ys = [10, 11, 12, 13, 14, 15, 16, 17, 18, 19];
+
+    let v: Vec<_> = xs.iter().zip(&ys).rev().step_by(3).collect();
+    assert_eq!(v, vec![(&9, &19), (&6, &16), (&3, &13), (&0, &10)]);
+}