Skip to content

implement Iter::group_by #1998

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Apr 29, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions builtin/builtin.mbti
Original file line number Diff line number Diff line change
Expand Up @@ -233,6 +233,7 @@ impl Iter {
flat_map[T, R](Self[T], (T) -> Self[R]) -> Self[R]
flatten[T](Self[Self[T]]) -> Self[T]
fold[T, B](Self[T], init~ : B, (B, T) -> B) -> B
group_by[T, K : Eq + Hash](Self[T], (T) -> K) -> Map[K, Array[T]]
head[A](Self[A]) -> A?
intersperse[A](Self[A], A) -> Self[A]
iter[T](Self[T]) -> Self[T]
Expand Down
38 changes: 38 additions & 0 deletions builtin/iter.mbt
Original file line number Diff line number Diff line change
Expand Up @@ -998,3 +998,41 @@ pub fn Iter::minimum[T : Compare](self : Iter[T]) -> T? {
}
res
}

///|
/// 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
///
/// A Map where keys are the result of applying the discriminator function to elements,
/// and values are arrays containing all elements that share the same key.
///
/// # Example
///
/// ```moonbit
/// test "group_by" {
/// let iter = [1, 1, 2, 3, 2, 2, 1].iter()
/// let result = iter.group_by(fn(x) { x })
/// assert_eq!(result.get(1), Some([1, 1, 1]))
/// assert_eq!(result.get(2), Some([2, 2, 2]))
/// assert_eq!(result.get(3), Some([3]))
/// }
pub fn Iter::group_by[T, K : Eq + Hash](
self : Iter[T],
f : (T) -> K
) -> Map[K, Array[T]] {
let result = Map::new()
for element in self {
let key = f(element)
match result.get(key) {
Some(arr) => result.set(key, arr + [element])
None => result.set(key, [element])
}
}
result
}
67 changes: 67 additions & 0 deletions builtin/iter_test.mbt
Original file line number Diff line number Diff line change
Expand Up @@ -758,3 +758,70 @@ test "Float::until negative step" {
let result = 10.0.until(0.0, step=-2.0).collect()
inspect!(result, content="[10, 8, 6, 4, 2]")
}

///|
test "group_by with consecutive identical elements" {
let iter = [1, 1, 2, 2, 3, 3].iter()
let grouped = iter.group_by(fn(x) { x })
assert_eq!(grouped.get(1), Some([1, 1]))
assert_eq!(grouped.get(2), Some([2, 2]))
assert_eq!(grouped.get(3), Some([3, 3]))
}

///|
test "group_by with non-consecutive identical elements" {
let iter = [1, 2, 1, 3, 2, 1].iter()
let grouped = iter.group_by(fn(x) { x })
assert_eq!(grouped.get(1), Some([1, 1, 1]))
assert_eq!(grouped.get(2), Some([2, 2]))
assert_eq!(grouped.get(3), Some([3]))
}

///|
test "group_by with empty input" {
let iter : Iter[Int] = Iter::empty()
let grouped = iter.group_by(fn(x) { x })
assert_eq!(grouped.size(), 0)
}

///|
test "group_by with single element input" {
let iter = [42].iter()
let grouped = iter.group_by(fn(x) { x })
assert_eq!(grouped.get(42), Some([42]))
}

///|
test "group_by with custom key function" {
let iter = [1, 2, 3, 4].iter()
let grouped = iter.group_by(fn(x) { x % 2 })
assert_eq!(grouped.get(0), Some([2, 4]))
assert_eq!(grouped.get(1), Some([1, 3]))
}

///|
test "group_by with strings" {
let iter = ["apple", "avocado", "banana", "cherry", "blueberry"].iter()
let grouped = iter.group_by(fn(s) { s[0] })
assert_eq!(grouped.get('a'), Some(["apple", "avocado"]))
assert_eq!(grouped.get('b'), Some(["banana", "blueberry"]))
assert_eq!(grouped.get('c'), Some(["cherry"]))
}

///|
test "group_by with complex objects" {
struct Person {
name : String
age : Int
}
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 grouped = people.group_by(fn(p) { p.age })
let groups = grouped.values().map(fn(a) { a.map(fn(p) { p.name }) }).collect()
assert_eq!(groups, [["Alice", "Bob"], ["Charlie", "Eve"], ["Dave"]])
}
Loading