feat(corelib): Iterator::peekable

corelib/src/iter.cairo

@@ -231,4 +231,5 @@
 //! [`map`]: Iterator::map
 mod adapters;
 mod traits;
+pub use adapters::PeekableTrait;
 pub use traits::{FromIterator, IntoIterator, Iterator};

corelib/src/iter/adapters.cairo

@@ -12,3 +12,8 @@ mod zip;
 pub use zip::Zip;
 #[allow(unused_imports)]
 pub(crate) use zip::zipped_iterator;
+
+mod peekable;
+#[allow(unused_imports)]
+pub(crate) use peekable::peekable_iterator;
+pub use peekable::{Peekable, PeekableTrait};

corelib/src/iter/adapters/peekable.cairo

@@ -0,0 +1,90 @@
+/// An iterator with a `peek()` that returns a copy to the next
+/// element.
+///
+/// This `struct` is created by the [`peekable`] method on [`Iterator`]. See its
+/// documentation for more.
+///
+/// [`peekable`]: Iterator::peekable
+#[must_use]
+#[derive(Drop, Debug)]
+pub struct Peekable<I, R> {
+    /// The underlying iterator.
+    iter: I,
+    /// Caches the next value of the iterator when `peek()` is called,
+    /// consuming the iterator only once even when several consecutive calls
+    /// to `peek()` are made.
+    ///
+    /// - `None`: Indicates that no peek has been performed.
+    /// - `Some(None)`: Indicates that a peek has been performed but the iterator is exhausted.
+    /// - `Some(Some(R))`: Contains the next element in the iterator, ready to be returned by
+    /// `peek()`.
+    ///
+    /// `peeked` is reset to `None` when `next()` is called.
+    peeked: Option<Option<R>>,
+}
+
+pub fn peekable_iterator<I, R>(iter: I) -> Peekable<I, R> {
+    Peekable { iter, peeked: Option::None }
+}
+
+impl PeekableIterator<
+    I, impl IterI: Iterator<I>, +Destruct<I>, +Destruct<IterI::Item>,
+> of Iterator<Peekable<I, IterI::Item>> {
+    type Item = IterI::Item;
+
+    fn next(ref self: Peekable<I, IterI::Item>) -> Option<Self::Item> {
+        let next_value = match self.peeked {
+            Option::Some(v) => v,
+            Option::None => self.iter.next(),
+        };
+        self.peeked = Option::None;
+        next_value
+    }
+}
+
+#[generate_trait]
+pub impl PeekableTraitImpl<
+    I, impl IterI: Iterator<I>, +Copy<IterI::Item>, +Drop<IterI::Item>,
+> of PeekableTrait<I> {
+    /// Returns a copy to the next() value without advancing the iterator.
+    ///
+    /// Like [`next`], if there is a value, it is wrapped in a `Some(T)`.
+    /// But if the iteration is over, `None` is returned.
+    ///
+    /// [`next`]: Iterator::next
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// let mut iter = (1..4_u8).into_iter().peekable();
+    ///
+    /// // peek() lets us see one step into the future
+    /// assert_eq!(iter.peek(), Some(1));
+    /// assert_eq!(iter.next(), Some(1));
+    ///
+    /// assert_eq!(iter.next(), Some(2));
+    ///
+    /// // The iterator does not advance even if we `peek` multiple times
+    /// assert_eq!(iter.peek(), Some(3));
+    /// assert_eq!(iter.peek(), Some(3));
+    ///
+    /// assert_eq!(iter.next(), Some(3));
+    ///
+    /// // After the iterator is finished, so is `peek()`
+    /// assert_eq!(iter.peek(), None);
+    /// assert_eq!(iter.next(), None);
+    /// ```
+    #[inline]
+    fn peek(ref self: Peekable<I, IterI::Item>) -> Option<IterI::Item> {
+        match self.peeked {
+            Option::Some(value) => value,
+            Option::None => {
+                let next_value = self.iter.next();
+                self.peeked = Option::Some(next_value);
+                next_value
+            },
+        }
+    }
+}

corelib/src/iter/traits/iterator.cairo

@@ -1,5 +1,6 @@
 use crate::iter::adapters::{
-    Enumerate, Map, Zip, enumerated_iterator, mapped_iterator, zipped_iterator,
+    Enumerate, Map, Peekable, Zip, enumerated_iterator, mapped_iterator, peekable_iterator,
+    zipped_iterator,
 };
 
 /// A trait for dealing with iterators.
@@ -407,4 +408,41 @@ pub trait Iterator<T> {
     ) -> B {
         FromIterator::<B, Self::Item>::from_iter::<T, Self>(self)
     }
+
+    /// Creates an iterator which can use the [`peek`] method to look at the next element of the
+    /// iterator. See its documentation for more information.
+    ///
+    /// Note that the underlying iterator is still advanced when [`peek`] is called for the first
+    /// time: In order to retrieve the next element, [`next`] is called on the underlying iterator,
+    /// hence any side effects (i.e. anything other than fetching the next value) of the [`next`]
+    /// method will occur.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// let mut iter = (1..4_u8).into_iter().peekable();
+    ///
+    /// // peek() lets us see one step into the future
+    /// assert_eq!(iter.peek(), Option::Some(1));
+    /// assert_eq!(iter.next(), Option::Some(1));
+    ///
+    /// assert_eq!(iter.next(), Option::Some(2));
+    ///
+    /// // we can peek() multiple times, the iterator won't advance
+    /// assert_eq!(iter.peek(), Option::Some(3));
+    /// assert_eq!(iter.peek(), Option::Some(3));
+    ///
+    /// assert_eq!(iter.next(), Option::Some(3));
+    ///
+    /// // after the iterator is finished, so is peek()
+    /// assert_eq!(iter.peek(), Option::None);
+    /// assert_eq!(iter.next(), Option::None);
+    /// ```
+    #[inline]
+    #[must_use]
+    fn peekable(self: T) -> Peekable<T, Self::Item> {
+        peekable_iterator(self)
+    }
 }

corelib/src/test/iter_test.cairo

@@ -1,3 +1,5 @@
+use core::iter::PeekableTrait;
+
 #[test]
 fn test_iter_count() {
     let mut empty_iter = ArrayTrait::<usize>::new().into_iter();
@@ -70,3 +72,22 @@ fn test_iter_adapter_fold() {
 fn test_iter_adapter_collect() {
     assert_eq!((0..3_u32).into_iter().collect(), array![0, 1, 2]);
 }
+
+#[test]
+fn test_iter_adapter_peekable() {
+    let mut iter = (1..4_u8).into_iter().peekable();
+
+    // peek() lets us see one step into the future
+    assert_eq!(iter.peek(), Option::Some(1));
+    assert_eq!(iter.next(), Option::Some(1));
+    assert_eq!(iter.next(), Option::Some(2));
+
+    // The iterator does not advance even if we `peek` multiple times
+    assert_eq!(iter.peek(), Option::Some(3));
+    assert_eq!(iter.peek(), Option::Some(3));
+    assert_eq!(iter.next(), Option::Some(3));
+
+    // After the iterator is finished, so is `peek()`
+    assert_eq!(iter.peek(), Option::None);
+    assert_eq!(iter.next(), Option::None);
+}