TypeScriptToLua · Perryvw · Feb 21, 2021 · Feb 21, 2021 · Feb 21, 2021
diff --git a/language-extensions/index.d.ts b/language-extensions/index.d.ts
@@ -430,3 +430,81 @@ declare type LuaLength<TOperand, TReturn> = ((operand: TOperand) => TReturn) & L
  * @param TReturn The resulting (return) type of the operation.
  */
 declare type LuaLengthMethod<TReturn> = (() => TReturn) & LuaExtension<"__luaLengthMethodBrand">;
+
+/**
+ * Calls to functions with this type are translated to `table[key]`.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TTable The type to access as a Lua table.
+ * @param TKey The type of the key to use to access the table.
+ * @param TValue The type of the value stored in the table.
+ */
+declare type LuaTableGet<TTable extends object, TKey extends {}, TValue> = ((table: TTable, key: TKey) => TValue) &
+    LuaExtension<"__luaTableGetBrand">;
+
+/**
+ * Calls to methods with this type are translated to `table[key]`, where `table` is the object with the method.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TKey The type of the key to use to access the table.
+ * @param TValue The type of the value stored in the table.
+ */
+declare type LuaTableGetMethod<TKey extends {}, TValue> = ((key: TKey) => TValue) &
+    LuaExtension<"__luaTableGetMethodBrand">;
+
+/**
+ * Calls to functions with this type are translated to `table[key] = value`.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TTable The type to access as a Lua table.
+ * @param TKey The type of the key to use to access the table.
+ * @param TValue The type of the value to assign to the table.
+ */
+declare type LuaTableSet<TTable extends object, TKey extends {}, TValue> = ((
+    table: TTable,
+    key: TKey,
+    value: TValue
+) => void) &
+    LuaExtension<"__luaTableSetBrand">;
+
+/**
+ * Calls to methods with this type are translated to `table[key] = value`, where `table` is the object with the method.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TKey The type of the key to use to access the table.
+ * @param TValue The type of the value to assign to the table.
+ */
+declare type LuaTableSetMethod<TKey extends {}, TValue> = ((key: TKey, value: TValue) => void) &
+    LuaExtension<"__luaTableSetMethodBrand">;
+
+/**
+ * A convenience type for working directly with a Lua table.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TKey The type of the keys used to access the table.
+ * @param TValue The type of the values stored in the table.
+ */
+declare interface LuaTable<TKey extends {} = {}, TValue = any> {
+    length: LuaLengthMethod<number>;
+    get: LuaTableGetMethod<TKey, TValue>;
+    set: LuaTableSetMethod<TKey, TValue>;
+}
+
+/**
+ * A convenience type for working directly with a Lua table.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TKey The type of the keys used to access the table.
+ * @param TValue The type of the values stored in the table.
+ */
+declare type LuaTableConstructor = (new <TKey extends {} = {}, TValue = any>() => LuaTable<TKey, TValue>) &
+    LuaExtension<"__luaTableNewBrand">;
+
+/**
+ * A convenience type for working directly with a Lua table.
+ * For more information see: https://typescripttolua.github.io/docs/advanced/language-extensions
+ *
+ * @param TKey The type of the keys used to access the table.
+ * @param TValue The type of the values stored in the table.
+ */
+declare const LuaTable: LuaTableConstructor;
diff --git a/src/transformation/utils/diagnostics.ts b/src/transformation/utils/diagnostics.ts
@@ -174,6 +174,14 @@ export const invalidOperatorMappingUse = createErrorDiagnosticFactory(
     "This function must always be directly called and cannot be referred to."
 );
 
