Fix optimized vararg spread with casts

GlassBricks · GlassBricks · commit b63bf86d27ba · 2021-12-19T12:27:45.000-08:00
Also fixes for other outer expressions.
diff --git a/src/transformation/visitors/errors.ts b/src/transformation/visitors/errors.ts
@@ -44,31 +44,38 @@ export const transformTryStatement: FunctionVisitor<ts.TryStatement> = (statemen
         const catchParameter = statement.catchClause.variableDeclaration
             ? transformIdentifier(context, statement.catchClause.variableDeclaration.name as ts.Identifier)
             : undefined;
-        const catchParameters = () => (catchParameter ? [lua.cloneIdentifier(catchParameter)] : []);
-
+        const catchFunction = lua.createFunctionExpression(
+            catchBlock,
+            catchParameter ? [lua.cloneIdentifier(catchParameter)] : []
+        );
         const catchIdentifier = lua.createIdentifier("____catch");
-        const catchFunction = lua.createFunctionExpression(catchBlock, catchParameters());
         result.push(lua.createVariableDeclarationStatement(catchIdentifier, catchFunction));
 
+        const hasReturn = tryScope.functionReturned ?? catchScope.functionReturned;
+
         const tryReturnIdentifiers = [tryResultIdentifier]; // ____try
-        if (returnedIdentifier) {
-            tryReturnIdentifiers.push(returnedIdentifier); // ____returned or catch variable
-            if (tryScope.functionReturned || catchScope.functionReturned) {
+        if (hasReturn || statement.catchClause.variableDeclaration) {
+            tryReturnIdentifiers.push(returnedIdentifier); // ____returned
+            if (hasReturn) {
                 tryReturnIdentifiers.push(returnValueIdentifier); // ____returnValue
                 returnCondition = lua.cloneIdentifier(returnedIdentifier);
             }
         }
         result.push(lua.createVariableDeclarationStatement(tryReturnIdentifiers, tryCall));
 
-        // Wrap catch in function if try or catch has return
-        const catchCall = lua.createCallExpression(catchIdentifier, [lua.cloneIdentifier(returnedIdentifier)]);
-        const catchAssign = lua.createAssignmentStatement(
-            [lua.cloneIdentifier(returnedIdentifier), lua.cloneIdentifier(returnValueIdentifier)],
-            catchCall
+        const catchCall = lua.createCallExpression(
+            catchIdentifier,
+            statement.catchClause.variableDeclaration ? [lua.cloneIdentifier(returnedIdentifier)] : []
         );
+        const catchCallStatement = hasReturn
+            ? lua.createAssignmentStatement(
+                  [lua.cloneIdentifier(returnedIdentifier), lua.cloneIdentifier(returnValueIdentifier)],
+                  catchCall
+              )
+            : lua.createExpressionStatement(catchCall);
 
         const notTryCondition = lua.createUnaryExpression(tryResultIdentifier, lua.SyntaxKind.NotOperator);
-        result.push(lua.createIfStatement(notTryCondition, lua.createBlock([catchAssign])));
+        result.push(lua.createIfStatement(notTryCondition, lua.createBlock([catchCallStatement])));
     } else if (tryScope.functionReturned) {
         // try with return, but no catch
         // returnedIdentifier = lua.createIdentifier("____returned");
diff --git a/src/transformation/visitors/spread.ts b/src/transformation/visitors/spread.ts
@@ -16,8 +16,15 @@ import { isMultiReturnCall } from "./language-extensions/multi";
 import { annotationRemoved } from "../utils/diagnostics";
 import { isGlobalVarargConstant } from "./language-extensions/vararg";
 
+function skipOuterExpressionParents(node: ts.Node) {
+    while (ts.isOuterExpression(node)) {
+        node = node.parent;
+    }
+    return node;
+}
+
 export function isOptimizedVarArgSpread(context: TransformationContext, symbol: ts.Symbol, identifier: ts.Identifier) {
-    if (!ts.isSpreadElement(identifier.parent)) {
+    if (!ts.isSpreadElement(skipOuterExpressionParents(identifier.parent))) {
         return false;
     }
 
@@ -63,20 +70,21 @@ export function isOptimizedVarArgSpread(context: TransformationContext, symbol:
 
 // TODO: Currently it's also used as an array member
 export const transformSpreadElement: FunctionVisitor<ts.SpreadElement> = (node, context) => {
-    if (ts.isIdentifier(node.expression)) {
-        if (isVarargType(context, node.expression)) {
+    const tsInnerExpression = ts.skipOuterExpressions(node.expression);
+    if (ts.isIdentifier(tsInnerExpression)) {
+        if (isVarargType(context, tsInnerExpression)) {
             context.diagnostics.push(annotationRemoved(node, AnnotationKind.Vararg));
         }
-        const symbol = context.checker.getSymbolAtLocation(node.expression);
-        if (symbol && isOptimizedVarArgSpread(context, symbol, node.expression)) {
+        const symbol = context.checker.getSymbolAtLocation(tsInnerExpression);
+        if (symbol && isOptimizedVarArgSpread(context, symbol, tsInnerExpression)) {
             return lua.createDotsLiteral(node);
         }
     }
 
     const innerExpression = context.transformExpression(node.expression);
-    if (isMultiReturnCall(context, node.expression)) return innerExpression;
+    if (isMultiReturnCall(context, tsInnerExpression)) return innerExpression;
 
-    const type = context.checker.getTypeAtLocation(node.expression);
+    const type = context.checker.getTypeAtLocation(node.expression); // not ts-inner expression, in case of casts
     if (isArrayType(context, type)) {
         return createUnpackCall(context, innerExpression, node);
     }
diff --git a/src/typescript-internal.d.ts b/src/typescript-internal.d.ts
@@ -1,5 +1,4 @@
 import { DiagnosticsProducingTypeChecker } from "./transformation/context";
-
 export {};
 
 declare module "typescript" {
@@ -31,6 +30,13 @@ declare module "typescript" {
 
     function transformJsx(context: TransformationContext): (x: SourceFile) => SourceFile;
 
-    function skipParentheses(node: Expression): Expression;
-    function skipParentheses(node: Node): Node;
+    export type OuterExpression =
+        | ParenthesizedExpression
+        | TypeAssertion
+        | AsExpression
+        | NonNullExpression
+        | PartiallyEmittedExpression;
+
+    function skipOuterExpressions(node: Expression, kinds?: OuterExpressionKinds): Expression;
+    export function isOuterExpression(node: Node, kinds?: OuterExpressionKinds): node is OuterExpression;
 }
diff --git a/test/unit/__snapshots__/spread.spec.ts.snap b/test/unit/__snapshots__/spread.spec.ts.snap
@@ -17,6 +17,21 @@ end
 return ____exports"
 `;
 
+exports[`vararg spread optimization With cast 1`] = `
+"local ____exports = {}
+function ____exports.__main(self)
+    local function pick(self, ...)
+        local args = {...}
+        return args[2]
+    end
+    local function test(self, ...)
+        return pick(nil, ...)
+    end
+    return test(nil, \\"a\\", \\"b\\", \\"c\\")
+end
+return ____exports"
+`;
+
 exports[`vararg spread optimization basic use 1`] = `
 "local ____exports = {}
 function ____exports.__main(self)
diff --git a/test/unit/spread.spec.ts b/test/unit/spread.spec.ts
@@ -265,6 +265,18 @@ describe("vararg spread optimization", () => {
             .expectLuaToMatchSnapshot()
             .expectToMatchJsResult();
     });
+
+    test("With cast", () => {
+        util.testFunction`
+            function pick(...args: any[]) { return args[1]; }
+            function test<F extends (...args: any)=>any>(...args: Parameters<F>) {
+                return pick(...(args as any[]));
+            }
+            return test<(...args: string[])=>void>("a", "b", "c");
+        `
+            .expectLuaToMatchSnapshot()
+            .expectToMatchJsResult();
+    });
 });
 
 describe("vararg spread de-optimization", () => {
