11import * as ts from "typescript" ;
22import * as lua from "../../LuaAST" ;
3- import { FunctionVisitor } from "../context" ;
3+ import { FunctionVisitor , TransformationContext } from "../context" ;
44import { isInTupleReturnFunction , isTupleReturnCall } from "../utils/annotations" ;
55import { validateAssignment } from "../utils/assignment-validation" ;
66import { createUnpackCall , wrapInTable } from "../utils/lua-ast" ;
77import { ScopeType , walkScopesUp } from "../utils/scope" ;
88import { isArrayType } from "../utils/typescript" ;
99
10+ function transformExpressionsInReturn (
11+ context : TransformationContext ,
12+ node : ts . Expression ,
13+ insideTryCatch : boolean
14+ ) : lua . Expression [ ] {
15+ if ( ! isInTupleReturnFunction ( context , node ) ) {
16+ return [ context . transformExpression ( node ) ] ;
17+ }
18+
19+ let results : lua . Expression [ ] ;
20+ const expressionType = context . checker . getTypeAtLocation ( node ) ;
21+
22+ // Parent function is a TupleReturn function
23+ if ( ts . isArrayLiteralExpression ( node ) ) {
24+ // If return expression is an array literal, leave out brackets.
25+ results = node . elements . map ( e => context . transformExpression ( e ) ) ;
26+ } else if ( ! isTupleReturnCall ( context , node ) && isArrayType ( context , expressionType ) ) {
27+ // If return expression is an array-type and not another TupleReturn call, unpack it
28+ results = [ createUnpackCall ( context , context . transformExpression ( node ) , node ) ] ;
29+ } else {
30+ results = [ context . transformExpression ( node ) ] ;
31+ }
32+
33+ // Wrap tupleReturn results when returning inside try/catch
34+ if ( insideTryCatch ) {
35+ results = [ wrapInTable ( ...results ) ] ;
36+ }
37+
38+ return results ;
39+ }
40+
41+ export function transformExpressionBodyToReturnStatement (
42+ context : TransformationContext ,
43+ node : ts . Expression
44+ ) : lua . Statement {
45+ const expressions = transformExpressionsInReturn ( context , node , false ) ;
46+ return lua . createReturnStatement ( expressions , node ) ;
47+ }
48+
1049export const transformReturnStatement : FunctionVisitor < ts . ReturnStatement > = ( statement , context ) => {
1150 // Bubble up explicit return flag and check if we're inside a try/catch block
1251 let insideTryCatch = false ;
@@ -29,27 +68,7 @@ export const transformReturnStatement: FunctionVisitor<ts.ReturnStatement> = (st
2968 validateAssignment ( context , statement , expressionType , returnType ) ;
3069 }
3170
32- if ( isInTupleReturnFunction ( context , statement ) ) {
33- // Parent function is a TupleReturn function
34- if ( ts . isArrayLiteralExpression ( statement . expression ) ) {
35- // If return expression is an array literal, leave out brackets.
36- results = statement . expression . elements . map ( e => context . transformExpression ( e ) ) ;
37- } else if ( ! isTupleReturnCall ( context , statement . expression ) && isArrayType ( context , expressionType ) ) {
38- // If return expression is an array-type and not another TupleReturn call, unpack it
39- results = [
40- createUnpackCall ( context , context . transformExpression ( statement . expression ) , statement . expression ) ,
41- ] ;
42- } else {
43- results = [ context . transformExpression ( statement . expression ) ] ;
44- }
45-
46- // Wrap tupleReturn results when returning inside try/catch
47- if ( insideTryCatch ) {
48- results = [ wrapInTable ( ...results ) ] ;
49- }
50- } else {
51- results = [ context . transformExpression ( statement . expression ) ] ;
52- }
71+ results = transformExpressionsInReturn ( context , statement . expression , insideTryCatch ) ;
5372 } else {
5473 // Empty return
5574 results = [ ] ;
0 commit comments