TypeScriptToLua · Perryvw · Feb 11, 2019 · Feb 10, 2019 · Feb 10, 2019 · Feb 10, 2019
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -986,47 +986,71 @@ export class LuaTransformer {
                     ));
                 }
             } else {
-                const table: tstl.IdentifierOrTableIndexExpression  =
-                    this.transformIdentifierExpression(enumDeclaration.name);
-                const property = tstl.createTableIndexExpression(table, memberName, undefined);
+                const enumTable = this.transformIdentifierExpression(enumDeclaration.name);
+                const property = tstl.createTableIndexExpression(enumTable, memberName);
                 result.push(tstl.createAssignmentStatement(property, enumMember.value, enumMember.original));
+
+                const valueIndex = tstl.createTableIndexExpression(enumTable, enumMember.value);
+                result.push(tstl.createAssignmentStatement(valueIndex, memberName, enumMember.original));
             }
         }
 
         return result;
     }
 
     public computeEnumMembers(node: ts.EnumDeclaration):
-        Array<{name: ts.PropertyName, value: tstl.NumericLiteral | tstl.StringLiteral, original: ts.Node}> {
+        Array<{name: ts.PropertyName, value: tstl.Expression, original: ts.Node}> {
         let numericValue = 0;
         let hasStringInitializers = false;
 
+        const valueMap = new Map<ts.PropertyName, tstl.Expression>();
+
         return node.members.map(member => {
-            let valueLiteral: tstl.NumericLiteral | tstl.StringLiteral;
+            let valueExpression: tstl.Expression;
             if (member.initializer) {
-                if (ts.isNumericLiteral(member.initializer)) {
+                if (ts.isNumericLiteral(member.initializer))
+                {
                     numericValue = Number(member.initializer.text);
-                    valueLiteral = tstl.createNumericLiteral(numericValue);
-                } else if (ts.isStringLiteral(member.initializer)) {
+                    valueExpression = this.transformNumericLiteral(member.initializer);
+                    numericValue++;
+                }
+                else if (ts.isStringLiteral(member.initializer))
+                {
                     hasStringInitializers = true;
-                    valueLiteral = tstl.createStringLiteral(member.initializer.text);
-                } else {
-                    throw TSTLErrors.InvalidEnumMember(member.initializer);
+                    valueExpression = this.transformStringLiteral(member.initializer);
                 }
-            } else if (hasStringInitializers) {
+                else
+                {
+                    if (ts.isIdentifier(member.initializer)) {
+                        const [isEnumMember, originalName] = tsHelper.isEnumMember(node, member.initializer);
+                        if (isEnumMember) {
+                            valueExpression = valueMap.get(originalName);
+                        } else {
+                            valueExpression = this.transformExpression(member.initializer);
+                        }
+                    } else {
+                        valueExpression = this.transformExpression(member.initializer);
+                    }
+                }
+            }
+            else if (hasStringInitializers)
+            {
                 throw TSTLErrors.HeterogeneousEnum(node);
-            } else {
-                valueLiteral = tstl.createNumericLiteral(numericValue);
             }
+            else
+            {
+                valueExpression = tstl.createNumericLiteral(numericValue);
+                numericValue++;
+            }
+
+            valueMap.set(member.name, valueExpression);
 
             const enumMember = {
                 name: member.name,
                 original: member,
-                value: valueLiteral,
+                value: valueExpression,
             };
 
-            numericValue++;
-
             return enumMember;
         });
     }
@@ -2725,26 +2749,7 @@ export class LuaTransformer {
         }
 
         if (type.symbol && (type.symbol.flags & ts.SymbolFlags.ConstEnum)) {
-            const propertyValueDeclaration = this.checker.getTypeAtLocation(node).symbol.valueDeclaration;
-
-            if (propertyValueDeclaration && propertyValueDeclaration.kind === ts.SyntaxKind.EnumMember) {
-                const enumMember = propertyValueDeclaration as ts.EnumMember;
-
-                if (enumMember.initializer) {
-                    return this.transformExpression(enumMember.initializer);
-                } else {
-                    const enumMembers = this.computeEnumMembers(enumMember.parent);
-                    const memberPosition = enumMember.parent.members.indexOf(enumMember);
-
-                    if (memberPosition === -1) {
-                        throw TSTLErrors.UnsupportedProperty(type.symbol.name, property, node);
-                    }
-
-                    const value = tstl.cloneNode(enumMembers[memberPosition].value);
-                    tstl.setNodeOriginal(value, enumMember);
-                    return value;
-                }
-            }
+            return this.transformConstEnumValue(type, property, node);
         }
 
         this.checkForLuaLibType(type);
@@ -2840,6 +2845,13 @@ export class LuaTransformer {
         const index = this.transformExpression(node.argumentExpression);
 
         const type = this.checker.getTypeAtLocation(node.expression);
+
+        if (type.symbol && (type.symbol.flags & ts.SymbolFlags.ConstEnum)
+            && ts.isStringLiteral(node.argumentExpression))
+        {
+            return this.transformConstEnumValue(type, node.argumentExpression.text, node);
+        }
+
         if (tsHelper.isArrayType(type, this.checker)) {
             return tstl.createTableIndexExpression(table, this.expressionPlusOne(index), node);
         } else if (tsHelper.isStringType(type)) {
@@ -2853,6 +2865,45 @@ export class LuaTransformer {
         }
     }
 
+    private transformConstEnumValue(enumType: ts.EnumType, memberName: string, tsOriginal: ts.Node): tstl.Expression {
+        // Assumption: the enum only has one declaration
+        const enumDeclaration = enumType.symbol.declarations.find(d => ts.isEnumDeclaration(d)) as ts.EnumDeclaration;
+        const enumMember = enumDeclaration.members
+            .find(m => ts.isIdentifier(m.name) && m.name.text === memberName);
+
+        if (enumMember) {
+            if (enumMember.initializer) {
+                if (ts.isIdentifier(enumMember.initializer)) {
+                    const [isEnumMember, valueName] = tsHelper.isEnumMember(enumDeclaration, enumMember.initializer);
+                    if (isEnumMember) {
+                        if (ts.isIdentifier(valueName)) {
+                            return this.transformConstEnumValue(enumType, valueName.text, tsOriginal);
+                        }
+                    } else {
+                        return tstl.setNodeOriginal(this.transformExpression(enumMember.initializer), tsOriginal);
+                    }
+                } else {
+                    return tstl.setNodeOriginal(this.transformExpression(enumMember.initializer), tsOriginal);
+                }
+            } else {
+                let enumValue = 0;
+                for (const member of enumDeclaration.members) {
+                    if (member === enumMember) {
+                        return tstl.createNumericLiteral(enumValue, tsOriginal);
+                    }
+                    if (member.initializer === undefined) {
+                        enumValue++;
+                    } else if (ts.isNumericLiteral(member.initializer)) {
+                        enumValue = Number(member.initializer.text) + 1;
+                    }
+                }
+
+                throw TSTLErrors.CouldNotFindEnumMember(enumDeclaration, memberName, tsOriginal);
+            }
+        }
+        throw TSTLErrors.CouldNotFindEnumMember(enumDeclaration, memberName, tsOriginal);
+    }
+
     public transformStringCallExpression(node: ts.CallExpression): tstl.Expression {
         const expression = node.expression as ts.PropertyAccessExpression;
         const params = this.transformArguments(node.arguments);

diff --git a/src/TSHelper.ts b/src/TSHelper.ts
@@ -625,4 +625,21 @@ export class TSHelper {
         const firstDeclaration = this.getFirstDeclaration(symbol);
         return firstDeclaration === node;
     }
+
+    public static isEnumMember(enumDeclaration: ts.EnumDeclaration, value: ts.Expression): [boolean, ts.PropertyName] {
+        if (ts.isIdentifier(value)) {
+            const enumMember = enumDeclaration.members.find(m => ts.isIdentifier(m.name) && m.name.text === value.text);
+            if (enumMember !== undefined) {
+                if (enumMember.initializer && ts.isIdentifier(enumMember.initializer)) {
+                    return this.isEnumMember(enumDeclaration, enumMember.initializer);
+                } else {
+                    return [true, enumMember.name];
+                }
+            } else {
+                return [false, undefined];
+            }
+        } else {
+            return [false, undefined];
+        }
+    }
 }
diff --git a/src/TSTLErrors.ts b/src/TSTLErrors.ts
@@ -4,6 +4,11 @@ import {TranspileError} from "./TranspileError";
 import {TSHelper as tsHelper} from "./TSHelper";
 
 export class TSTLErrors {
+    public static CouldNotFindEnumMember =
+        (enumDeclaration: ts.EnumDeclaration, enumMember: string, node: ts.Node) => new TranspileError(
+            `Could not find ${enumMember} in ${enumDeclaration.name.text}`, node
+        );
+
     public static DefaultImportsNotSupported = (node: ts.Node) =>
         new TranspileError(`Default Imports are not supported, please use named imports instead!`, node);
 

diff --git a/test/translation/lua/enum.lua b/test/translation/lua/enum.lua
@@ -1,4 +1,7 @@
 TestEnum = {};
 TestEnum.val1 = 0;
+TestEnum[0] = "val1";
 TestEnum.val2 = 2;
+TestEnum[2] = "val2";
 TestEnum.val3 = 3;
+TestEnum[3] = "val3";
diff --git a/test/translation/lua/enumHeterogeneous.lua b/test/translation/lua/enumHeterogeneous.lua
@@ -1,4 +1,7 @@
 TestEnum = {};
 TestEnum.val1 = 0;
+TestEnum[0] = "val1";
 TestEnum.val2 = 3;
+TestEnum[3] = "val2";
 TestEnum.val3 = "baz";
+TestEnum.baz = "val3";
diff --git a/test/translation/lua/enumString.lua b/test/translation/lua/enumString.lua
@@ -1,4 +1,7 @@
 TestEnum = {};
 TestEnum.val1 = "foo";
+TestEnum.foo = "val1";
 TestEnum.val2 = "bar";
+TestEnum.bar = "val2";
 TestEnum.val3 = "baz";
+TestEnum.baz = "val3";
diff --git a/test/translation/lua/modulesNamespaceExportEnum.lua b/test/translation/lua/modulesNamespaceExportEnum.lua
@@ -4,6 +4,8 @@ local test = exports.test;
 do
     test.TestEnum = {};
     test.TestEnum.foo = "foo";
+    test.TestEnum.foo = "foo";
+    test.TestEnum.bar = "bar";
     test.TestEnum.bar = "bar";
 end
 return exports;
diff --git a/test/unit/enum.spec.ts b/test/unit/enum.spec.ts
@@ -26,10 +26,10 @@ export class EnumTests {
                 MEMBER_TWO = "test2"
             }
 
-            const valueOne = TestEnum.MEMBER_ONE;
+            const valueOne = TestEnum.MEMBER_TWO;
         `;
 
-        Expect(util.transpileString(testCode)).toBe(`local valueOne = "test";`);
+        Expect(util.transpileString(testCode)).toBe(`local valueOne = "test2";`);
     }
 
     @Test("Const enum without initializer")
@@ -40,10 +40,10 @@ export class EnumTests {
                 MEMBER_TWO
             }
 
-            const valueOne = TestEnum.MEMBER_ONE;
+            const valueOne = TestEnum.MEMBER_TWO;
         `;
 
-        Expect(util.transpileString(testCode)).toBe(`local valueOne = 0;`);
+        Expect(util.transpileString(testCode)).toBe(`local valueOne = 1;`);
     }
 
     @Test("Const enum without initializer in some values")
@@ -76,20 +76,6 @@ export class EnumTests {
                         + "member values, or specify values (of the same type) for all members.");
     }
 
-    @Test("Unsuported enum")
-    public unsuportedEnum(): void {
-        // Transpile & Assert
-        Expect(() => {
-            const lua = util.transpileString(
-                `enum TestEnum {
-                    val1 = [],
-                    val2 = "ok",
-                    val3 = "bye"
-                }`
-            );
-        }).toThrowError(TranspileError, "Only numeric or string initializers allowed for enums.");
-    }
-
     @Test("String literal name in enum")
     public stringLiteralNameEnum(): void {
         const code = `enum TestEnum {
@@ -99,4 +85,107 @@ export class EnumTests {
         const result = util.transpileAndExecute(code);
         Expect(result).toBe("foo");
     }
+
+    @Test("Enum identifier value internal")
+    public enumIdentifierValueInternal(): void {
+        const result = util.transpileAndExecute(
+            `enum testEnum {
+                abc,
+                def,
+                ghi = def,
+                jkl,
+            }
+            return \`\${testEnum.abc},\${testEnum.def},\${testEnum.ghi},\${testEnum.jkl}\`;`
+        );
+
+        Expect(result).toBe("0,1,1,2");
+    }
+
+    @Test("Enum identifier value internal recursive")
+    public enumIdentifierValueInternalRecursive(): void {
+        const result = util.transpileAndExecute(
+            `enum testEnum {
+                abc,
+                def,
+                ghi = def,
+                jkl = ghi,
+            }
+            return \`\${testEnum.abc},\${testEnum.def},\${testEnum.ghi},\${testEnum.jkl}\`;`
+        );
+
+        Expect(result).toBe("0,1,1,1");
+    }
+
+    @Test("Enum identifier value external")
+    public enumIdentifierValueExternal(): void {
+        const result = util.transpileAndExecute(
+            `const ext = 6;
+            enum testEnum {
+                abc,
+                def,
+                ghi = ext,
+            }
+            return \`\${testEnum.abc},\${testEnum.def},\${testEnum.ghi}\`;`
+        );
+
+        Expect(result).toBe("0,1,6");
+    }
+
+    @Test("Enum reverse mapping")
+    public enumReverseMapping(): void {
+        const result = util.transpileAndExecute(
+            `enum testEnum {
+                abc,
+                def,
+                ghi
+            }
+            return testEnum[testEnum.abc] + testEnum[testEnum.ghi]`
+        );
+
+        Expect(result).toBe("abcghi");
+    }
+
+    @Test("Const enum index")
+    public constEnumIndex(): void {
+        const result = util.transpileAndExecute(
+            `const enum testEnum {
+                abc,
+                def,
+                ghi
+            }
+            return testEnum["def"];`
+        );
+
+        Expect(result).toBe(1);
+    }
+
+    @Test("Const enum index identifier value")
+    public constEnumIndexIdnetifierValue(): void {
+        const result = util.transpileAndExecute(
+            `const enum testEnum {
+                abc,
+                def = 4,
+                ghi,
+                jkl = ghi
+            }
+            return testEnum["jkl"];`
+        );
+
+        Expect(result).toBe(5);
+    }
+
+    @Test("Const enum index identifier chain")
+    public constEnumIndexIdnetifierChain(): void {
+        const result = util.transpileAndExecute(
+            `const enum testEnum {
+                abc = 3,
+                def,
+                ghi = def,
+                jkl = ghi,
+            }
+            return testEnum["ghi"];`
+        );
+
+        Expect(result).toBe(4);
+    }
 }