diff --git a/corelib/src/iter/traits/iterator.cairo b/corelib/src/iter/traits/iterator.cairo index adf8d44bbd0..a7e617929ff 100644 --- a/corelib/src/iter/traits/iterator.cairo +++ b/corelib/src/iter/traits/iterator.cairo @@ -41,6 +41,47 @@ pub trait Iterator { /// ``` fn next(ref self: T) -> Option; + /// 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 a `Err(NonZero)` with value `k` if [`None`] is encountered, + /// where `k` is remaining number of steps that could not be advanced because the iterator ran + /// out. + /// If `self` is empty and `n` is non-zero, then this returns `Err(n)`. + /// Otherwise, `k` is always less than `n`. + /// + /// [`None`]: Option::None + /// [`next`]: Iterator::next + /// + /// # Examples + /// + /// ``` + /// let mut iter = array![1_u8, 2, 3, 4].into_iter(); + /// + /// assert_eq!(iter.advance_by(2), Result::Ok(())); + /// assert_eq!(iter.next(), Option::Some(3)); + /// assert_eq!(iter.advance_by(0), Result::Ok(())); + /// assert_eq!(iter.advance_by(100), Result::Err(99)); + /// ``` + fn advance_by<+Destruct, +Destruct>( + ref self: T, n: usize, + ) -> Result< + (), NonZero, + > { + if let Option::Some(nz_n) = n.try_into() { + if let Option::Some(_) = Self::next(ref self) { + return Self::advance_by(ref self, n - 1); + } else { + Result::Err(nz_n) + } + } else { + Result::Ok(()) + } + } + /// Takes a closure and creates an iterator which calls that closure on each /// element. /// diff --git a/corelib/src/test/iter_test.cairo b/corelib/src/test/iter_test.cairo index 4b77f19c610..599b03f6d58 100644 --- a/corelib/src/test/iter_test.cairo +++ b/corelib/src/test/iter_test.cairo @@ -1,3 +1,13 @@ +#[test] +fn test_advance_by() { + let mut iter = array![1_u8, 2, 3, 4].into_iter(); + + assert_eq!(iter.advance_by(2), Result::Ok(())); + assert_eq!(iter.next(), Option::Some(3)); + assert_eq!(iter.advance_by(0), Result::Ok(())); + assert_eq!(iter.advance_by(100), Result::Err(99)); +} + #[test] fn test_iter_adapter_map() { let mut iter = array![1, 2, 3].into_iter().map(|x| 2 * x); diff --git a/tests/bug_samples/issue7031.cairo b/tests/bug_samples/issue7031.cairo index e6df5cdc41d..21d030e5c7c 100644 --- a/tests/bug_samples/issue7031.cairo +++ b/tests/bug_samples/issue7031.cairo @@ -1,5 +1,5 @@ pub trait IteratorEx, +Destruct, +Drop> { - fn advance_by( + fn advance_by_( ref self: T, n: usize, ) -> Result< (), NonZero, @@ -19,7 +19,7 @@ impl ItratorExImpl, +Destruct, +Drop> of Iter #[test] fn test_advance_by() { let mut iter = array![1_u8, 2, 3, 4].into_iter(); - assert_eq!(iter.advance_by(2), Result::Ok(())); + assert_eq!(iter.advance_by_(2), Result::Ok(())); assert_eq!(iter.next(), Option::Some(3)); - assert_eq!(iter.advance_by(0), Result::Ok(())); + assert_eq!(iter.advance_by_(0), Result::Ok(())); }