fixed return-forwarding of luaIterator functions (#395)

tomblind · web-flow · commit 978d025a1a87 · 2019-02-12T08:49:40.000-07:00
fixes #389
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -1495,7 +1495,10 @@ export class LuaTransformer {
                     // If return expression is an array literal, leave out brackets.
                     return tstl.createReturnStatement(statement.expression.elements
                         .map(elem => this.transformExpression(elem)));
-                } else if (!tsHelper.isTupleReturnCall(statement.expression, this.checker)) {
+                } else if (
+                    !tsHelper.isTupleReturnCall(statement.expression, this.checker)
+                    && !tsHelper.isInLuaIteratorFunction(statement, this.checker))
+                {
                     // If return expression is not another TupleReturn call, unpack it
                     const expression = this.createUnpackCall(
                         this.transformExpression(statement.expression),
diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -177,15 +177,20 @@ export class TSHelper {
     }
 
     public static isInTupleReturnFunction(node: ts.Node, checker: ts.TypeChecker): boolean {
-        const declaration = TSHelper.findFirstNodeAbove(
-            node,
-            (n): n is ts.Node => ts.isFunctionDeclaration(n) || ts.isMethodDeclaration(n)
-        );
+        const declaration = TSHelper.findFirstNodeAbove(node, ts.isFunctionLike);
+        if (declaration) {
+            const decorators = TSHelper.getCustomDecorators(checker.getTypeAtLocation(declaration), checker);
+            return decorators.has(DecoratorKind.TupleReturn);
+        } else {
+            return false;
+        }
+    }
+
+    public static isInLuaIteratorFunction(node: ts.Node, checker: ts.TypeChecker): boolean {
+        const declaration = TSHelper.findFirstNodeAbove(node, ts.isFunctionLike);
         if (declaration) {
             const decorators = TSHelper.getCustomDecorators(checker.getTypeAtLocation(declaration), checker);
-            return decorators.has(DecoratorKind.TupleReturn)
-                   // Lua iterators are not 'true' tupleReturn functions as they actually return a function
-                   && !decorators.has(DecoratorKind.LuaIterator);
+            return decorators.has(DecoratorKind.LuaIterator);
         } else {
             return false;
         }
diff --git a/test/unit/loops.spec.ts b/test/unit/loops.spec.ts
@@ -680,6 +680,61 @@ export class LuaLoopTests
             + "the TupleReturn decorator.");
     }
 
+    @Test("forof forwarded lua iterator")
+    public forofForwardedLuaIterator(): void {
+        const code =
+            `const arr = ["a", "b", "c"];
+            /** @luaIterator */
+            function luaIter(): Iterable<string> {
+                let i = 0;
+                function iter() { return arr[i++]; }
+                return iter as any;
+            }
+            /** @luaIterator */
+            function forward(): Iterable<string> {
+                return luaIter();
+            }
+            let result = "";
+            for (let a of forward()) { result += a; }
+            return result;`;
+        const compilerOptions = {
+            luaLibImport: LuaLibImportKind.Require,
+            luaTarget: LuaTarget.Lua53,
+            target: ts.ScriptTarget.ES2015,
+        };
+        const result = util.transpileAndExecute(code, compilerOptions);
+        Expect(result).toBe("abc");
+    }
+
+    @Test("forof forwarded lua iterator with tupleReturn")
+    public forofForwardedLuaIteratorWithTupleReturn(): void {
+        const code =
+            `const arr = ["a", "b", "c"];
+            /** @luaIterator */
+            /** @tupleReturn */
+            function luaIter(): Iterable<[string, string]> {
+                let i = 0;
+                /** @tupleReturn */
+                function iter() { return arr[i] && [i.toString(), arr[i++]] || []; }
+                return iter as any;
+            }
+            /** @luaIterator */
+            /** @tupleReturn */
+            function forward(): Iterable<[string, string]> {
+                return luaIter();
+            }
+            let result = "";
+            for (let [a, b] of forward()) { result += a + b; }
+            return result;`;
+        const compilerOptions = {
+            luaLibImport: LuaLibImportKind.Require,
+            luaTarget: LuaTarget.Lua53,
+            target: ts.ScriptTarget.ES2015,
+        };
+        const result = util.transpileAndExecute(code, compilerOptions);
+        Expect(result).toBe("0a1b2c");
+    }
+
     @TestCase("while (a < b) { i++; }")
     @TestCase("do { i++; } while (a < b)")
     @TestCase("for (let i = 0; i < 3; i++) {}")
