Implement Map.groupBy and Object.groupBy (#1546)

Perryvw · web-flow · commit 90b378444167 · 2024-03-17T16:00:45.000+01:00
* Add support for Object.groupBy

* Implement Map.groupBy
diff --git a/src/LuaLib.ts b/src/LuaLib.ts
@@ -58,6 +58,7 @@ export enum LuaLibFeature {
     Iterator = "Iterator",
     LuaIteratorSpread = "LuaIteratorSpread",
     Map = "Map",
+    MapGroupBy = "MapGroupBy",
     Match = "Match",
     MathAtan2 = "MathAtan2",
     MathModf = "MathModf",
@@ -77,6 +78,7 @@ export enum LuaLibFeature {
     ObjectFromEntries = "ObjectFromEntries",
     ObjectGetOwnPropertyDescriptor = "ObjectGetOwnPropertyDescriptor",
     ObjectGetOwnPropertyDescriptors = "ObjectGetOwnPropertyDescriptors",
+    ObjectGroupBy = "ObjectGroupBy",
     ObjectKeys = "ObjectKeys",
     ObjectRest = "ObjectRest",
     ObjectValues = "ObjectValues",
diff --git a/src/lualib/MapGroupBy.ts b/src/lualib/MapGroupBy.ts
@@ -0,0 +1,22 @@
+export function __TS__MapGroupBy<K, T>(
+    this: void,
+    items: Iterable<T>,
+    keySelector: (item: T, index: number) => K
+): Map<K, T[]> {
+    const result = new Map<K, T[]>();
+
+    let i = 0;
+    for (const item of items) {
+        const key = keySelector(item, i);
+
+        if (result.has(key)) {
+            result.get(key)!.push(item);
+        } else {
+            result.set(key, [item]);
+        }
+
+        i++;
+    }
+
+    return result;
+}
diff --git a/src/lualib/ObjectGroupBy.ts b/src/lualib/ObjectGroupBy.ts
@@ -0,0 +1,22 @@
+export function __TS__ObjectGroupBy<K extends PropertyKey, T>(
+    this: void,
+    items: Iterable<T>,
+    keySelector: (item: T, index: number) => K
+): Partial<Record<K, T[]>> {
+    const result: Partial<Record<K, T[]>> = {};
+
+    let i = 0;
+    for (const item of items) {
+        const key = keySelector(item, i);
+
+        if (key in result) {
+            result[key]!.push(item);
+        } else {
+            result[key] = [item];
+        }
+
+        i++;
+    }
+
+    return result;
+}
diff --git a/src/transformation/builtins/index.ts b/src/transformation/builtins/index.ts
@@ -18,6 +18,7 @@ import { transformStringConstructorCall, transformStringProperty, transformStrin
 import { transformSymbolConstructorCall } from "./symbol";
 import { unsupportedBuiltinOptionalCall } from "../utils/diagnostics";
 import { LuaTarget } from "../../CompilerOptions";
+import { transformMapConstructorCall } from "./map";
 
 export function transformBuiltinPropertyAccessExpression(
     context: TransformationContext,
@@ -93,6 +94,9 @@ function tryTransformBuiltinGlobalMethodCall(
         case "Console":
             result = transformConsoleCall(context, node, calledMethod);
             break;
+        case "MapConstructor":
+            result = transformMapConstructorCall(context, node, calledMethod);
+            break;
         case "Math":
             result = transformMathCall(context, node, calledMethod);
             break;
diff --git a/src/transformation/builtins/map.ts b/src/transformation/builtins/map.ts
@@ -0,0 +1,22 @@
+import * as lua from "../../LuaAST";
+import * as ts from "typescript";
+import { TransformationContext } from "../context";
+import { unsupportedProperty } from "../utils/diagnostics";
+import { transformArguments } from "../visitors/call";
+import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
+
+export function transformMapConstructorCall(
+    context: TransformationContext,
+    node: ts.CallExpression,
+    calledMethod: ts.PropertyAccessExpression
+): lua.Expression | undefined {
+    const args = transformArguments(context, node.arguments);
+    const methodName = calledMethod.name.text;
+
+    switch (methodName) {
+        case "groupBy":
+            return transformLuaLibFunction(context, LuaLibFeature.MapGroupBy, node, ...args);
+        default:
+            context.diagnostics.push(unsupportedProperty(calledMethod.name, "Map", methodName));
+    }
+}
diff --git a/src/transformation/builtins/object.ts b/src/transformation/builtins/object.ts
@@ -26,6 +26,8 @@ export function transformObjectConstructorCall(
             return transformLuaLibFunction(context, LuaLibFeature.ObjectGetOwnPropertyDescriptor, node, ...args);
         case "getOwnPropertyDescriptors":
             return transformLuaLibFunction(context, LuaLibFeature.ObjectGetOwnPropertyDescriptors, node, ...args);
+        case "groupBy":
+            return transformLuaLibFunction(context, LuaLibFeature.ObjectGroupBy, node, ...args);
         case "keys":
             return transformLuaLibFunction(context, LuaLibFeature.ObjectKeys, node, ...args);
         case "values":
diff --git a/test/unit/builtins/map.spec.ts b/test/unit/builtins/map.spec.ts
@@ -209,3 +209,47 @@ describe.each(iterationMethods)("map.%s() preserves insertion order", iterationM
         `.expectToMatchJsResult();
     });
 });
+
+describe("Map.groupBy", () => {
+    test("empty", () => {
+        util.testFunction`
+            const array = [];
+
+            const map = Map.groupBy(array, (num, index) => {
+                return num % 2 === 0 ? "even": "odd";
+            });
+
+            return Object.fromEntries(map.entries());
+        `.expectToEqual([]);
+    });
+
+    test("groupBy", () => {
+        util.testFunction`
+            const array = [0, 1, 2, 3, 4, 5];
+
+            const map = Map.groupBy(array, (num, index) => {
+                return num % 2 === 0 ? "even": "odd";
+            });
+
+            return Object.fromEntries(map.entries());
+        `.expectToEqual({
+            even: [0, 2, 4],
+            odd: [1, 3, 5],
+        });
+    });
+
+    test("groupBy index", () => {
+        util.testFunction`
+            const array = [0, 1, 2, 3, 4, 5];
+
+            const map = Map.groupBy(array, (num, index) => {
+                return index < 3 ? "low": "high";
+            });
+
+            return Object.fromEntries(map.entries());
+        `.expectToEqual({
+            low: [0, 1, 2],
+            high: [3, 4, 5],
+        });
+    });
+});
diff --git a/test/unit/builtins/object.spec.ts b/test/unit/builtins/object.spec.ts
@@ -264,3 +264,41 @@ describe("delete from object", () => {
             .expectToMatchJsResult();
     });
 });
+
+describe("Object.groupBy", () => {
+    test("empty", () => {
+        util.testFunction`
+            const array = [];
+
+            return Object.groupBy(array, (num, index) => {
+                return num % 2 === 0 ? "even": "odd";
+            });
+        `.expectToEqual([]);
+    });
+
+    test("groupBy", () => {
+        util.testFunction`
+            const array = [0, 1, 2, 3, 4, 5];
+
+            return Object.groupBy(array, (num, index) => {
+                return num % 2 === 0 ? "even": "odd";
+            });
+        `.expectToEqual({
+            even: [0, 2, 4],
+            odd: [1, 3, 5],
+        });
+    });
+
+    test("groupBy index", () => {
+        util.testFunction`
+            const array = [0, 1, 2, 3, 4, 5];
+
+            return Object.groupBy(array, (num, index) => {
+                return index < 3 ? "low": "high";
+            });
+        `.expectToEqual({
+            low: [0, 1, 2],
+            high: [3, 4, 5],
+        });
+    });
+});
