implement Iter::chunk_by

Author: juliano

builtin/builtin.mbti

@@ -218,6 +218,7 @@ impl Iter {
   all[T](Self[T], (T) -> Bool) -> Bool
   any[T](Self[T], (T) -> Bool) -> Bool
   append[T](Self[T], T) -> Self[T]
+  chunk_by[T, K : Eq](Self[T], (T) -> K) -> Self[(K, Self[T])]
   collect[T](Self[T]) -> Array[T]
   concat[T](Self[T], Self[T]) -> Self[T]
   contains[A : Eq](Self[A], A) -> Bool

builtin/iter.mbt

@@ -1036,3 +1036,68 @@ pub fn Iter::group_by[T, K : Eq + Hash](
   }
   result
 }
+
+///|
+/// Groups elements of an iterator according to a discriminator function.
+///
+/// # Parameters
+///
+/// * `self` - The input iterator.
+/// * `f` - The discriminator function that maps elements to keys.
+///
+/// # Returns
+///
+/// An iterator of tuples where each tuple contains a key and an iterator of elements that share that key.
+///
+/// # Example
+///
+/// ```moonbit
+/// test "chunk_by" {
+///   let iter = [1, 1, 2, 3, 2, 2, 1].iter()
+///   let chunked = iter.chunk_by(fn(x) { x })
+///   let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+///   assert_eq!(result, [(1, [1, 1]), (2, [2]), (3, [3]), (2, [2, 2]), (1, [1])])
+/// }
+pub fn Iter::chunk_by[T, K : Eq](
+  self : Iter[T],
+  f : (T) -> K
+) -> Iter[(K, Iter[T])] {
+  fn(yield_) {
+    let mut has_current = false
+    let mut current_key : K? = None
+    let mut buffer : Array[T] = []
+    fn emit_current_chunk() -> IterResult {
+      if buffer.is_empty() {
+        return IterContinue
+      }
+      let key = current_key.unwrap()
+      let elements = buffer
+      buffer = Array::new(capacity=16)
+      yield_((key, elements.iter()))
+    }
+
+    self.run(fn(x) {
+      let key = f(x)
+      if has_current && key == current_key.unwrap() {
+        buffer.push(x)
+        IterContinue
+      } else {
+        if has_current {
+          if emit_current_chunk() == IterEnd {
+            return IterEnd
+          }
+        }
+        has_current = true
+        current_key = Some(key)
+        buffer.push(x)
+        IterContinue
+      }
+    })
+    |> ignore
+    if has_current {
+      emit_current_chunk()
+    } else {
+      IterContinue
+    }
+  }
+}

builtin/iter_test.mbt

@@ -825,3 +825,88 @@ test "group_by with complex objects" {
   let groups = grouped.values().map(fn(a) { a.map(fn(p) { p.name }) }).collect()
   assert_eq!(groups, [["Alice", "Bob"], ["Charlie", "Eve"], ["Dave"]])
 }
+
+///|
+test "chunk_by with consecutive identical elements" {
+  let iter = [1, 1, 2, 2, 3, 3].iter()
+  let chunked = iter.chunk_by(fn(x) { x })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [(1, [1, 1]), (2, [2, 2]), (3, [3, 3])])
+}
+
+///|
+test "chunk_by with non-consecutive identical elements" {
+  let iter = [1, 2, 1, 3, 2, 1].iter()
+  let chunked = iter.chunk_by(fn(x) { x })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [
+    (1, [1]),
+    (2, [2]),
+    (1, [1]),
+    (3, [3]),
+    (2, [2]),
+    (1, [1]),
+  ])
+}
+
+///|
+test "chunk_by with empty input" {
+  let iter = Iter::empty()
+  let chunked = iter.chunk_by(fn(x) { x })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [])
+}
+
+///|
+test "chunk_by with single element input" {
+  let iter = [42].iter()
+  let chunked = iter.chunk_by(fn(x) { x })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [(42, [42])])
+}
+
+///|
+test "chunk_by with custom key function" {
+  let iter = [1, 2, 3, 4].iter()
+  let chunked = iter.chunk_by(fn(x) { x % 2 })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [(1, [1]), (0, [2]), (1, [3]), (0, [4])])
+}
+
+///|
+test "chunk_by with strings" {
+  let iter = ["apple", "avocado", "banana", "cherry", "blueberry"].iter()
+  let chunked = iter.chunk_by(fn(s) { s[0] })
+  let result = chunked.map(fn(g) { (g.0, g.1.collect()) }).collect()
+  assert_eq!(result, [
+    ('a', ["apple", "avocado"]),
+    ('b', ["banana"]),
+    ('c', ["cherry"]),
+    ('b', ["blueberry"]),
+  ])
+}
+
+///|
+test "chunk_by with complex objects" {
+  struct Person {
+    name : String
+    age : Int
+  } derive(Show)
+  let people = [
+    Person::{ name: "Alice", age: 25 },
+    Person::{ name: "Bob", age: 25 },
+    Person::{ name: "Charlie", age: 30 },
+    Person::{ name: "Dave", age: 35 },
+    Person::{ name: "Eve", age: 30 },
+  ].iter()
+  let chunked = people.chunk_by(fn(p) { p.age })
+  let groups = chunked
+    .map(fn(g) { (g.0, g.1.map(fn(p) { p.name }).collect()) })
+    .collect()
+  assert_eq!(groups, [
+    (25, ["Alice", "Bob"]),
+    (30, ["Charlie"]),
+    (35, ["Dave"]),
+    (30, ["Eve"]),
+  ])
+}