From ab881eff9ee2666a10d1be80847bdb0924d37c51 Mon Sep 17 00:00:00 2001
From: MagisterDallis <194415551+MagisterDallis@users.noreply.github.com>
Date: Wed, 15 Jan 2025 17:22:13 +0100
Subject: [PATCH] feat: Iterator::count

---
 corelib/src/iter/traits/iterator.cairo | 40 ++++++++++++++++++++++++++
 corelib/src/test/iter_test.cairo       | 12 ++++++++
 2 files changed, 52 insertions(+)

diff --git a/corelib/src/iter/traits/iterator.cairo b/corelib/src/iter/traits/iterator.cairo
index daffe4d4e7a..f984db66d36 100644
--- a/corelib/src/iter/traits/iterator.cairo
+++ b/corelib/src/iter/traits/iterator.cairo
@@ -41,6 +41,46 @@ pub trait Iterator<T> {
     /// ```
     fn next(ref self: T) -> Option<Self::Item>;
 
+    /// Consumes the iterator, counting the number of iterations and returning it.
+    ///
+    /// This method will call [`next`] repeatedly until [`None`] is encountered,
+    /// returning the number of times it saw [`Some`]. Note that [`next`] has to be
+    /// called at least once even if the iterator does not have any elements.
+    ///
+    /// [`next`]: Iterator::next
+    ///
+    /// # Overflow Behavior
+    ///
+    /// The method does no guarding against overflows, so counting elements of
+    /// an iterator with more than [`Bounded::<usize>::MAX`] elements either produces the
+    /// wrong result or panics.
+    ///
+    /// [`Bounded::<usize>::MAX`]: core::num::traits::Bounded
+    ///
+    /// # Panics
+    ///
+    /// This function might panic if the iterator has more than [`Bounded::<usize>::MAX`]
+    /// elements.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// let mut a = array![1, 2, 3].into_iter();
+    /// assert_eq!(a.count(), 3);
+    ///
+    /// let mut a = array![1, 2, 3, 4, 5].into_iter();
+    /// assert_eq!(a.count(), 5);
+    /// ```
+    #[inline]
+    fn count<+Destruct<T>, +Destruct<Self::Item>>(
+        self: T,
+    ) -> usize {
+        let mut self = self;
+        Self::fold(ref self, 0_usize, |count, _x| {
+            count + 1
+        })
+    }
+
     /// Advances the iterator by `n` elements.
     ///
     /// This method will eagerly skip `n` elements by calling [`next`] up to `n`
diff --git a/corelib/src/test/iter_test.cairo b/corelib/src/test/iter_test.cairo
index 6053c440625..530aa7cde33 100644
--- a/corelib/src/test/iter_test.cairo
+++ b/corelib/src/test/iter_test.cairo
@@ -1,3 +1,15 @@
+#[test]
+fn test_iter_count() {
+    let mut empty_iter = ArrayTrait::<usize>::new().into_iter();
+    let count = empty_iter.count();
+    assert_eq!(count, 0);
+
+    let mut iter = array![1, 2, 3].into_iter();
+    let count = iter.count();
+
+    assert_eq!(count, 3);
+}
+
 #[test]
 fn test_advance_by() {
     let mut iter = array![1_u8, 2, 3, 4].into_iter();