diff --git a/builtin/builtin.mbti b/builtin/builtin.mbti index 0167e70c0..b5558bc24 100644 --- a/builtin/builtin.mbti +++ b/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 diff --git a/builtin/iter.mbt b/builtin/iter.mbt index ee211f938..f4e3ca7ff 100644 --- a/builtin/iter.mbt +++ b/builtin/iter.mbt @@ -1036,3 +1036,54 @@ 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 current_key : K? = None + let mut buffer : Array[T] = [] + for x in self { + let key = f(x) + match current_key { + None => current_key = Some(key) + Some(old_key) => + if key != old_key { + guard yield_((old_key, buffer.iter())) is IterContinue else { + return IterEnd + } + buffer = Array::new(capacity=16) + current_key = Some(key) + } + } + buffer.push(x) + } + if current_key is Some(old_key) { + yield_((old_key, buffer.iter())) + } else { + IterContinue + } + } +} diff --git a/builtin/iter_test.mbt b/builtin/iter_test.mbt index 2b918a365..91817ba68 100644 --- a/builtin/iter_test.mbt +++ b/builtin/iter_test.mbt @@ -825,3 +825,100 @@ 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 break" { + let iter = [1, 2].iter() + let chunked = iter.chunk_by(fn(x) { x }).append((3, [3].iter())) + let mut c = 0 + for _ in chunked { + c += 1 + break + } + assert_eq!(c, 1) +} + +///| +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"]), + ]) +}