+export const invalidTableExtensionUse = createErrorDiagnosticFactory(
+    "This function must be called directly and cannot be referred to."
+);
+
+export const invalidTableSetExpression = createErrorDiagnosticFactory(
+    "Table set extension can only be called as a stand-alone statement. It cannot be used as an expression in another statement."
+);
+
 export const annotationDeprecated = createWarningDiagnosticFactory(
     (kind: AnnotationKind) =>
         `'@${kind}' is deprecated and will be removed in a future update. Please update your code before upgrading to the next release, otherwise your project will no longer compile. ` +

diff --git a/src/transformation/utils/language-extensions.ts b/src/transformation/utils/language-extensions.ts
@@ -42,6 +42,11 @@ export enum ExtensionKind {
     BitwiseNotOperatorMethodType = "BitwiseNotOperatorMethodType",
     LengthOperatorType = "LengthOperatorType",
     LengthOperatorMethodType = "LengthOperatorMethodType",
+    TableNewType = "TableNewType",
+    TableGetType = "TableGetType",
+    TableGetMethodType = "TableGetMethodType",
+    TableSetType = "TableSetType",
+    TableSetMethodType = "TableSetMethodType",
 }
 
 const extensionKindToFunctionName: { [T in ExtensionKind]?: string } = {
@@ -90,6 +95,11 @@ const extensionKindToTypeBrand: { [T in ExtensionKind]: string } = {
     [ExtensionKind.BitwiseNotOperatorMethodType]: "__luaBitwiseNotMethodBrand",
     [ExtensionKind.LengthOperatorType]: "__luaLengthBrand",
     [ExtensionKind.LengthOperatorMethodType]: "__luaLengthMethodBrand",
+    [ExtensionKind.TableNewType]: "__luaTableNewBrand",
+    [ExtensionKind.TableGetType]: "__luaTableGetBrand",
+    [ExtensionKind.TableGetMethodType]: "__luaTableGetMethodBrand",
+    [ExtensionKind.TableSetType]: "__luaTableSetBrand",
+    [ExtensionKind.TableSetMethodType]: "__luaTableSetMethodBrand",
 };
 
 export function isExtensionType(type: ts.Type, extensionKind: ExtensionKind): boolean {

diff --git a/src/transformation/utils/typescript/index.ts b/src/transformation/utils/typescript/index.ts
@@ -79,3 +79,15 @@ export function getAllCallSignatures(type: ts.Type): readonly ts.Signature[] {
 export function isExpressionWithEvaluationEffect(node: ts.Expression): boolean {
     return !(ts.isLiteralExpression(node) || ts.isIdentifier(node) || node.kind === ts.SyntaxKind.ThisKeyword);
 }
+
+export function getFunctionTypeForCall(context: TransformationContext, node: ts.CallExpression) {
+    const signature = context.checker.getResolvedSignature(node);
+    if (!signature || !signature.declaration) {
+        return;
+    }
+    const typeDeclaration = findFirstNodeAbove(signature.declaration, ts.isTypeAliasDeclaration);
+    if (!typeDeclaration) {
+        return;
+    }
+    return context.checker.getTypeFromTypeNode(typeDeclaration.type);
+}
diff --git a/src/transformation/visitors/call.ts b/src/transformation/visitors/call.ts
@@ -13,6 +13,13 @@ import { transformElementAccessArgument } from "./access";
 import { transformLuaTableCallExpression } from "./lua-table";
 import { shouldMultiReturnCallBeWrapped, returnsMultiType } from "./language-extensions/multi";
 import { isOperatorMapping, transformOperatorMappingExpression } from "./language-extensions/operators";
+import {
+    isTableGetCall,
+    isTableSetCall,
+    transformTableGetExpression,
+    transformTableSetExpression,
+} from "./language-extensions/table";
+import { invalidTableSetExpression } from "../utils/diagnostics";
 
 export type PropertyCallExpression = ts.CallExpression & { expression: ts.PropertyAccessExpression };
 
@@ -214,6 +221,18 @@ export const transformCallExpression: FunctionVisitor<ts.CallExpression> = (node
         return transformOperatorMappingExpression(context, node);
     }
 
+    if (isTableGetCall(context, node)) {
+        return transformTableGetExpression(context, node);
+    }
+
+    if (isTableSetCall(context, node)) {
+        context.diagnostics.push(invalidTableSetExpression(node));
+        return createImmediatelyInvokedFunctionExpression(
+            [transformTableSetExpression(context, node)],
+            lua.createNilLiteral()
+        );
+    }
+
     if (ts.isPropertyAccessExpression(node.expression)) {
         const result = transformPropertyCall(context, node as PropertyCallExpression);
         return wrapResult ? wrapInTable(result) : result;

diff --git a/src/transformation/visitors/class/new.ts b/src/transformation/visitors/class/new.ts
@@ -5,6 +5,7 @@ import { AnnotationKind, getTypeAnnotations } from "../../utils/annotations";
 import { annotationInvalidArgumentCount, extensionCannotConstruct } from "../../utils/diagnostics";
 import { importLuaLibFeature, LuaLibFeature, transformLuaLibFunction } from "../../utils/lualib";
 import { transformArguments } from "../call";
+import { isTableNewCall } from "../language-extensions/table";
 import { transformLuaTableNewExpression } from "../lua-table";
 
 const builtinErrorTypeNames = new Set([
@@ -53,6 +54,10 @@ export const transformNewExpression: FunctionVisitor<ts.NewExpression> = (node,
         return luaTableResult;
     }
 
+    if (isTableNewCall(context, node)) {
+        return lua.createTableExpression(undefined, node);
+    }
+
     const name = context.transformExpression(node.expression);
     const signature = context.checker.getResolvedSignature(node);
     const params = node.arguments

diff --git a/src/transformation/visitors/expression-statement.ts b/src/transformation/visitors/expression-statement.ts
@@ -2,6 +2,7 @@ import * as ts from "typescript";
 import * as lua from "../../LuaAST";
 import { FunctionVisitor } from "../context";
 import { transformBinaryExpressionStatement } from "./binary-expression";
+import { isTableSetCall, transformTableSetExpression } from "./language-extensions/table";
 import { transformLuaTableExpressionStatement } from "./lua-table";
 import { transformUnaryExpressionStatement } from "./unary-expression";
 
@@ -11,6 +12,10 @@ export const transformExpressionStatement: FunctionVisitor<ts.ExpressionStatemen
         return luaTableResult;
     }
 
+    if (ts.isCallExpression(node.expression) && isTableSetCall(context, node.expression)) {
+        return transformTableSetExpression(context, node.expression);
+    }
+
     const unaryExpressionResult = transformUnaryExpressionStatement(context, node);
     if (unaryExpressionResult) {
         return unaryExpressionResult;

diff --git a/src/transformation/visitors/identifier.ts b/src/transformation/visitors/identifier.ts
@@ -8,6 +8,7 @@ import {
     invalidMultiFunctionUse,
     invalidOperatorMappingUse,
     invalidRangeUse,
+    invalidTableExtensionUse,
 } from "../utils/diagnostics";
 import { createExportedIdentifier, getSymbolExportScope } from "../utils/export";
 import { createSafeName, hasUnsafeIdentifierName } from "../utils/safe-names";
@@ -16,6 +17,7 @@ import { findFirstNodeAbove } from "../utils/typescript";
 import { isMultiFunctionNode } from "./language-extensions/multi";
 import { isOperatorMapping } from "./language-extensions/operators";
 import { isRangeFunctionNode } from "./language-extensions/range";
+import { isTableExtensionIdentifier } from "./language-extensions/table";
 
 export function transformIdentifier(context: TransformationContext, identifier: ts.Identifier): lua.Identifier {
     if (isMultiFunctionNode(context, identifier)) {
@@ -27,6 +29,10 @@ export function transformIdentifier(context: TransformationContext, identifier:
         context.diagnostics.push(invalidOperatorMappingUse(identifier));
     }
 
+    if (isTableExtensionIdentifier(context, identifier)) {
+        context.diagnostics.push(invalidTableExtensionUse(identifier));
+    }
+
     if (isRangeFunctionNode(context, identifier)) {
         context.diagnostics.push(invalidRangeUse(identifier));
         return lua.createAnonymousIdentifier(identifier);

diff --git a/src/transformation/visitors/language-extensions/operators.ts b/src/transformation/visitors/language-extensions/operators.ts
@@ -3,7 +3,7 @@ import * as lua from "../../../LuaAST";
 import { TransformationContext } from "../../context";
 import * as extensions from "../../utils/language-extensions";
 import { assert } from "../../../utils";
-import { findFirstNodeAbove } from "../../utils/typescript";
+import { getFunctionTypeForCall } from "../../utils/typescript";
 import { LuaTarget } from "../../../CompilerOptions";
 import { unsupportedForTarget } from "../../utils/diagnostics";
 
@@ -66,43 +66,17 @@ const bitwiseOperatorMapExtensions = new Set<extensions.ExtensionKind>([
     extensions.ExtensionKind.BitwiseNotOperatorMethodType,
 ]);
 
-function getTypeDeclaration(declaration: ts.Declaration) {
-    return ts.isTypeAliasDeclaration(declaration)
-        ? declaration
-        : findFirstNodeAbove(declaration, ts.isTypeAliasDeclaration);
-}
-
 function getOperatorMapExtensionKindForCall(context: TransformationContext, node: ts.CallExpression) {
-    const signature = context.checker.getResolvedSignature(node);
-    if (!signature || !signature.declaration) {
-        return;
-    }
-    const typeDeclaration = getTypeDeclaration(signature.declaration);
-    if (!typeDeclaration) {
-        return;
-    }
-    const type = context.checker.getTypeFromTypeNode(typeDeclaration.type);
-    return operatorMapExtensions.find(extensionKind => extensions.isExtensionType(type, extensionKind));
-}
-
-function isOperatorMapType(context: TransformationContext, type: ts.Type): boolean {
-    if (type.isUnionOrIntersection()) {
-        return type.types.some(t => isOperatorMapType(context, t));
-    } else {
-        return operatorMapExtensions.some(extensionKind => extensions.isExtensionType(type, extensionKind));
-    }
-}
-
-function isOperatorMapIdentifier(context: TransformationContext, node: ts.Identifier) {
-    const type = context.checker.getTypeAtLocation(node);
-    return isOperatorMapType(context, type);
+    const type = getFunctionTypeForCall(context, node);
+    return type && operatorMapExtensions.find(extensionKind => extensions.isExtensionType(type, extensionKind));
 }
 
 export function isOperatorMapping(context: TransformationContext, node: ts.CallExpression | ts.Identifier) {
     if (ts.isCallExpression(node)) {
         return getOperatorMapExtensionKindForCall(context, node) !== undefined;
     } else {
-        return isOperatorMapIdentifier(context, node);
+        const type = context.checker.getTypeAtLocation(node);
+        return operatorMapExtensions.some(extensionKind => extensions.isExtensionType(type, extensionKind));
     }
 }
 

diff --git a/src/transformation/visitors/language-extensions/table.ts b/src/transformation/visitors/language-extensions/table.ts
@@ -0,0 +1,80 @@
+import * as ts from "typescript";
+import * as lua from "../../../LuaAST";
+import { TransformationContext } from "../../context";
+import * as extensions from "../../utils/language-extensions";
+import { getFunctionTypeForCall } from "../../utils/typescript";
+import { assert } from "../../../utils";
+
+const tableGetExtensions = [extensions.ExtensionKind.TableGetType, extensions.ExtensionKind.TableGetMethodType];
+
+const tableSetExtensions = [extensions.ExtensionKind.TableSetType, extensions.ExtensionKind.TableSetMethodType];
+
+const tableExtensions = [extensions.ExtensionKind.TableNewType, ...tableGetExtensions, ...tableSetExtensions];
+
+function getTableExtensionKindForCall(
+    context: TransformationContext,
+    node: ts.CallExpression,
+    validExtensions: extensions.ExtensionKind[]
+) {
+    const type = getFunctionTypeForCall(context, node);
+    return type && validExtensions.find(extensionKind => extensions.isExtensionType(type, extensionKind));
+}
+
+export function isTableExtensionIdentifier(context: TransformationContext, node: ts.Identifier) {
+    const type = context.checker.getTypeAtLocation(node);
+    return tableExtensions.some(extensionKind => extensions.isExtensionType(type, extensionKind));
+}
+
+export function isTableGetCall(context: TransformationContext, node: ts.CallExpression) {
+    return getTableExtensionKindForCall(context, node, tableGetExtensions) !== undefined;
+}
+
+export function isTableSetCall(context: TransformationContext, node: ts.CallExpression) {
+    return getTableExtensionKindForCall(context, node, tableSetExtensions) !== undefined;
+}
+
+export function isTableNewCall(context: TransformationContext, node: ts.NewExpression) {
+    const type = context.checker.getTypeAtLocation(node.expression);
+    return extensions.isExtensionType(type, extensions.ExtensionKind.TableNewType);
+}
+
+export function transformTableGetExpression(context: TransformationContext, node: ts.CallExpression): lua.Expression {
+    const extensionKind = getTableExtensionKindForCall(context, node, tableGetExtensions);
+    assert(extensionKind);
+
+    const args = node.arguments.slice();
+    if (
+        args.length === 1 &&
+        (ts.isPropertyAccessExpression(node.expression) || ts.isElementAccessExpression(node.expression))
+    ) {
+        args.unshift(node.expression.expression);
+    }
+
+    return lua.createTableIndexExpression(
+        context.transformExpression(args[0]),
+        context.transformExpression(args[1]),
+        node
+    );
+}
+
+export function transformTableSetExpression(context: TransformationContext, node: ts.CallExpression): lua.Statement {
+    const extensionKind = getTableExtensionKindForCall(context, node, tableSetExtensions);
+    assert(extensionKind);
+
+    const args = node.arguments.slice();
+    if (
+        args.length === 2 &&
+        (ts.isPropertyAccessExpression(node.expression) || ts.isElementAccessExpression(node.expression))
+    ) {
+        args.unshift(node.expression.expression);
+    }
+
+    return lua.createAssignmentStatement(
+        lua.createTableIndexExpression(
+            context.transformExpression(args[0]),
+            context.transformExpression(args[1]),
+            node
+        ),
+        context.transformExpression(args[2])
+    );
+}