diff --git a/corelib/src/iter/adapters.cairo b/corelib/src/iter/adapters.cairo index fdd38e85c67..ee647323441 100644 --- a/corelib/src/iter/adapters.cairo +++ b/corelib/src/iter/adapters.cairo @@ -2,3 +2,8 @@ mod map; pub use map::Map; #[allow(unused_imports)] pub(crate) use map::mapped_iterator; + +mod enumerate; +pub use enumerate::Enumerate; +#[allow(unused_imports)] +pub(crate) use enumerate::enumerated_iterator; diff --git a/corelib/src/iter/adapters/enumerate.cairo b/corelib/src/iter/adapters/enumerate.cairo new file mode 100644 index 00000000000..7157a9aaf05 --- /dev/null +++ b/corelib/src/iter/adapters/enumerate.cairo @@ -0,0 +1,41 @@ +/// An iterator that yields the current count and the element during iteration. +/// +/// This `struct` is created by the [`enumerate`] method on [`Iterator`]. See its +/// documentation for more. +/// +/// [`enumerate`]: Iterator::enumerate +/// [`Iterator`]: core::iter::Iterator +#[must_use] +#[derive(Drop, Clone, Debug)] +pub struct Enumerate { + iter: I, + count: usize, +} + +pub fn enumerated_iterator(iter: I) -> Enumerate { + Enumerate { iter, count: 0 } +} + +impl EnumerateIterator< + I, T, +Iterator[Item: T], +Destruct, +Destruct, +> of Iterator> { + type Item = (usize, T); + + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so enumerating more than + /// `Bounded::::MAX` elements will always panic. + /// + /// [`Bounded`]: core::num::traits::Bounded + /// + /// # Panics + /// + /// Will panic if the index of the element overflows a `usize`. + #[inline] + fn next(ref self: Enumerate) -> Option { + let a = self.iter.next()?; + let i = self.count; + self.count += 1; + Option::Some((i, a)) + } +} diff --git a/corelib/src/iter/traits/iterator.cairo b/corelib/src/iter/traits/iterator.cairo index 7ee26837f4c..adf8d44bbd0 100644 --- a/corelib/src/iter/traits/iterator.cairo +++ b/corelib/src/iter/traits/iterator.cairo @@ -1,4 +1,4 @@ -use crate::iter::adapters::{Map, mapped_iterator}; +use crate::iter::adapters::{Enumerate, Map, enumerated_iterator, mapped_iterator}; /// A trait for dealing with iterators. /// @@ -93,4 +93,39 @@ pub trait Iterator { ) -> Map { mapped_iterator(self, f) } + + /// Creates an iterator which gives the current iteration count as well as + /// the next value. + /// + /// The iterator returned yields pairs `(i, val)`, where `i` is the + /// current index of iteration and `val` is the value returned by the + /// iterator. + /// + /// `enumerate()` keeps its count as a [`usize`]. + /// + /// # Overflow Behavior + /// + /// The method does no guarding against overflows, so enumerating more than + /// `Bounded::::MAX` elements will always panic. + /// + /// [`Bounded`]: core::num::traits::Bounded + /// + /// # Panics + /// + /// Will panic if the to-be-returned index overflows a `usize`. + /// + /// # Examples + /// + /// ``` + /// let mut iter = array!['a', 'b', 'c'].into_iter().enumerate(); + /// + /// assert_eq!(iter.next(), Option::Some((0, 'a'))); + /// assert_eq!(iter.next(), Option::Some((1, 'b'))); + /// assert_eq!(iter.next(), Option::Some((2, 'c'))); + /// assert_eq!(iter.next(), Option::None); + /// ``` + #[inline] + fn enumerate(self: T) -> Enumerate { + enumerated_iterator(self) + } } diff --git a/corelib/src/test/iter_test.cairo b/corelib/src/test/iter_test.cairo index d5aa58efdd6..4b77f19c610 100644 --- a/corelib/src/test/iter_test.cairo +++ b/corelib/src/test/iter_test.cairo @@ -7,3 +7,12 @@ fn test_iter_adapter_map() { assert_eq!(iter.next(), Option::Some(6)); assert_eq!(iter.next(), Option::None); } + +fn test_iterator_enumerate() { + let mut iter = array!['a', 'b', 'c'].into_iter().enumerate(); + + assert_eq!(iter.next(), Option::Some((0, 'a'))); + assert_eq!(iter.next(), Option::Some((1, 'b'))); + assert_eq!(iter.next(), Option::Some((2, 'c'))); + assert_eq!(iter.next(), Option::None); +}