Skip to content

Commit 9b558d2

Browse files
authored
tupleReturn improvements (#445)
* tupleReturn improvements - fixed `@tupleReturn` on arrow functions - enabled inference of tupleReturn on arrow functions - preventing wrap+unpack when using spread operator on tupleReturn function * added tests for argument inference * better name
1 parent 68f3681 commit 9b558d2

File tree

3 files changed

+138
-13
lines changed

3 files changed

+138
-13
lines changed

src/LuaTransformer.ts

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2972,7 +2972,15 @@ export class LuaTransformer {
29722972
hasContext ? context : undefined
29732973
);
29742974

2975-
const body = ts.isBlock(node.body) ? node.body : ts.createBlock([ts.createReturn(node.body)]);
2975+
let body: ts.Block;
2976+
if (ts.isBlock(node.body)) {
2977+
body = node.body;
2978+
} else {
2979+
const returnExpression = ts.createReturn(node.body);
2980+
body = ts.createBlock([returnExpression]);
2981+
returnExpression.parent = body;
2982+
body.parent = node.body.parent;
2983+
}
29762984
const [transformedBody] = this.transformFunctionBody(node.parameters, body, spreadIdentifier);
29772985

29782986
return tstl.createFunctionExpression(
@@ -3053,9 +3061,10 @@ export class LuaTransformer {
30533061
const isTupleReturnForward =
30543062
node.parent && ts.isReturnStatement(node.parent) && tsHelper.isInTupleReturnFunction(node, this.checker);
30553063
const isInDestructingAssignment = tsHelper.isInDestructingAssignment(node);
3064+
const isInSpread = node.parent && ts.isSpreadElement(node.parent);
30563065
const returnValueIsUsed = node.parent && !ts.isExpressionStatement(node.parent);
3057-
const wrapResult = isTupleReturn && !isTupleReturnForward&& !isInDestructingAssignment
3058-
&& returnValueIsUsed && !isLuaIterator;
3066+
const wrapResult = isTupleReturn && !isTupleReturnForward && !isInDestructingAssignment
3067+
&& !isInSpread && returnValueIsUsed && !isLuaIterator;
30593068

30603069
if (ts.isPropertyAccessExpression(node.expression)) {
30613070
const result = this.transformPropertyCall(node);
@@ -3707,8 +3716,11 @@ export class LuaTransformer {
37073716

37083717
public transformSpreadElement(expression: ts.SpreadElement): ExpressionVisitResult {
37093718
const innerExpression = this.transformExpression(expression.expression);
3710-
3711-
return this.createUnpackCall(innerExpression, expression);
3719+
if (tsHelper.isTupleReturnCall(expression.expression, this.checker)) {
3720+
return innerExpression;
3721+
} else {
3722+
return this.createUnpackCall(innerExpression, expression);
3723+
}
37123724
}
37133725

37143726
public transformStringLiteral(literal: ts.StringLiteralLike): tstl.StringLiteral {

src/TSHelper.ts

Lines changed: 35 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,13 @@ export class TSHelper {
187187
public static isInTupleReturnFunction(node: ts.Node, checker: ts.TypeChecker): boolean {
188188
const declaration = TSHelper.findFirstNodeAbove(node, ts.isFunctionLike);
189189
if (declaration) {
190-
const decorators = TSHelper.getCustomDecorators(checker.getTypeAtLocation(declaration), checker);
190+
let functionType: ts.Type;
191+
if (ts.isFunctionExpression(declaration) || ts.isArrowFunction(declaration)) {
192+
functionType = TSHelper.inferAssignedType(declaration, checker);
193+
} else {
194+
functionType = checker.getTypeAtLocation(declaration);
195+
}
196+
const decorators = TSHelper.getCustomDecorators(functionType, checker);
191197
return decorators.has(DecoratorKind.TupleReturn);
192198
} else {
193199
return false;
@@ -446,15 +452,36 @@ export class TSHelper {
446452

447453
} else if (ts.isCallExpression(expression.parent)) {
448454
// Expression being passed as argument to a function
449-
let i = expression.parent.arguments.indexOf(expression);
450-
if (i >= 0) {
455+
const argumentIndex = expression.parent.arguments.indexOf(expression);
456+
if (argumentIndex >= 0) {
451457
const parentSignature = checker.getResolvedSignature(expression.parent);
452-
const parentSignatureDeclaration = parentSignature.getDeclaration();
453-
if (parentSignatureDeclaration) {
454-
if (this.getExplicitThisParameter(parentSignatureDeclaration)) {
455-
++i;
458+
if (parentSignature.parameters.length > 0) { // In case function type is 'any'
459+
const signatureIndex = Math.min(argumentIndex, parentSignature.parameters.length - 1);
460+
let parameterType = checker.getTypeOfSymbolAtLocation(
461+
parentSignature.parameters[signatureIndex],
462+
expression
463+
);
464+
if (TSHelper.isArrayType(parameterType, checker)) {
465+
// Check for elipses argument
466+
const parentSignatureDeclaration = parentSignature.getDeclaration();
467+
if (parentSignatureDeclaration) {
468+
let declarationIndex = signatureIndex;
469+
if (this.getExplicitThisParameter(parentSignatureDeclaration)) {
470+
// Ignore 'this' parameter
471+
++declarationIndex;
472+
}
473+
const parameterDeclaration = parentSignatureDeclaration.parameters[declarationIndex];
474+
if ((parameterType.flags & ts.TypeFlags.Object) !== 0
475+
&& ((parameterType as ts.ObjectType).objectFlags & ts.ObjectFlags.Reference) !== 0
476+
&& parameterDeclaration.dotDotDotToken)
477+
{
478+
// Determine type from elipsis array/tuple type
479+
const parameterTypeReference = (parameterType as ts.TypeReference);
480+
parameterType = parameterTypeReference.typeArguments[argumentIndex - declarationIndex];
481+
}
482+
}
456483
}
457-
return checker.getTypeAtLocation(parentSignatureDeclaration.parameters[i]);
484+
return parameterType;
458485
}
459486
}
460487

test/unit/tuples.spec.ts

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,4 +168,90 @@ export class TupleTests {
168168
// Assert
169169
Expect(result).toBe(5);
170170
}
171+
172+
@Test("Tuple Return on Arrow Function")
173+
public tupleReturnOnArrowFunction(): void {
174+
const code =
175+
`const fn = /** @tupleReturn */ (s: string) => [s, "bar"];
176+
const [a, b] = fn("foo");
177+
return a + b;`;
178+
const lua = util.transpileString(code);
179+
Expect(lua).not.toContain("unpack");
180+
const result = util.executeLua(lua);
181+
Expect(result).toBe("foobar");
182+
}
183+
184+
@Test("Tuple Return Inference")
185+
public tupleReturnInference(): void {
186+
const code =
187+
`/** @tupleReturn */ interface Fn { (s: string): [string, string] }
188+
const fn: Fn = s => [s, "bar"];
189+
const [a, b] = fn("foo");
190+
return a + b;`;
191+
const lua = util.transpileString(code);
192+
Expect(lua).not.toContain("unpack");
193+
const result = util.executeLua(lua);
194+
Expect(result).toBe("foobar");
195+
}
196+
197+
@Test("Tuple Return Inference as Argument")
198+
public tupleReturnInferenceAsArgument(): void {
199+
const code =
200+
`/** @tupleReturn */ interface Fn { (s: string): [string, string] }
201+
function foo(fn: Fn) {
202+
const [a, b] = fn("foo");
203+
return a + b;
204+
}
205+
return foo(s => [s, "bar"]);`;
206+
const lua = util.transpileString(code);
207+
Expect(lua).not.toContain("unpack");
208+
const result = util.executeLua(lua);
209+
Expect(result).toBe("foobar");
210+
}
211+
212+
@Test("Tuple Return Inference as Elipsis Argument")
213+
public tupleReturnInferenceAsElipsisArgument(): void {
214+
const code =
215+
`/** @tupleReturn */ interface Fn { (s: string): [string, string] }
216+
function foo(a: number, ...fn: Fn[]) {
217+
const [a, b] = fn[0]("foo");
218+
return a + b;
219+
}
220+
return foo(7, s => [s, "bar"]);`;
221+
const lua = util.transpileString(code);
222+
Expect(lua).not.toContain("unpack");
223+
const result = util.executeLua(lua);
224+
Expect(result).toBe("foobar");
225+
}
226+
227+
@Test("Tuple Return Inference as Elipsis Tuple Argument")
228+
public tupleReturnInferenceAsElipsisTupleArgument(): void {
229+
const code =
230+
`/** @tupleReturn */ interface Fn { (s: string): [string, string] }
231+
function foo(a: number, ...fn: [number, Fn]) {
232+
const [a, b] = fn[1]("foo");
233+
return a + b;
234+
}
235+
return foo(7, 17, s => [s, "bar"]);`;
236+
const lua = util.transpileString(code);
237+
Expect(lua).not.toContain("unpack");
238+
const result = util.executeLua(lua);
239+
Expect(result).toBe("foobar");
240+
}
241+
242+
@Test("Tuple Return in Spread")
243+
public tupleReturnInSpread(): void {
244+
const code =
245+
`/** @tupleReturn */ function foo(): [string, string] {
246+
return ["foo", "bar"];
247+
}
248+
function bar(a: string, b: string) {
249+
return a + b;
250+
}
251+
return bar(...foo());`;
252+
const lua = util.transpileString(code);
253+
Expect(lua).not.toContain("unpack");
254+
const result = util.executeLua(lua);
255+
Expect(result).toBe("foobar");
256+
}
171257
}

0 commit comments

Comments
 (0)