Rollup merge of #76909 - timvermeulen:advance_by, r=Amanieu

Add Iterator::advance_by and DoubleEndedIterator::advance_back_by

This PR adds the iterator method

```rust
fn advance_by(&mut self, n: usize) -> Result<(), usize>
```

that advances the iterator by `n` elements, returning `Ok(())` if this succeeds or `Err(len)` if the length of the iterator was less than `n`.

Currently `Iterator::nth` is the method to override for efficiently advancing an iterator by multiple elements at once. `advance_by` is superior for this purpose because
- it's simpler to implement: instead of advancing the iterator and producing the next element you only need to advance the iterator
- it composes better: iterators like `Chain` and `FlatMap` can implement `advance_by` in terms of `advance_by` on their inner iterators, but they cannot implement `nth` in terms of `nth` on their inner iterators (see #60395)
- the default implementation of `nth` can trivially be implemented in terms of `advance_by` and `next`, which this PR also does

This PR also adds `DoubleEndedIterator::advance_back_by` for all the same reasons.

I'll make a tracking issue if it's decided this is worth merging. Also let me know if anything can be improved, this went through several iterations so there might very well still be room for improvement (especially in the doc comments). I've written overrides of these methods for most iterators that already override `nth`/`nth_back`, but those still need tests so I'll add them in a later PR.

cc @cuviper @scottmcm @Amanieu

Author: Dylan-DPC

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

@@ -124,6 +124,11 @@ where
         self.iter.size_hint()
     }
 
+    #[inline]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        self.iter.advance_back_by(n)
+    }
+
     #[inline]
     fn nth(&mut self, n: usize) -> Option<<I as Iterator>::Item> {
         self.iter.nth_back(n)
@@ -164,6 +169,11 @@ where
         self.iter.next()
     }
 
+    #[inline]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        self.iter.advance_by(n)
+    }
+
     #[inline]
     fn nth_back(&mut self, n: usize) -> Option<<I as Iterator>::Item> {
         self.iter.nth(n)

library/core/src/iter/traits/double_ended.rs

@@ -91,6 +91,46 @@ pub trait DoubleEndedIterator: Iterator {
     #[stable(feature = "rust1", since = "1.0.0")]
     fn next_back(&mut self) -> Option<Self::Item>;
 
+    /// Advances the iterator from the back by `n` elements.
+    ///
+    /// `advance_back_by` is the reverse version of [`advance_by`]. This method will
+    /// eagerly skip `n` elements starting from the back by calling [`next_back`] up
+    /// to `n` times until [`None`] is encountered.
+    ///
+    /// `advance_back_by(n)` will return [`Ok(())`] if the iterator successfully advances by
+    /// `n` elements, or [`Err(k)`] if [`None`] is encountered, where `k` is the number of
+    /// elements the iterator is advanced by before running out of elements (i.e. the length
+    /// of the iterator). Note that `k` is always less than `n`.
+    ///
+    /// Calling `advance_back_by(0)` does not consume any elements and always returns [`Ok(())`].
+    ///
+    /// [`advance_by`]: Iterator::advance_by
+    /// [`next_back`]: DoubleEndedIterator::next_back
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_advance_by)]
+    ///
+    /// let a = [3, 4, 5, 6];
+    /// let mut iter = a.iter();
+    ///
+    /// assert_eq!(iter.advance_back_by(2), Ok(()));
+    /// assert_eq!(iter.next_back(), Some(&4));
+    /// assert_eq!(iter.advance_back_by(0), Ok(()));
+    /// assert_eq!(iter.advance_back_by(100), Err(1)); // only `&3` was skipped
+    /// ```
+    #[inline]
+    #[unstable(feature = "iter_advance_by", reason = "recently added", issue = "none")]
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        for i in 0..n {
+            self.next_back().ok_or(i)?;
+        }
+        Ok(())
+    }
+
     /// Returns the `n`th element from the end of the iterator.
     ///
     /// This is essentially the reversed version of [`Iterator::nth()`].
