From 8fd1bb57d5bdbf4851f2190a1764522a0c4c565b Mon Sep 17 00:00:00 2001
From: Yuki Okushi <huyuumi.dev@gmail.com>
Date: Tue, 14 Jan 2020 13:26:00 +0900
Subject: [PATCH] Implement specialized `nth_back()` for Zip

---
 src/libcore/iter/adapters/zip.rs | 60 ++++++++++++++++++++++++++++++++
 src/libcore/tests/iter.rs        | 37 ++++++++++++++++++++
 2 files changed, 97 insertions(+)

diff --git a/src/libcore/iter/adapters/zip.rs b/src/libcore/iter/adapters/zip.rs
index b13e12e2e8608..8241ae9023025 100644
--- a/src/libcore/iter/adapters/zip.rs
+++ b/src/libcore/iter/adapters/zip.rs
@@ -36,6 +36,22 @@ impl<A: Iterator, B: Iterator> Zip<A, B> {
     }
 }
 
+impl<A: DoubleEndedIterator, B: DoubleEndedIterator> Zip<A, B>
+where
+    A: DoubleEndedIterator + ExactSizeIterator,
+    B: DoubleEndedIterator + ExactSizeIterator,
+{
+    fn super_nth_back(&mut self, mut n: usize) -> Option<(A::Item, B::Item)> {
+        while let Some(x) = DoubleEndedIterator::next_back(self) {
+            if n == 0 {
+                return Some(x);
+            }
+            n -= 1;
+        }
+        None
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<A, B> Iterator for Zip<A, B>
 where
@@ -70,6 +86,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<Self::Item> {
+        ZipImpl::nth_back(self, n)
+    }
 }
 
 // Zip specialization trait
@@ -84,6 +105,10 @@ trait ZipImpl<A, B> {
     where
         A: DoubleEndedIterator + ExactSizeIterator,
         B: DoubleEndedIterator + ExactSizeIterator;
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator;
 }
 
 // General Zip impl
@@ -142,6 +167,15 @@ where
         }
     }
 
+    #[inline]
+    default fn nth_back(&mut self, n: usize) -> Option<Self::Item>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator,
+    {
+        self.super_nth_back(n)
+    }
+
     #[inline]
     default fn size_hint(&self) -> (usize, Option<usize>) {
         let (a_lower, a_upper) = self.a.size_hint();
@@ -248,6 +282,32 @@ where
             None
         }
     }
+
+    #[inline]
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item>
+    where
+        A: DoubleEndedIterator + ExactSizeIterator,
+        B: DoubleEndedIterator + ExactSizeIterator,
+    {
+        let delta = cmp::min(n, self.index);
+        let end = self.index - delta;
+        while end < self.index {
+            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_back(n - delta)
+    }
 }
 
 #[stable(feature = "rust1", since = "1.0.0")]
diff --git a/src/libcore/tests/iter.rs b/src/libcore/tests/iter.rs
index 8b8dc941534ee..8b514ae3d003a 100644
--- a/src/libcore/tests/iter.rs
+++ b/src/libcore/tests/iter.rs
@@ -269,6 +269,22 @@ fn test_zip_nth() {
     assert_eq!(it.nth(3), None);
 }
 
+#[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_side_effects() {
     let mut a = Vec::new();
@@ -291,6 +307,27 @@ fn test_zip_nth_side_effects() {
     assert_eq!(b, vec![200, 300, 400, 500, 600]);
 }
 
+#[test]
+fn test_zip_nth_back_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
+        }))
+        .nth_back(3);
+    assert_eq!(value, Some((30, 4000)));
+    assert_eq!(a, vec![6, 6, 5, 5, 4, 4, 3]);
+    assert_eq!(b, vec![800, 700, 700, 600, 600, 500, 500, 400]);
+}
+
 #[test]
 fn test_iterator_step_by() {
     // Identity