TypeScriptToLua · tomblind · Feb 23, 2019 · Feb 23, 2019 · Feb 23, 2019 · Feb 23, 2019
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -2972,7 +2972,15 @@ export class LuaTransformer {
             hasContext ? context : undefined
         );
 
-        const body = ts.isBlock(node.body) ? node.body : ts.createBlock([ts.createReturn(node.body)]);
+        let body: ts.Block;
+        if (ts.isBlock(node.body)) {
+            body = node.body;
+        } else {
+            const returnExpression = ts.createReturn(node.body);
+            body = ts.createBlock([returnExpression]);
+            returnExpression.parent = body;
+            body.parent = node.body.parent;
+        }
         const [transformedBody] = this.transformFunctionBody(node.parameters, body, spreadIdentifier);
 
         return tstl.createFunctionExpression(
@@ -3053,9 +3061,10 @@ export class LuaTransformer {
         const isTupleReturnForward =
             node.parent && ts.isReturnStatement(node.parent) && tsHelper.isInTupleReturnFunction(node, this.checker);
         const isInDestructingAssignment = tsHelper.isInDestructingAssignment(node);
+        const isInSpread = node.parent && ts.isSpreadElement(node.parent);
         const returnValueIsUsed = node.parent && !ts.isExpressionStatement(node.parent);
-        const wrapResult = isTupleReturn && !isTupleReturnForward&& !isInDestructingAssignment
-            && returnValueIsUsed && !isLuaIterator;
+        const wrapResult = isTupleReturn && !isTupleReturnForward && !isInDestructingAssignment
+            && !isInSpread && returnValueIsUsed && !isLuaIterator;
 
         if (ts.isPropertyAccessExpression(node.expression)) {
             const result = this.transformPropertyCall(node);
@@ -3707,8 +3716,11 @@ export class LuaTransformer {
 
     public transformSpreadElement(expression: ts.SpreadElement): ExpressionVisitResult {
         const innerExpression = this.transformExpression(expression.expression);
-
-        return this.createUnpackCall(innerExpression, expression);
+        if (tsHelper.isTupleReturnCall(expression.expression, this.checker)) {
+            return innerExpression;
+        } else {
+            return this.createUnpackCall(innerExpression, expression);
+        }
     }
 
     public transformStringLiteral(literal: ts.StringLiteralLike): tstl.StringLiteral {

diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -187,7 +187,13 @@ export class TSHelper {
     public static isInTupleReturnFunction(node: ts.Node, checker: ts.TypeChecker): boolean {
         const declaration = TSHelper.findFirstNodeAbove(node, ts.isFunctionLike);
         if (declaration) {
-            const decorators = TSHelper.getCustomDecorators(checker.getTypeAtLocation(declaration), checker);
+            let functionType: ts.Type;
+            if (ts.isFunctionExpression(declaration) || ts.isArrowFunction(declaration)) {
+                functionType = TSHelper.inferAssignedType(declaration, checker);
+            } else {
+                functionType = checker.getTypeAtLocation(declaration);
+            }
+            const decorators = TSHelper.getCustomDecorators(functionType, checker);
             return decorators.has(DecoratorKind.TupleReturn);
         } else {
             return false;
@@ -446,15 +452,36 @@ export class TSHelper {
 
         } else if (ts.isCallExpression(expression.parent)) {
             // Expression being passed as argument to a function
-            let i = expression.parent.arguments.indexOf(expression);
-            if (i >= 0) {
+            const argumentIndex = expression.parent.arguments.indexOf(expression);
+            if (argumentIndex >= 0) {
                 const parentSignature = checker.getResolvedSignature(expression.parent);
-                const parentSignatureDeclaration = parentSignature.getDeclaration();
-                if (parentSignatureDeclaration) {
-                    if (this.getExplicitThisParameter(parentSignatureDeclaration)) {
-                        ++i;
+                if (parentSignature.parameters.length > 0) { // In case function type is 'any'
+                    const signatureIndex = Math.min(argumentIndex, parentSignature.parameters.length - 1);
+                    let parameterType = checker.getTypeOfSymbolAtLocation(
+                        parentSignature.parameters[signatureIndex],
+                        expression
+                    );
+                    if (TSHelper.isArrayType(parameterType, checker)) {
+                        // Check for elipses argument
+                        const parentSignatureDeclaration = parentSignature.getDeclaration();
+                        if (parentSignatureDeclaration) {
+                            let declarationIndex = signatureIndex;
+                            if (this.getExplicitThisParameter(parentSignatureDeclaration)) {
+                                // Ignore 'this' parameter
+                                ++declarationIndex;
+                            }
+                            const parameterDeclaration = parentSignatureDeclaration.parameters[declarationIndex];
+                            if ((parameterType.flags & ts.TypeFlags.Object) !== 0
+                                && ((parameterType as ts.ObjectType).objectFlags & ts.ObjectFlags.Reference) !== 0
+                                && parameterDeclaration.dotDotDotToken)
+                            {
+                                // Determine type from elipsis array/tuple type
+                                const parameterTypeReference = (parameterType as ts.TypeReference);
+                                parameterType = parameterTypeReference.typeArguments[argumentIndex - declarationIndex];
+                            }
+                        }
                     }
-                    return checker.getTypeAtLocation(parentSignatureDeclaration.parameters[i]);
+                    return parameterType;
                 }
             }
 

diff --git a/test/unit/tuples.spec.ts b/test/unit/tuples.spec.ts
@@ -168,4 +168,90 @@ export class TupleTests {
         // Assert
         Expect(result).toBe(5);
     }
+
+    @Test("Tuple Return on Arrow Function")
+    public tupleReturnOnArrowFunction(): void {
+        const code =
+            `const fn = /** @tupleReturn */ (s: string) => [s, "bar"];
+            const [a, b] = fn("foo");
+            return a + b;`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
+
+    @Test("Tuple Return Inference")
+    public tupleReturnInference(): void {
+        const code =
+            `/** @tupleReturn */ interface Fn { (s: string): [string, string] }
+            const fn: Fn = s => [s, "bar"];
+            const [a, b] = fn("foo");
+            return a + b;`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
+
+    @Test("Tuple Return Inference as Argument")
+    public tupleReturnInferenceAsArgument(): void {
+        const code =
+            `/** @tupleReturn */ interface Fn { (s: string): [string, string] }
+            function foo(fn: Fn) {
+                const [a, b] = fn("foo");
+                return a + b;
+            }
+            return foo(s => [s, "bar"]);`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
+
+    @Test("Tuple Return Inference as Elipsis Argument")
+    public tupleReturnInferenceAsElipsisArgument(): void {
+        const code =
+            `/** @tupleReturn */ interface Fn { (s: string): [string, string] }
+            function foo(a: number, ...fn: Fn[]) {
+                const [a, b] = fn[0]("foo");
+                return a + b;
+            }
+            return foo(7, s => [s, "bar"]);`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
+
+    @Test("Tuple Return Inference as Elipsis Tuple Argument")
+    public tupleReturnInferenceAsElipsisTupleArgument(): void {
+        const code =
+            `/** @tupleReturn */ interface Fn { (s: string): [string, string] }
+            function foo(a: number, ...fn: [number, Fn]) {
+                const [a, b] = fn[1]("foo");
+                return a + b;
+            }
+            return foo(7, 17, s => [s, "bar"]);`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
+
+    @Test("Tuple Return in Spread")
+    public tupleReturnInSpread(): void {
+        const code =
+            `/** @tupleReturn */ function foo(): [string, string] {
+                return ["foo", "bar"];
+            }
+            function bar(a: string, b: string) {
+                return a + b;
+            }
+            return bar(...foo());`;
+        const lua = util.transpileString(code);
+        Expect(lua).not.toContain("unpack");
+        const result = util.executeLua(lua);
+        Expect(result).toBe("foobar");
+    }
 }