@@ -134,14 +174,9 @@ pub trait DoubleEndedIterator: Iterator {
     /// ```
     #[inline]
     #[stable(feature = "iter_nth_back", since = "1.37.0")]
-    fn nth_back(&mut self, mut n: usize) -> Option<Self::Item> {
-        for x in self.rev() {
-            if n == 0 {
-                return Some(x);
-            }
-            n -= 1;
-        }
-        None
+    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
+        self.advance_back_by(n).ok()?;
+        self.next_back()
     }
 
     /// This is the reverse version of [`Iterator::try_fold()`]: it takes
@@ -318,6 +353,9 @@ impl<'a, I: DoubleEndedIterator + ?Sized> DoubleEndedIterator for &'a mut I {
     fn next_back(&mut self) -> Option<I::Item> {
         (**self).next_back()
     }
+    fn advance_back_by(&mut self, n: usize) -> Result<(), usize> {
+        (**self).advance_back_by(n)
+    }
     fn nth_back(&mut self, n: usize) -> Option<I::Item> {
         (**self).nth_back(n)
     }

library/core/src/iter/traits/iterator.rs

@@ -284,6 +284,44 @@ pub trait Iterator {
         self.fold(None, some)
     }
 
+    /// Advances the iterator by `n` elements.
+    ///
+    /// This method will eagerly skip `n` elements by calling [`next`] up to `n`
+    /// times until [`None`] is encountered.
+    ///
+    /// `advance_by(n)` will return [`Ok(())`] if the iterator successfully advances by
+    /// `n` elements, or [`Err(k)`] if [`None`] is encountered, where `k` is the number
+    /// of elements the iterator is advanced by before running out of elements (i.e. the
+    /// length of the iterator). Note that `k` is always less than `n`.
+    ///
+    /// Calling `advance_by(0)` does not consume any elements and always returns [`Ok(())`].
+    ///
+    /// [`next`]: Iterator::next
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// #![feature(iter_advance_by)]
+    ///
+    /// let a = [1, 2, 3, 4];
+    /// let mut iter = a.iter();
+    ///
+    /// assert_eq!(iter.advance_by(2), Ok(()));
+    /// assert_eq!(iter.next(), Some(&3));
+    /// assert_eq!(iter.advance_by(0), Ok(()));
+    /// assert_eq!(iter.advance_by(100), Err(1)); // only `&4` was skipped
+    /// ```
+    #[inline]
+    #[unstable(feature = "iter_advance_by", reason = "recently added", issue = "none")]
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        for i in 0..n {
+            self.next().ok_or(i)?;
+        }
+        Ok(())
+    }
+
     /// Returns the `n`th element of the iterator.
     ///
     /// Like most indexing operations, the count starts from zero, so `nth(0)`
@@ -325,14 +363,9 @@ pub trait Iterator {
     /// ```
     #[inline]
     #[stable(feature = "rust1", since = "1.0.0")]
-    fn 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 nth(&mut self, n: usize) -> Option<Self::Item> {
+        self.advance_by(n).ok()?;
+        self.next()
     }
 
     /// Creates an iterator starting at the same point, but stepping by
@@ -3265,6 +3298,9 @@ impl<I: Iterator + ?Sized> Iterator for &mut I {
     fn size_hint(&self) -> (usize, Option<usize>) {
         (**self).size_hint()
     }
+    fn advance_by(&mut self, n: usize) -> Result<(), usize> {
+        (**self).advance_by(n)
+    }
     fn nth(&mut self, n: usize) -> Option<Self::Item> {
         (**self).nth(n)
     }

library/core/tests/iter.rs

@@ -1570,6 +1570,66 @@ fn test_iterator_rev_nth() {
     assert_eq!(v.iter().rev().nth(v.len()), None);
 }
 
+#[test]
+fn test_iterator_advance_by() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+
+    for i in 0..v.len() {
+        let mut iter = v.iter();
+        assert_eq!(iter.advance_by(i), Ok(()));
+        assert_eq!(iter.next().unwrap(), &v[i]);
+        assert_eq!(iter.advance_by(100), Err(v.len() - 1 - i));
+    }
+
+    assert_eq!(v.iter().advance_by(v.len()), Ok(()));
+    assert_eq!(v.iter().advance_by(100), Err(v.len()));
+}
+
+#[test]
+fn test_iterator_advance_back_by() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+
+    for i in 0..v.len() {
+        let mut iter = v.iter();
+        assert_eq!(iter.advance_back_by(i), Ok(()));
+        assert_eq!(iter.next_back().unwrap(), &v[v.len() - 1 - i]);
+        assert_eq!(iter.advance_back_by(100), Err(v.len() - 1 - i));
+    }
+
+    assert_eq!(v.iter().advance_back_by(v.len()), Ok(()));
+    assert_eq!(v.iter().advance_back_by(100), Err(v.len()));
+}
+
+#[test]
+fn test_iterator_rev_advance_by() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+
+    for i in 0..v.len() {
+        let mut iter = v.iter().rev();
+        assert_eq!(iter.advance_by(i), Ok(()));
+        assert_eq!(iter.next().unwrap(), &v[v.len() - 1 - i]);
+        assert_eq!(iter.advance_by(100), Err(v.len() - 1 - i));
+    }
+
+    assert_eq!(v.iter().rev().advance_by(v.len()), Ok(()));
+    assert_eq!(v.iter().rev().advance_by(100), Err(v.len()));
+}
+
+#[test]
+fn test_iterator_rev_advance_back_by() {
+    let v: &[_] = &[0, 1, 2, 3, 4];
+
+    for i in 0..v.len() {
+        let mut iter = v.iter().rev();
+        assert_eq!(iter.advance_back_by(i), Ok(()));
+        assert_eq!(iter.next_back().unwrap(), &v[i]);
+        assert_eq!(iter.advance_back_by(100), Err(v.len() - 1 - i));
+    }
+
+    assert_eq!(v.iter().rev().advance_back_by(v.len()), Ok(()));
+    assert_eq!(v.iter().rev().advance_back_by(100), Err(v.len()));
+}
+
 #[test]
 fn test_iterator_last() {
     let v: &[_] = &[0, 1, 2, 3, 4];

library/core/tests/lib.rs

@@ -41,6 +41,7 @@
 #![feature(slice_partition_dedup)]
 #![feature(int_error_matching)]
 #![feature(array_value_iter)]
+#![feature(iter_advance_by)]
 #![feature(iter_partition_in_place)]
 #![feature(iter_is_partitioned)]
 #![feature(iter_order_by)]