feat(collections/group_by): accept iterable input, add index param to callback (#3390)

collections/group_by.ts

@@ -28,14 +28,15 @@
  * ```
  */
 export function groupBy<T, K extends string>(
-  array: readonly T[],
-  selector: (el: T) => K,
+  iterable: Iterable<T>,
+  selector: (element: T, index: number) => K,
 ): Partial<Record<K, T[]>> {
   const ret: Partial<Record<K, T[]>> = {};
+  let i = 0;
 
-  for (const element of array) {
-    const key = selector(element);
-    const arr = ret[key] ??= [] as T[];
+  for (const element of iterable) {
+    const key = selector(element, i++);
+    const arr: T[] = ret[key] ??= [];
     arr.push(element);
   }
 

collections/group_by_test.ts

@@ -85,3 +85,31 @@ Deno.test({
     );
   },
 });
+
+Deno.test({
+  name: "[collections/groupBy] callback index",
+  fn() {
+    const actual = groupBy(
+      ["a", "b", "c", "d"],
+      (_, i) => i % 2 === 0 ? "even" : "odd",
+    );
+
+    const expected = { even: ["a", "c"], odd: ["b", "d"] };
+
+    assertEquals(actual, expected);
+  },
+});
+
+Deno.test({
+  name: "[collections/groupBy] iterable input",
+  fn() {
+    function* count(): Generator<number, void> {
+      for (let i = 0; i < 5; i += 1) yield i;
+    }
+
+    const actual = groupBy(count(), (n) => n % 2 === 0 ? "even" : "odd");
+    const expected = { even: [0, 2, 4], odd: [1, 3] };
+
+    assertEquals(actual, expected);
+  },
+});