Fix LuaMultiReturn not being recognized inside more complicated iterables (#1595)

Perryvw · web-flow · commit 43d827410f20 · 2024-11-08T20:27:00.000+01:00
* Fix LuaMultiReturn not being recognized inside more complicated iterable types

* Fix it in a way that doesn't break other stuff
diff --git a/package-lock.json b/package-lock.json
diff --git a/src/transformation/visitors/language-extensions/iterable.ts b/src/transformation/visitors/language-extensions/iterable.ts
@@ -62,8 +62,13 @@ export function transformForOfIterableStatement(
     statement: ts.ForOfStatement,
     block: lua.Block
 ): lua.Statement {
-    const type = context.checker.getTypeAtLocation(statement.expression);
-    if (type.aliasTypeArguments?.length === 2 && isMultiReturnType(type.aliasTypeArguments[0])) {
+    const iteratedExpressionType = context.checker.getTypeAtLocation(statement.expression);
+    const iterableType =
+        iteratedExpressionType.isIntersection() &&
+        iteratedExpressionType.types.find(t => t.symbol.escapedName === "Iterable");
+    const iterableTypeArguments = (iterableType as ts.TypeReference)?.typeArguments;
+
+    if (iterableTypeArguments && iterableTypeArguments.length > 0 && isMultiReturnType(iterableTypeArguments[0])) {
         const luaIterator = context.transformExpression(statement.expression);
         return transformForOfMultiIterableStatement(
             context,
diff --git a/test/unit/language-extensions/multi.spec.ts b/test/unit/language-extensions/multi.spec.ts
@@ -373,3 +373,31 @@ describe("LuaMultiReturn returns all values even when indexed with [0] #1411", (
             .expectToEqual(1);
     });
 });
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1591
+test("LuaMultiReturn in LuaIterable (#1591)", () => {
+    const lua = util.testModule`
+        type IterableAlias = LuaIterable<LuaMultiReturn<[number, number]>>;
+
+        declare const iterable: IterableAlias;
+
+        for (const [a, b] of iterable) {}
+    `
+        .withLanguageExtensions()
+        .getMainLuaCodeChunk();
+
+    expect(lua).toContain("a, b");
+});
+
+// https://github.com/TypeScriptToLua/TypeScriptToLua/issues/1591
+test("LuaMultiReturn in LuaIterable intersection (#1591)", () => {
+    const lua = util.testModule`
+        declare function iterator(): { a: string } & LuaIterable<LuaMultiReturn<[string, string]>>;
+
+        for (const [a, b] of iterator()) {}
+    `
+        .withLanguageExtensions()
+        .getMainLuaCodeChunk();
+
+    expect(lua).toContain("a, b");
+});
