refactor: improve groupBy (#2532)

3 files changed, 82 insertions, 13 deletions

diff --git a/scripts/apidoc/writer.ts b/scripts/apidoc/writer.ts
index 23586230..67bf1d7a 100644
--- a/scripts/apidoc/writer.ts
+++ b/scripts/apidoc/writer.ts
@@ -169,14 +169,13 @@ function writeApiDocsModuleData(
  * @param pages The pages to write into the index.
  */
 export async function writeApiPagesIndex(pages: Page[]): Promise<void> {
-  const pagesByCategory = groupBy(pages, (page) => page.category);
+  const pagesByCategory: Record<string, DefaultTheme.SidebarItem[]> = groupBy(
+    pages,
+    (page) => page.category,
+    ({ text, link }) => ({ text, link })
+  );
   const pageTree = Object.entries(pagesByCategory).flatMap(
-    ([category, items]): DefaultTheme.SidebarItem[] => {
-      const cleanedItems = items.map(({ text, link }) => ({ text, link }));
-      return category
-        ? [{ text: category, items: cleanedItems }]
-        : cleanedItems;
-    }
+    ([category, items]) => (category ? [{ text: category, items }] : items)
   );
 
   // Write api-pages.ts
diff --git a/src/internal/group-by.ts b/src/internal/group-by.ts
index ff3242c9..1fc86e8c 100644
--- a/src/internal/group-by.ts
+++ b/src/internal/group-by.ts
@@ -4,21 +4,50 @@
  * @internal
  *
  * @param values The values to group.
- * @param keyFunction The function to get the key from the value.
+ * @param keyMapper The function to get the key from the value.
  */
 export function groupBy<TValue>(
   values: ReadonlyArray<TValue>,
-  keyFunction: (value: TValue) => string | number
-): Record<string, TValue[]> {
-  const result: Record<string, TValue[]> = {};
+  keyMapper: (value: TValue) => string | number
+): Record<string, TValue[]>;
+/**
+ * Groups the values by the key function and maps the values.
+ *
+ * @internal
+ *
+ * @param values The values to group.
+ * @param keyMapper The function to get the key from the value.
+ * @param valueMapper The function to get the value from the value.
+ */
+export function groupBy<TOriginalValue, TMappedValue>(
+  values: ReadonlyArray<TOriginalValue>,
+  keyMapper: (value: TOriginalValue) => string | number,
+  valueMapper: (value: TOriginalValue) => TMappedValue
+): Record<string, TMappedValue[]>;
+/**
+ * Groups the values by the key function and maps the values.
+ *
+ * @internal
+ *
+ * @param values The values to group.
+ * @param keyMapper The function to get the key from the value.
+ * @param valueMapper The function to map the value.
+ */
+export function groupBy<TOriginalValue, TMappedValue>(
+  values: ReadonlyArray<TOriginalValue>,
+  keyMapper: (value: TOriginalValue) => string | number,
+  valueMapper: (value: TOriginalValue) => TMappedValue = (value) =>
+    value as unknown as TMappedValue
+): Record<string, TMappedValue[]> {
+  const result: Record<string, TMappedValue[]> = {};
 
   for (const value of values) {
-    const key = keyFunction(value);
+    const key = keyMapper(value);
     if (result[key] === undefined) {
       result[key] = [];
     }
 
-    result[key].push(value);
+    result[key].push(valueMapper(value));
   }
 
   return result;
diff --git a/test/internal/group-by.spec.ts b/test/internal/group-by.spec.ts
new file mode 100644
index 00000000..4f186e3d
--- /dev/null
+++ b/test/internal/group-by.spec.ts
@@ -0,0 +1,41 @@
+import { describe, expect, it } from 'vitest';
+import { groupBy } from '../../src/internal/group-by';
+
+describe('groupBy()', () => {
+  it('should group values by key', () => {
+    const values = [
+      { id: 1, name: 'John' },
+      { id: 2, name: 'Jane' },
+      { id: 3, name: 'John' },
+    ];
+
+    const result = groupBy(values, ({ name }) => name);
+
+    expect(result).toEqual({
+      John: [
+        { id: 1, name: 'John' },
+        { id: 3, name: 'John' },
+      ],
+      Jane: [{ id: 2, name: 'Jane' }],
+    });
+  });
+
+  it('should group by key and map values', () => {
+    const values = [
+      { id: 1, name: 'John' },
+      { id: 2, name: 'Jane' },
+      { id: 3, name: 'John' },
+    ];
+
+    const result = groupBy(
+      values,
+      ({ name }) => name,
+      ({ id }) => id
+    );
+
+    expect(result).toEqual({
+      John: [1, 3],
+      Jane: [2],
+    });
+  });
+});