Add advance_by and advance_back_by

Author: timvermeulen

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 [`nth`]. Although like most indexing
@@ -135,14 +175,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 [`try_fold()`]: it takes elements
@@ -323,6 +358,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
@@ -3262,6 +3295,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

@@ -1473,6 +1473,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

@@ -38,6 +38,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)]