feat(corelib): core::iter::chain

corelib/src/iter.cairo

@@ -1,2 +1,4 @@
+mod adapters;
 mod traits;
+pub use adapters::{Chain, chained_iterator};
 pub use traits::{IntoIterator, Iterator};

corelib/src/iter/adapters.cairo

@@ -0,0 +1,2 @@
+mod chain_adapter;
+pub use chain_adapter::{Chain, chained_iterator};

corelib/src/iter/adapters/chain_adapter.cairo

@@ -0,0 +1,51 @@
+/// An iterator that links two iterators together, in a chain.
+///
+/// This `struct` is created by [`Iterator::chain`]. See the
+/// documentation for more.
+#[derive(Drop)]
+pub struct Chain<A, B> {
+    // These are "fused" with `Option` so we don't need separate state to track which part is
+    // already exhausted, and we may also get niche layout for `None`.
+    //
+    // Only the "first" iterator is actually set `None` when exhausted.
+    a: Option<A>,
+    b: Option<B>,
+}
+
+#[inline]
+pub fn chained_iterator<A, B>(a: A, b: B) -> Chain<A, B> {
+    Chain { a: Option::Some(a), b: Option::Some(b) }
+}
+
+impl ChainIterator<
+    A,
+    B,
+    impl IterA: Iterator<A>,
+    +Iterator<B>[Item: IterA::Item],
+    +Drop<A>,
+    +Drop<B>,
+    +Drop<IterA::Item>,
+> of Iterator<Chain<A, B>> {
+    type Item = IterA::Item;
+
+    fn next(ref self: Chain<A, B>) -> Option<Self::Item> {
+        // First iterate over first container values
+        if self.a.is_some() {
+            let mut first_container = self.a.unwrap();
+            let value = first_container.next();
+            if value.is_some() {
+                self.a = Option::Some(first_container);
+                return value;
+            } else {
+                self.a = Option::None;
+            }
+        }
+
+        // Then iterate over second container values
+        let mut second_container = self.b.unwrap();
+        let value = second_container.next();
+        self.b = Option::Some(second_container);
+
+        value
+    }
+}

corelib/src/iter/traits/collect.cairo

@@ -91,6 +91,8 @@ pub trait IntoIterator<T> {
     fn into_iter(self: T) -> Self::IntoIter;
 }
 
+/// Trying to convert an already existing iterator into an iterator
+/// just returns the iterator itself
 impl IteratorIntoIterator<T, +Iterator<T>> of IntoIterator<T> {
     type IntoIter = T;
     fn into_iter(self: T) -> T {

corelib/src/iter/traits/iterator.cairo

@@ -1,3 +1,6 @@
+use core::iter::adapters::{Chain, chained_iterator};
+use core::metaprogramming::TypeEqual;
+
 /// A trait for dealing with iterators.
 ///
 /// This is the main iterator trait. For more about the concept of iterators
@@ -38,4 +41,64 @@ pub trait Iterator<T> {
     /// assert_eq!(Option::None, iter.next());
     /// ```
     fn next(ref self: T) -> Option<Self::Item>;
+
+    /// Takes two iterators and creates a new iterator over both in sequence.
+    ///
+    /// `chain()` will return a new iterator which will first iterate over
+    /// values from the first iterator and then over values from the second
+    /// iterator.
+    ///
+    /// In other words, it links two iterators together, in a chain. 🔗
+    ///
+    /// Arguments do not have to be of the same type as long as the underlying iterated
+    /// over items are.
+    ///
+    /// # Examples
+    ///
+    /// Basic usage:
+    ///
+    /// ```
+    /// use core::ops::Range;
+    ///
+    /// let a: Array<u8> = array![7, 8, 9];
+    /// let b: Range<u8> = 0..5;
+    ///
+    /// let mut iter = a.into_iter().chain(b.into_iter());
+    ///
+    /// assert_eq!(iter.next(), Option::Some(7));
+    /// assert_eq!(iter.next(), Option::Some(8));
+    /// assert_eq!(iter.next(), Option::Some(9));
+    /// assert_eq!(iter.next(), Option::Some(0));
+    /// assert_eq!(iter.next(), Option::Some(1));
+    /// assert_eq!(iter.next(), Option::Some(2));
+    /// assert_eq!(iter.next(), Option::Some(3));
+    /// assert_eq!(iter.next(), Option::Some(4));
+    /// assert_eq!(iter.next(), Option::None);
+    /// ```
+    ///
+    /// Since the argument to `chain()` uses [`IntoIterator`], we can pass
+    /// anything that can be converted into an [`Iterator`], not just an
+    /// [`Iterator`] itself. For example, arrays implement
+    /// [`IntoIterator`], and so can be passed to `chain()` directly:
+    ///
+    /// ```
+    /// let a = array![1, 2, 3];
+    /// let b = array![4, 5, 6];
+    ///
+    /// let mut iter = a.into_iter().chain(b);
+    ///
+    /// assert_eq!(iter.next(), Option::Some(1));
+    /// assert_eq!(iter.next(), Option::Some(2));
+    /// assert_eq!(iter.next(), Option::Some(3));
+    /// assert_eq!(iter.next(), Option::Some(4));
+    /// assert_eq!(iter.next(), Option::Some(5));
+    /// assert_eq!(iter.next(), Option::Some(6));
+    /// assert_eq!(iter.next(), Option::None);
+    /// ```
+    fn chain<U, impl IntoIterU: IntoIterator<U>, +TypeEqual<Self::Item, IntoIterU::Iterator::Item>>(
+        self: T, other: U,
+    ) -> Chain<T, IntoIterU::IntoIter>
+    {
+        chained_iterator(self, other.into_iter())
+    }
 }

corelib/src/test.cairo

@@ -13,6 +13,7 @@ mod felt_test;
 mod fmt_test;
 mod hash_test;
 mod integer_test;
+mod iter_test;
 mod keccak_test;
 mod math_test;
 mod nullable_test;

corelib/src/test/iter_test.cairo

@@ -0,0 +1,35 @@
+use core::ops::Range;
+
+#[test]
+fn test_iterator_chain_different_types() {
+    let a: Array<u8> = array![7, 8, 9];
+    let b: Range<u8> = 0..5;
+
+    let mut iter = a.into_iter().chain(b.into_iter());
+
+    assert_eq!(iter.next(), Option::Some(7));
+    assert_eq!(iter.next(), Option::Some(8));
+    assert_eq!(iter.next(), Option::Some(9));
+    assert_eq!(iter.next(), Option::Some(0));
+    assert_eq!(iter.next(), Option::Some(1));
+    assert_eq!(iter.next(), Option::Some(2));
+    assert_eq!(iter.next(), Option::Some(3));
+    assert_eq!(iter.next(), Option::Some(4));
+    assert_eq!(iter.next(), Option::None);
+}
+
+#[test]
+fn test_iterator_chain_without_into_iter() {
+    let a = array![1, 2, 3];
+    let b = array![4, 5, 6];
+
+    let mut iter = a.into_iter().chain(b);
+
+    assert_eq!(iter.next(), Option::Some(1));
+    assert_eq!(iter.next(), Option::Some(2));
+    assert_eq!(iter.next(), Option::Some(3));
+    assert_eq!(iter.next(), Option::Some(4));
+    assert_eq!(iter.next(), Option::Some(5));
+    assert_eq!(iter.next(), Option::Some(6));
+    assert_eq!(iter.next(), Option::None);
+}