From ecacc7534b6bf50205c37c89402565b82d95a257 Mon Sep 17 00:00:00 2001 From: Tim Vermeulen Date: Sat, 19 Sep 2020 10:26:46 +0200 Subject: [PATCH] Add advance_by and advance_back_by --- library/core/src/iter/adapters/mod.rs | 10 ++++ library/core/src/iter/traits/double_ended.rs | 54 +++++++++++++++--- library/core/src/iter/traits/iterator.rs | 52 ++++++++++++++--- library/core/tests/iter.rs | 60 ++++++++++++++++++++ library/core/tests/lib.rs | 1 + 5 files changed, 161 insertions(+), 16 deletions(-) diff --git a/library/core/src/iter/adapters/mod.rs b/library/core/src/iter/adapters/mod.rs index ab27fe15a8e..7cddc6773ff 100644 --- a/library/core/src/iter/adapters/mod.rs +++ b/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<::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<::Item> { self.iter.nth(n) diff --git a/library/core/src/iter/traits/double_ended.rs b/library/core/src/iter/traits/double_ended.rs index a025bc8b560..a5cfc022859 100644 --- a/library/core/src/iter/traits/double_ended.rs +++ b/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; + /// 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 { - for x in self.rev() { - if n == 0 { - return Some(x); - } - n -= 1; - } - None + fn nth_back(&mut self, n: usize) -> Option { + 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 { (**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 { (**self).nth_back(n) } diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs index b8a09f822b6..8ff4c273fcd 100644 --- a/library/core/src/iter/traits/iterator.rs +++ b/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 { - while let Some(x) = self.next() { - if n == 0 { - return Some(x); - } - n -= 1; - } - None + fn nth(&mut self, n: usize) -> Option { + self.advance_by(n).ok()?; + self.next() } /// Creates an iterator starting at the same point, but stepping by @@ -3262,6 +3295,9 @@ impl Iterator for &mut I { fn size_hint(&self) -> (usize, Option) { (**self).size_hint() } + fn advance_by(&mut self, n: usize) -> Result<(), usize> { + (**self).advance_by(n) + } fn nth(&mut self, n: usize) -> Option { (**self).nth(n) } diff --git a/library/core/tests/iter.rs b/library/core/tests/iter.rs index 00e3972c42f..7a8450116fb 100644 --- a/library/core/tests/iter.rs +++ b/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]; diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index a5b1b51e06c..fd3f71f8786 100644 --- a/library/core/tests/lib.rs +++ b/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)]