Fix enum transpile bug for binary expressions

rokucommunity · Apr 6, 2022 · 617249c · 617249c
1 parent 38811bc
commit 617249c
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 13 deletions.
diff --git a/src/bscPlugin/transpile/BrsFilePreTranspileProcessor.ts b/src/bscPlugin/transpile/BrsFilePreTranspileProcessor.ts
@@ -1,6 +1,8 @@
+import { isBinaryExpression } from '../../astUtils/reflection';
 import { Cache } from '../../Cache';
 import type { BrsFile } from '../../files/BrsFile';
 import type { BeforeFileTranspileEvent } from '../../interfaces';
+import type { Expression } from '../../parser/Expression';
 import util from '../../util';
 
 export class BrsFilePreTranspileProcessor {
@@ -21,19 +23,29 @@ export class BrsFilePreTranspileProcessor {
         if ((enumLookup?.size ?? 0) === 0) {
             return;
         }
-        for (const expression of this.event.file.parser.references.expressions) {
-            const parts = util.getAllDottedGetParts(expression)?.map(x => x.toLowerCase());
-            if (parts) {
-                //get the name of the enum member
-                const memberName = parts.pop();
-                //get the name of the enum (including leading namespace if applicable)
-                const enumName = parts.join('.');
-                const lowerEnumName = enumName.toLowerCase();
-                const theEnum = enumLookup.get(lowerEnumName)?.item;
-                if (theEnum) {
-                    const members = membersByEnum.getOrAdd(lowerEnumName, () => theEnum.getMemberValueMap());
-                    const value = members?.get(memberName);
-                    this.event.editor.overrideTranspileResult(expression, value);
+        for (const referenceExpression of this.event.file.parser.references.expressions) {
+            const actualExpressions: Expression[] = [];
+            //binary expressions actually have two expressions (left and right), so handle them independently
+            if (isBinaryExpression(referenceExpression)) {
+                actualExpressions.push(referenceExpression.left, referenceExpression.right);
+            } else {
+                //assume all other expressions are a single chain
+                actualExpressions.push(referenceExpression);
+            }
+            for (const expression of actualExpressions) {
+                const parts = util.getAllDottedGetParts(expression)?.map(x => x.toLowerCase());
+                if (parts) {
+                    //get the name of the enum member
+                    const memberName = parts.pop();
+                    //get the name of the enum (including leading namespace if applicable)
+                    const enumName = parts.join('.');
+                    const lowerEnumName = enumName.toLowerCase();
+                    const theEnum = enumLookup.get(lowerEnumName)?.item;
+                    if (theEnum) {
+                        const members = membersByEnum.getOrAdd(lowerEnumName, () => theEnum.getMemberValueMap());
+                        const value = members?.get(memberName);
+                        this.event.editor.overrideTranspileResult(expression, value);
+                    }
                 }
             }
         }

diff --git a/src/parser/tests/statement/Enum.spec.ts b/src/parser/tests/statement/Enum.spec.ts
@@ -868,6 +868,57 @@ describe('EnumStatement', () => {
                 kind: CompletionItemKind.Enum
             }]);
         });
+
+        it('handles both sides of a logical expression', () => {
+            testTranspile(`
+                sub main()
+                    dir = m.direction = Direction.up
+                    dir = Direction.up = m.direction
+                end sub
+                enum Direction
+                    up = "up"
+                    down = "down"
+                end enum
+            `, `
+                sub main()
+                    dir = m.direction = "up"
+                    dir = "up" = m.direction
+                end sub
+            `);
+        });
+
+        it('replaces enum values in if statements', () => {
+            testTranspile(`
+                sub main()
+                    if m.direction = Direction.up
+                        print Direction.up
+                    end if
+                end sub
+                enum Direction
+                    up = "up"
+                    down = "down"
+                end enum
+            `, `
+                sub main()
+                    if m.direction = "up"
+                        print "up"
+                    end if
+                end sub
+            `);
+        });
+
+        it('replaces enum values in function default parameter value expressions', () => {
+            testTranspile(`
+                sub speak(dir = Direction.up)
+                end sub
+                enum Direction
+                    up = "up"
+                end enum
+            `, `
+                sub speak(dir = "up")
+                end sub
+            `);
+        });
     });
 
 });