Fixes to return for TupleReturn functions (#232)

tomblind · Perryvw · commit 8e9dfe8c4e71 · 2018-10-13T16:32:56.000+02:00
* TupleReturn functions now correctly unpack return values or forward results from other TupleReturn call

* indentation fix
diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -82,6 +82,20 @@ export class TSHelper {
         }
     }
 
+    public static isInTupleReturnFunction(node: ts.Node, checker: ts.TypeChecker): boolean {
+        const declaration = this.findFirstNodeAbove(node, (n): n is ts.Node =>
+            ts.isFunctionDeclaration(n) || ts.isMethodDeclaration(n));
+        if (declaration) {
+            const decorators = this.getCustomDecorators(
+                checker.getTypeAtLocation(declaration),
+                checker
+            );
+            return decorators.has(DecoratorKind.TupleReturn);
+        } else {
+            return false;
+        }
+    }
+
     public static getCustomDecorators(type: ts.Type, checker: ts.TypeChecker): Map<DecoratorKind, Decorator> {
         if (type.symbol) {
             const comments = type.symbol.getDocumentationComment(checker);
diff --git a/src/Transpiler.ts b/src/Transpiler.ts
@@ -685,20 +685,15 @@ export abstract class LuaTranspiler {
 
     public transpileReturn(node: ts.ReturnStatement): string {
         if (node.expression) {
-            // If parent function is a TupleReturn function
-            // and return expression is an array literal, leave out brackets.
-            const declaration = tsHelper.findFirstNodeAbove(node, (n): n is ts.Node =>
-                ts.isFunctionDeclaration(n) || ts.isMethodDeclaration(n));
-            let isTupleReturn = false;
-            if (declaration) {
-                const decorators = tsHelper.getCustomDecorators(
-                    this.checker.getTypeAtLocation(declaration),
-                    this.checker
-                );
-                isTupleReturn = decorators.has(DecoratorKind.TupleReturn);
-            }
-            if (isTupleReturn && ts.isArrayLiteralExpression(node.expression)) {
-                return "return " + node.expression.elements.map(elem => this.transpileExpression(elem)).join(",");
+            if (tsHelper.isInTupleReturnFunction(node, this.checker)) {
+                // Parent function is a TupleReturn function
+                if (ts.isArrayLiteralExpression(node.expression)) {
+                    // If return expression is an array literal, leave out brackets.
+                    return "return " + node.expression.elements.map(elem => this.transpileExpression(elem)).join(",");
+                } else if (!tsHelper.isTupleReturnCall(node.expression, this.checker)) {
+                    // If return expression is not another TupleReturn call, unpack it
+                    return `return ${this.transpileDestructingAssignmentValue(node.expression)}`;
+                }
             }
 
             return "return " + this.transpileExpression(node.expression);
@@ -1085,12 +1080,15 @@ export abstract class LuaTranspiler {
         let callPath;
 
         const isTupleReturn = tsHelper.isTupleReturnCall(node, this.checker);
+        const isTupleReturnForward = node.parent && ts.isReturnStatement(node.parent)
+            && tsHelper.isInTupleReturnFunction(node, this.checker);
         const isInDestructingAssignment = tsHelper.isInDestructingAssignment(node);
         const returnValueIsUsed = node.parent && !ts.isExpressionStatement(node.parent);
 
         if (ts.isPropertyAccessExpression(node.expression)) {
             const result = this.transpilePropertyCall(node);
-            return isTupleReturn && !isInDestructingAssignment && returnValueIsUsed ? `({ ${result} })` : result;
+            return isTupleReturn && !isTupleReturnForward && !isInDestructingAssignment && returnValueIsUsed
+                ? `({ ${result} })` : result;
         }
 
         // Handle super calls properly
@@ -1102,7 +1100,7 @@ export abstract class LuaTranspiler {
 
         callPath = this.transpileExpression(node.expression);
         params = this.transpileArguments(node.arguments);
-        return isTupleReturn && !isInDestructingAssignment && returnValueIsUsed
+        return isTupleReturn && !isTupleReturnForward && !isInDestructingAssignment && returnValueIsUsed
             ? `({ ${callPath}(${params}) })` : `${callPath}(${params})`;
     }
 
diff --git a/test/translation/lua/tupleReturn.lua b/test/translation/lua/tupleReturn.lua
@@ -1,3 +1,6 @@
+function tupleReturn()
+    return 0,"foobar"
+end
 tupleReturn();
 noTupleReturn();
 local a,b=tupleReturn();
@@ -10,3 +13,16 @@ e = ({ tupleReturn() });
 f = noTupleReturn();
 foo(({ tupleReturn() }));
 foo(noTupleReturn());
+function tupleReturnFromVar()
+    local r = {1,"baz"};
+    return table.unpack(r)
+end
+function tupleReturnForward()
+    return tupleReturn()
+end
+function tupleNoForward()
+    return ({ tupleReturn() })
+end
+function tupleReturnUnpack()
+    return table.unpack(tupleNoForward())
+end
diff --git a/test/translation/ts/tupleReturn.ts b/test/translation/ts/tupleReturn.ts
@@ -1,5 +1,7 @@
 /** !TupleReturn */
-declare function tupleReturn(): [number, string];
+function tupleReturn(): [number, string] {
+    return [0, "foobar"];
+}
 declare function noTupleReturn(): [number, string];
 declare function foo(a: [number, string]): void;
 tupleReturn();
@@ -14,3 +16,19 @@ e = tupleReturn();
 f = noTupleReturn();
 foo(tupleReturn());
 foo(noTupleReturn());
+/** !TupleReturn */
+function tupleReturnFromVar(): [number, string] {
+    const r: [number, string] = [1, "baz"];
+    return r;
+}
+/** !TupleReturn */
+function tupleReturnForward(): [number, string] {
+    return tupleReturn();
+}
+function tupleNoForward(): [number, string] {
+    return tupleReturn();
+}
+/** !TupleReturn */
+function tupleReturnUnpack(): [number, string] {
+    return tupleNoForward();
+}
