From f149692a4bb1640f48650cd8980f16fa26bd77e6 Mon Sep 17 00:00:00 2001
From: hazzard993 <hazzard993@gmail.com>
Date: Mon, 4 Jan 2021 20:21:55 +1000
Subject: [PATCH] Allow numeric access to MultiReturn

---
 src/transformation/utils/diagnostics.ts       |  4 +++
 src/transformation/visitors/access.ts         | 21 ++++++++++++-
 .../__snapshots__/multi.spec.ts.snap          | 31 +++++++++++++++++++
 test/unit/language-extensions/multi.spec.ts   | 20 ++++++++++++
 4 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/src/transformation/utils/diagnostics.ts b/src/transformation/utils/diagnostics.ts
index 8f69eadd9..d7b7b6467 100644
--- a/src/transformation/utils/diagnostics.ts
+++ b/src/transformation/utils/diagnostics.ts
@@ -164,6 +164,10 @@ export const invalidMultiTypeArrayLiteralElementInitializer = createErrorDiagnos
     "This array literal pattern cannot have initializers."
 );
 
+export const invalidMultiReturnAccess = createErrorDiagnosticFactory(
+    "The MultiReturn type can only be accessed via an element access expression of a numeric type."
+);
+
 export const unsupportedMultiFunctionAssignment = createErrorDiagnosticFactory(
     "Omitted expressions and BindingElements are expected here."
 );
diff --git a/src/transformation/visitors/access.ts b/src/transformation/visitors/access.ts
index eaa5591b7..6cf61a62c 100644
--- a/src/transformation/visitors/access.ts
+++ b/src/transformation/visitors/access.ts
@@ -3,10 +3,12 @@ import * as lua from "../../LuaAST";
 import { transformBuiltinPropertyAccessExpression } from "../builtins";
 import { FunctionVisitor, TransformationContext } from "../context";
 import { AnnotationKind, getTypeAnnotations } from "../utils/annotations";
+import { invalidMultiReturnAccess } from "../utils/diagnostics";
 import { addToNumericExpression } from "../utils/lua-ast";
 import { LuaLibFeature, transformLuaLibFunction } from "../utils/lualib";
 import { isArrayType, isNumberType, isStringType } from "../utils/typescript";
 import { tryGetConstEnumValue } from "./enum";
+import { returnsMultiType } from "./language-extensions/multi";
 import { transformLuaTablePropertyAccessExpression, validateLuaTableElementAccessExpression } from "./lua-table";
 
 export function transformElementAccessArgument(
@@ -41,7 +43,20 @@ export const transformElementAccessExpression: FunctionVisitor<ts.ElementAccessE
         return transformLuaLibFunction(context, LuaLibFeature.StringAccess, node, table, index);
     }
 
-    return lua.createTableIndexExpression(table, transformElementAccessArgument(context, node), node);
+    const accessExpression = transformElementAccessArgument(context, node);
+
+    if (ts.isCallExpression(node.expression) && returnsMultiType(context, node.expression)) {
+        const accessType = context.checker.getTypeAtLocation(node.argumentExpression);
+        if (!isNumberType(context, accessType)) {
+            context.diagnostics.push(invalidMultiReturnAccess(node));
+        }
+
+        const selectIdentifier = lua.createIdentifier("select");
+        const selectCall = lua.createCallExpression(selectIdentifier, [accessExpression, table]);
+        return selectCall;
+    }
+
+    return lua.createTableIndexExpression(table, accessExpression, node);
 };
 
 export const transformPropertyAccessExpression: FunctionVisitor<ts.PropertyAccessExpression> = (
@@ -63,6 +78,10 @@ export const transformPropertyAccessExpression: FunctionVisitor<ts.PropertyAcces
         return builtinResult;
     }
 
+    if (ts.isCallExpression(expression.expression) && returnsMultiType(context, expression.expression)) {
+        context.diagnostics.push(invalidMultiReturnAccess(expression));
+    }
+
     const property = expression.name.text;
     const type = context.checker.getTypeAtLocation(expression.expression);
 
diff --git a/test/unit/language-extensions/__snapshots__/multi.spec.ts.snap b/test/unit/language-extensions/__snapshots__/multi.spec.ts.snap
index 587cdfd90..b745edf27 100644
--- a/test/unit/language-extensions/__snapshots__/multi.spec.ts.snap
+++ b/test/unit/language-extensions/__snapshots__/multi.spec.ts.snap
@@ -1,5 +1,36 @@
 // Jest Snapshot v1, https://goo.gl/fbAQLP
 
+exports[`disallow MultiReturn non-numeric access: code 1`] = `
+"local ____exports = {}
+function ____exports.__main(self)
+    local function multi(self, ...)
+        local args = {...}
+        return table.unpack(args)
+    end
+    return select(
+        \\"forEach\\",
+        multi(nil)
+    )
+end
+return ____exports"
+`;
+
+exports[`disallow MultiReturn non-numeric access: code 2`] = `
+"local ____exports = {}
+function ____exports.__main(self)
+    local function multi(self, ...)
+        local args = {...}
+        return table.unpack(args)
+    end
+    return multi(nil).forEach
+end
+return ____exports"
+`;
+
+exports[`disallow MultiReturn non-numeric access: diagnostics 1`] = `"main.ts(7,16): error TSTL: The MultiReturn type can only be accessed via an element access expression of a numeric type."`;
+
+exports[`disallow MultiReturn non-numeric access: diagnostics 2`] = `"main.ts(7,16): error TSTL: The MultiReturn type can only be accessed via an element access expression of a numeric type."`;
+
 exports[`invalid $multi call ($multi()): code 1`] = `"____(_G)"`;
 
 exports[`invalid $multi call ($multi()): diagnostics 1`] = `"main.ts(2,9): error TSTL: The $multi function must be called in an expression that is returned."`;
diff --git a/test/unit/language-extensions/multi.spec.ts b/test/unit/language-extensions/multi.spec.ts
index ed5353bbc..c697cbad0 100644
--- a/test/unit/language-extensions/multi.spec.ts
+++ b/test/unit/language-extensions/multi.spec.ts
@@ -5,6 +5,7 @@ import {
     invalidMultiFunctionUse,
     invalidMultiTypeToNonArrayBindingPattern,
     invalidMultiTypeArrayBindingPatternElementInitializer,
+    invalidMultiReturnAccess,
 } from "../../../src/transformation/utils/diagnostics";
 
 const multiProjectOptions: tstl.CompilerOptions = {
@@ -126,3 +127,22 @@ test("allow $multi call in ArrowFunction body", () => {
         .setOptions(multiProjectOptions)
         .expectToEqual(1);
 });
+
+test.each(["0", "i"])("allow MultiReturn numeric access", expression => {
+    util.testFunction`
+        ${multiFunction}
+        const i = 0;
+        return multi(1)[${expression}];
+    `
+        .setOptions(multiProjectOptions)
+        .expectToEqual(1);
+});
+
+test.each(["multi()['forEach']", "multi().forEach"])("disallow MultiReturn non-numeric access", expression => {
+    util.testFunction`
+        ${multiFunction}
+        return ${expression};
+    `
+        .setOptions(multiProjectOptions)
+        .expectDiagnosticsToMatchSnapshot([invalidMultiReturnAccess.code]);
+});