Simplify const enum transform (#623)

ark120202 · Perryvw · commit 4ff68194b3f7 · 2019-06-14T18:51:50.000+02:00
* Simplify const enum transform

* Rename getConstEnumValue to tryGetConstEnumValue
diff --git a/src/LuaTransformer.ts b/src/LuaTransformer.ts
@@ -3889,6 +3889,11 @@ export class LuaTransformer {
     public transformPropertyAccessExpression(expression: ts.PropertyAccessExpression): ExpressionVisitResult {
         const property = this.getIdentifierText(expression.name);
 
+        const constEnumValue = this.tryGetConstEnumValue(expression);
+        if (constEnumValue) {
+            return constEnumValue;
+        }
+
         // Check for primitive types to override
         const type = this.checker.getTypeAtLocation(expression.expression);
         if (tsHelper.isStringType(type)) {
@@ -3898,8 +3903,6 @@ export class LuaTransformer {
             if (arrayPropertyAccess) {
                 return arrayPropertyAccess;
             }
-        } else if (type.symbol && type.symbol.flags & ts.SymbolFlags.ConstEnum) {
-            return this.transformConstEnumValue(type, property, expression);
         }
 
         this.checkForLuaLibType(type);
@@ -4062,16 +4065,12 @@ export class LuaTransformer {
         const table = this.transformExpression(expression.expression);
         const index = this.transformExpression(expression.argumentExpression);
 
-        const type = this.checker.getTypeAtLocation(expression.expression);
-
-        if (
-            type.symbol &&
-            type.symbol.flags & ts.SymbolFlags.ConstEnum &&
-            ts.isStringLiteral(expression.argumentExpression)
-        ) {
-            return this.transformConstEnumValue(type, expression.argumentExpression.text, expression);
+        const constEnumValue = this.tryGetConstEnumValue(expression);
+        if (constEnumValue) {
+            return constEnumValue;
         }
 
+        const type = this.checker.getTypeAtLocation(expression.expression);
         if (tsHelper.isArrayType(type, this.checker, this.program)) {
             return tstl.createTableIndexExpression(table, this.expressionPlusOne(index), expression);
         } else if (tsHelper.isStringType(type)) {
@@ -4085,46 +4084,15 @@ export class LuaTransformer {
         }
     }
 
-    protected transformConstEnumValue(
-        enumType: ts.EnumType,
-        memberName: string,
-        tsOriginal: ts.Node
-    ): ExpressionVisitResult {
-        // 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 && valueName) {
-                        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);
-            }
+    private tryGetConstEnumValue(
+        node: ts.PropertyAccessExpression | ts.ElementAccessExpression
+    ): tstl.Expression | undefined {
+        const value = this.checker.getConstantValue(node);
+        if (typeof value === "string") {
+            return tstl.createStringLiteral(value, node);
+        } else if (typeof value === "number") {
+            return tstl.createNumericLiteral(value, node);
         }
-        throw TSTLErrors.CouldNotFindEnumMember(enumDeclaration, memberName, tsOriginal);
     }
 
     protected transformStringCallExpression(node: ts.CallExpression): tstl.Expression {
diff --git a/src/TSTLErrors.ts b/src/TSTLErrors.ts
@@ -8,9 +8,6 @@ export class TSTLErrors {
     public static CouldNotCast = (castName: string) =>
         new Error(`Failed to cast all elements to expected type using ${castName}.`);
 
-    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);
 
