From 333651c689b48007bfc947b0c527db0d4a7b484c Mon Sep 17 00:00:00 2001
From: GlassBricks <24237065+GlassBricks@users.noreply.github.com>
Date: Wed, 10 Aug 2022 09:50:36 -0700
Subject: [PATCH] Allow MultiReturn to be typed as any

---
 .../visitors/language-extensions/multi.ts     |  6 +++-
 src/transformation/visitors/return.ts         | 22 +++++++++-----
 test/unit/language-extensions/multi.spec.ts   | 30 +++++++++++++++++++
 3 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/transformation/visitors/language-extensions/multi.ts b/src/transformation/visitors/language-extensions/multi.ts
index 5dd239deb..73dec6d9c 100644
--- a/src/transformation/visitors/language-extensions/multi.ts
+++ b/src/transformation/visitors/language-extensions/multi.ts
@@ -14,7 +14,11 @@ export function isMultiReturnType(type: ts.Type): boolean {
 }
 
 export function canBeMultiReturnType(type: ts.Type): boolean {
-    return isMultiReturnType(type) || (type.isUnion() && type.types.some(t => canBeMultiReturnType(t)));
+    return (
+        (type.flags & ts.TypeFlags.Any) !== 0 ||
+        isMultiReturnType(type) ||
+        (type.isUnion() && type.types.some(t => canBeMultiReturnType(t)))
+    );
 }
 
 export function isMultiFunctionCall(context: TransformationContext, expression: ts.CallExpression): boolean {
diff --git a/src/transformation/visitors/return.ts b/src/transformation/visitors/return.ts
index 13fceeecf..14d785166 100644
--- a/src/transformation/visitors/return.ts
+++ b/src/transformation/visitors/return.ts
@@ -23,16 +23,20 @@ function transformExpressionsInReturn(
 ): lua.Expression[] {
     const expressionType = context.checker.getTypeAtLocation(node);
 
-    if (ts.isCallExpression(node)) {
+    // skip type assertions
+    // don't skip parenthesis as it may arise confusion with lua behavior (where parenthesis are significant)
+    const innerNode = ts.skipOuterExpressions(node, ts.OuterExpressionKinds.Assertions);
+
+    if (ts.isCallExpression(innerNode)) {
         // $multi(...)
-        if (isMultiFunctionCall(context, node)) {
+        if (isMultiFunctionCall(context, innerNode)) {
             // Don't allow $multi to be implicitly cast to something other than LuaMultiReturn
             const type = context.checker.getContextualType(node);
             if (type && !canBeMultiReturnType(type)) {
-                context.diagnostics.push(invalidMultiFunctionReturnType(node));
+                context.diagnostics.push(invalidMultiFunctionReturnType(innerNode));
             }
 
-            let returnValues = transformArguments(context, node.arguments);
+            let returnValues = transformArguments(context, innerNode.arguments);
             if (insideTryCatch) {
                 returnValues = [wrapInTable(...returnValues)]; // Wrap results when returning inside try/catch
             }
@@ -40,12 +44,16 @@ function transformExpressionsInReturn(
         }
 
         // Force-wrap LuaMultiReturn when returning inside try/catch
-        if (insideTryCatch && returnsMultiType(context, node) && !shouldMultiReturnCallBeWrapped(context, node)) {
+        if (
+            insideTryCatch &&
+            returnsMultiType(context, innerNode) &&
+            !shouldMultiReturnCallBeWrapped(context, innerNode)
+        ) {
             return [wrapInTable(context.transformExpression(node))];
         }
-    } else if (isInMultiReturnFunction(context, node) && isMultiReturnType(expressionType)) {
+    } else if (isInMultiReturnFunction(context, innerNode) && isMultiReturnType(expressionType)) {
         // Unpack objects typed as LuaMultiReturn
-        return [createUnpackCall(context, context.transformExpression(node), node)];
+        return [createUnpackCall(context, context.transformExpression(innerNode), innerNode)];
     }
 
     return [context.transformExpression(node)];
diff --git a/test/unit/language-extensions/multi.spec.ts b/test/unit/language-extensions/multi.spec.ts
index 50784a3db..ed7213b2b 100644
--- a/test/unit/language-extensions/multi.spec.ts
+++ b/test/unit/language-extensions/multi.spec.ts
@@ -184,6 +184,36 @@ test("forward $multi call in ArrowFunction body", () => {
         .expectToEqual([1, 2]);
 });
 
+test("$multi call in function typed as any", () => {
+    util.testFunction`
+        function foo(): any { return $multi(1, 2); }
+        return foo()
+    `
+        .withLanguageExtensions()
+        .expectToHaveNoDiagnostics()
+        .expectToEqual(1);
+});
+
+test("$multi call cast to any", () => {
+    util.testFunction`
+        function foo() { return $multi(1, 2) as any; }
+        return foo()
+    `
+        .withLanguageExtensions()
+        .expectToHaveNoDiagnostics()
+        .expectToEqual(1);
+});
+
+test("$multi call cast to MultiReturn type", () => {
+    util.testFunction`
+        function foo() { return $multi(1, 2) as unknown as LuaMultiReturn<number[]>; }
+        return foo()
+    `
+        .withLanguageExtensions()
+        .expectToHaveNoDiagnostics()
+        .expectToEqual(1);
+});
+
 test.each(["0", "i"])("allow LuaMultiReturn numeric access (%s)", expression => {
     util.testFunction`
         ${multiFunction}