Auto merge of rust-lang#77594 - timvermeulen:chain_advance_by, r=scottmcm

Implement advance_by, advance_back_by for iter::Chain

Part of rust-lang#77404.

This PR does two things:
- implement `Chain::advance[_back]_by` in terms of `advance[_back]_by` on `self.a` and `advance[_back]_by` on `self.b`
- change `Chain::nth[_back]` to use `advance[_back]_by` on `self.a` and `nth[_back]` on `self.b`

This ensures that `Chain::nth` can take advantage of an efficient `nth` implementation on the second iterator, in case it doesn't implement `advance_by`.

cc `@scottmcm` in case you want to review this

Author: bors

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

@@ -126,16 +126,42 @@ where
     }
 
     #[inline]
-    fn nth(&mut self, mut n: usize) -> Option<A::Item> {
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        let mut rem = n;
+
         if let Some(ref mut a) = self.a {
-            while let Some(x) = a.next() {
-                if n == 0 {
-                    return Some(x);
-                }
-                n -= 1;
+            match a.advance_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
             }
             self.a = None;
         }
+
+        if let Some(ref mut b) = self.b {
+            match b.advance_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
+            }
+            // we don't fuse the second iterator
+        }
+
+        if rem == 0 { Ok(()) } else { Err(n - rem) }
+    }
+
+    #[inline]
+    fn nth(&mut self, mut n: usize) -> Option<Self::Item> {
+        if let Some(ref mut a) = self.a {
+            match a.advance_by(n) {
+                Ok(()) => match a.next() {
+                    None => n = 0,
+                    x => return x,
+                },
+                Err(k) => n -= k,
+            }
+
+            self.a = None;
+        }
+
         maybe!(self.b.nth(n))
     }
 
@@ -202,16 +228,42 @@ where
     }
 
     #[inline]
-    fn nth_back(&mut self, mut n: usize) -> Option<A::Item> {
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        let mut rem = n;
+
         if let Some(ref mut b) = self.b {
-            while let Some(x) = b.next_back() {
-                if n == 0 {
-                    return Some(x);
-                }
-                n -= 1;
+            match b.advance_back_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
             }
             self.b = None;
         }
+
+        if let Some(ref mut a) = self.a {
+            match a.advance_back_by(rem) {
+                Ok(()) => return Ok(()),
+                Err(k) => rem -= k,
+            }
+            // we don't fuse the second iterator
+        }
+
+        if rem == 0 { Ok(()) } else { Err(n - rem) }
+    }
+
+    #[inline]
+    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
+        if let Some(ref mut b) = self.b {
+            match b.advance_back_by(n) {
+                Ok(()) => match b.next_back() {
+                    None => n = 0,
+                    x => return x,
+                },
+                Err(k) => n -= k,
+            }
+
+            self.b = None;
+        }
+
         maybe!(self.a.nth_back(n))
     }
 

library/core/tests/iter.rs

@@ -4,6 +4,43 @@ use core::cell::Cell;
 use core::convert::TryFrom;
 use core::iter::*;
 
+/// An iterator wrapper that panics whenever `next` or `next_back` is called
+/// after `None` has been returned.
+struct Unfuse<I> {
+    iter: I,
+    exhausted: bool,
+}
+
+fn unfuse<I: IntoIterator>(iter: I) -> Unfuse<I::IntoIter> {
+    Unfuse { iter: iter.into_iter(), exhausted: false }
+}
+
+impl<I> Iterator for Unfuse<I>
+where
+    I: Iterator,
+{
+    type Item = I::Item;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        assert!(!self.exhausted);
+        let next = self.iter.next();
+        self.exhausted = next.is_none();
+        next
+    }
+}
+
+impl<I> DoubleEndedIterator for Unfuse<I>
+where
+    I: DoubleEndedIterator,
+{
+    fn next_back(&mut self) -> Option<Self::Item> {
+        assert!(!self.exhausted);
+        let next = self.iter.next_back();
+        self.exhausted = next.is_none();
+        next
+    }
+}
+
 #[test]
 fn test_lt() {
     let empty: [isize; 0] = [];
@@ -142,6 +179,72 @@ fn test_iterator_chain() {
     assert_eq!(i, expected.len());
 }
 
+#[test]
+fn test_iterator_chain_advance_by() {
+    fn test_chain(xs: &[i32], ys: &[i32]) {
+        let len = xs.len() + ys.len();
+
+        for i in 0..xs.len() {
+            let mut iter = unfuse(xs).chain(unfuse(ys));
+            iter.advance_by(i).unwrap();
+            assert_eq!(iter.next(), Some(&xs[i]));
+            assert_eq!(iter.advance_by(100), Err(len - i - 1));
+        }
+
+        for i in 0..ys.len() {
+            let mut iter = unfuse(xs).chain(unfuse(ys));
+            iter.advance_by(xs.len() + i).unwrap();
+            assert_eq!(iter.next(), Some(&ys[i]));
+            assert_eq!(iter.advance_by(100), Err(ys.len() - i - 1));
+        }
+
+        let mut iter = xs.iter().chain(ys);
+        iter.advance_by(len).unwrap();
+        assert_eq!(iter.next(), None);
+
+        let mut iter = xs.iter().chain(ys);
+        assert_eq!(iter.advance_by(len + 1), Err(len));
+    }
+
+    test_chain(&[], &[]);
+    test_chain(&[], &[0, 1, 2, 3, 4, 5]);
+    test_chain(&[0, 1, 2, 3, 4, 5], &[]);
+    test_chain(&[0, 1, 2, 3, 4, 5], &[30, 40, 50, 60]);
+}
+
+#[test]
+fn test_iterator_chain_advance_back_by() {
+    fn test_chain(xs: &[i32], ys: &[i32]) {
+        let len = xs.len() + ys.len();
+
+        for i in 0..ys.len() {
+            let mut iter = unfuse(xs).chain(unfuse(ys));
+            iter.advance_back_by(i).unwrap();
+            assert_eq!(iter.next_back(), Some(&ys[ys.len() - i - 1]));
+            assert_eq!(iter.advance_back_by(100), Err(len - i - 1));
+        }
+
+        for i in 0..xs.len() {
+            let mut iter = unfuse(xs).chain(unfuse(ys));
+            iter.advance_back_by(ys.len() + i).unwrap();
+            assert_eq!(iter.next_back(), Some(&xs[xs.len() - i - 1]));
+            assert_eq!(iter.advance_back_by(100), Err(xs.len() - i - 1));
+        }
+
+        let mut iter = xs.iter().chain(ys);
+        iter.advance_back_by(len).unwrap();
+        assert_eq!(iter.next_back(), None);
+
+        let mut iter = xs.iter().chain(ys);
+        assert_eq!(iter.advance_back_by(len + 1), Err(len));
+    }
+
+    test_chain(&[], &[]);
+    test_chain(&[], &[0, 1, 2, 3, 4, 5]);
+    test_chain(&[0, 1, 2, 3, 4, 5], &[]);
+    test_chain(&[0, 1, 2, 3, 4, 5], &[30, 40, 50, 60]);
+}
+
 #[test]
 fn test_iterator_chain_nth() {
     let xs = [0, 1, 2, 3, 4, 5];