Spread LuaPairsIterable and LuaPairsKeysIterable (#1409)

Perryvw · web-flow · commit 3b30b09fce2c · 2023-02-26T11:29:00.000+01:00
* Spread LuaPairsIterable and LuaPairsKeysIterable

* Use expect.toContainEqual in tests
diff --git a/src/transformation/visitors/spread.ts b/src/transformation/visitors/spread.ts
@@ -1,8 +1,9 @@
 import * as ts from "typescript";
 import { LuaTarget } from "../../CompilerOptions";
 import * as lua from "../../LuaAST";
+import { assertNever } from "../../utils";
 import { FunctionVisitor, TransformationContext } from "../context";
-import { isLuaIterable } from "../utils/language-extensions";
+import { getIterableExtensionKindForNode, IterableExtensionKind } from "../utils/language-extensions";
 import { createUnpackCall } from "../utils/lua-ast";
 import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import {
@@ -76,11 +77,22 @@ export const transformSpreadElement: FunctionVisitor<ts.SpreadElement> = (node,
     const innerExpression = context.transformExpression(node.expression);
     if (isMultiReturnCall(context, tsInnerExpression)) return innerExpression;
 
-    const type = context.checker.getTypeAtLocation(node.expression); // not ts-inner expression, in case of casts
-    if (isLuaIterable(context, type)) {
-        return transformLuaLibFunction(context, LuaLibFeature.LuaIteratorSpread, node, innerExpression);
+    const iterableExtensionType = getIterableExtensionKindForNode(context, node.expression);
+    if (iterableExtensionType) {
+        if (iterableExtensionType === IterableExtensionKind.Iterable) {
+            return transformLuaLibFunction(context, LuaLibFeature.LuaIteratorSpread, node, innerExpression);
+        } else if (iterableExtensionType === IterableExtensionKind.Pairs) {
+            const objectEntries = transformLuaLibFunction(context, LuaLibFeature.ObjectEntries, node, innerExpression);
+            return createUnpackCall(context, objectEntries, node);
+        } else if (iterableExtensionType === IterableExtensionKind.PairsKey) {
+            const objectKeys = transformLuaLibFunction(context, LuaLibFeature.ObjectKeys, node, innerExpression);
+            return createUnpackCall(context, objectKeys, node);
+        } else {
+            assertNever(iterableExtensionType);
+        }
     }
 
+    const type = context.checker.getTypeAtLocation(node.expression); // not ts-inner expression, in case of casts
     if (isArrayType(context, type)) {
         return createUnpackCall(context, innerExpression, node);
     }
diff --git a/test/unit/spread.spec.ts b/test/unit/spread.spec.ts
@@ -509,7 +509,7 @@ test.each(["pairs", "ipairs"])("can spread %s (#1244)", func => {
 });
 
 // https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1244
-test.each(["LuaTable", "LuaMap"])("can spread %s (#1244)", type => {
+test.each(["LuaTable", "LuaMap"])("can spread %s with pairs (#1244)", type => {
     const result: Array<[string, string]> = util.testFunction`
         const tbl = new ${type}();
         tbl.set("foo", "bar");
@@ -524,6 +524,40 @@ test.each(["LuaTable", "LuaMap"])("can spread %s (#1244)", type => {
 
     // We don't know the order so match like this
     expect(result).toHaveLength(2);
-    expect(result.some(([k, v]) => k === "foo" && v === "bar")).toBe(true);
-    expect(result.some(([k, v]) => k === "fizz" && v === "buzz")).toBe(true);
+    expect(result).toContainEqual(["foo", "bar"]);
+    expect(result).toContainEqual(["fizz", "buzz"]);
+});
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1384
+test.each(["LuaTable", "LuaMap"])("can spread %s (#1384)", type => {
+    const result: Array<[string, string]> = util.testFunction`
+        const tbl = new ${type}();
+        tbl.set("foo", "bar");
+        tbl.set("fizz", "buzz");
+        return [...tbl];
+    `
+        .withLanguageExtensions()
+        .getLuaExecutionResult();
+
+    // We don't know the order so match like this
+    expect(result).toHaveLength(2);
+    expect(result).toContainEqual(["foo", "bar"]);
+    expect(result).toContainEqual(["fizz", "buzz"]);
+});
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1384
+test("can spread LuaSet (#1384)", () => {
+    const result: Array<[string, string]> = util.testFunction`
+        const tbl = new LuaSet();
+        tbl.add("foo");
+        tbl.add("bar");
+        return [...tbl];
+    `
+        .withLanguageExtensions()
+        .getLuaExecutionResult();
+
+    // We don't know the order so match like this
+    expect(result).toHaveLength(2);
+    expect(result).toContain("foo");
+    expect(result).toContain("bar");
 });
