Fix semantic tokens for enums in if statements (#584)

rokucommunity · May 12, 2022 · b8a0885 · b8a0885
1 parent d598550
commit b8a0885
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 26 deletions.
diff --git a/src/bscPlugin/semanticTokens/BrsFileSemanticTokensProcessor.spec.ts b/src/bscPlugin/semanticTokens/BrsFileSemanticTokensProcessor.spec.ts
@@ -76,4 +76,68 @@ describe('BrsFileSemanticTokensProcessor', () => {
             tokenType: SemanticTokenTypes.enumMember
         }]);
     });
+
+    it('matches each namespace section for enum', () => {
+        const file = program.setFile<BrsFile>('source/main.bs', `
+            sub main()
+                print Sentients.Humanoids.HumanoidType.Cylon
+            end sub
+            namespace Sentients.Humanoids
+                enum HumanoidType
+                    Human
+                    Alien
+                    Cylon
+                end enum
+            end namespace
+        `);
+        program.validate();
+        expectZeroDiagnostics(program);
+        expect(
+            util.sortByRange(program.getSemanticTokens(file.srcPath))
+        ).to.eql([{
+            range: util.createRange(2, 22, 2, 31),
+            tokenType: SemanticTokenTypes.namespace
+        }, {
+            range: util.createRange(2, 32, 2, 41),
+            tokenType: SemanticTokenTypes.namespace
+        }, {
+            range: util.createRange(2, 42, 2, 54),
+            tokenType: SemanticTokenTypes.enum
+        }, {
+            range: util.createRange(2, 55, 2, 60),
+            tokenType: SemanticTokenTypes.enumMember
+        }]);
+    });
+
+    it('matches enums in if statements', () => {
+        const file = program.setFile<BrsFile>('source/main.bs', `
+            sub main()
+                if Humanoids.HumanoidType.Cylon = "Cylon" then
+                    print true
+                end if
+            end sub
+            namespace Humanoids
+                enum HumanoidType
+                    Human = "Human"
+                    Alien = "Alien"
+                    Cylon = "Cylon"
+                end enum
+            end namespace
+        `);
+        program.validate();
+        expectZeroDiagnostics(program);
+        expect(
+            util.sortByRange(program.getSemanticTokens(file.srcPath))
+        ).to.eql([{
+            range: util.createRange(2, 19, 2, 28),
+            tokenType: SemanticTokenTypes.namespace
+        }, {
+            range: util.createRange(2, 29, 2, 41),
+            tokenType: SemanticTokenTypes.enum
+        }, {
+            range: util.createRange(2, 42, 2, 47),
+            tokenType: SemanticTokenTypes.enumMember
+        }]);
+    });
+
 });
diff --git a/src/bscPlugin/semanticTokens/BrsFileSemanticTokensProcessor.ts b/src/bscPlugin/semanticTokens/BrsFileSemanticTokensProcessor.ts
@@ -1,8 +1,9 @@
 import type { Range } from 'vscode-languageserver-protocol';
 import { SemanticTokenTypes } from 'vscode-languageserver-protocol';
-import { isCustomType } from '../../astUtils/reflection';
+import { isBinaryExpression, isCustomType } from '../../astUtils/reflection';
 import type { BrsFile } from '../../files/BrsFile';
 import type { OnGetSemanticTokensEvent } from '../../interfaces';
+import type { Expression } from '../../parser/Expression';
 import { ParseMode } from '../../parser/Parser';
 import util from '../../util';
 
@@ -69,33 +70,43 @@ export class BrsFileSemanticTokensProcessor {
 
     private handleEnums() {
         const enumLookup = this.event.file.program.getFirstScopeForFile(this.event.file)?.getEnumMap();
-        for (const expression of this.event.file.parser.references.expressions) {
-            const parts = util.getAllDottedGetParts(expression)?.map(x => x.toLowerCase());
-            if (parts) {
-                //discard the enum member name
-                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 tokens = util.splitGetRange('.', lowerEnumName + '.' + memberName, expression.range);
-                    //enum member name
-                    this.event.semanticTokens.push({
-                        range: tokens.pop().range,
-                        tokenType: SemanticTokenTypes.enumMember
-                    });
-                    //enum name
-                    this.event.semanticTokens.push({
-                        range: tokens.pop().range,
-                        tokenType: SemanticTokenTypes.enum
-                    });
-                    //namespace parts
-                    for (const token of tokens) {
+        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) {
+                    //discard the enum member name
+                    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 tokens = util.splitGetRange('.', lowerEnumName + '.' + memberName, expression.range);
+                        //enum member name
+                        this.event.semanticTokens.push({
+                            range: tokens.pop().range,
+                            tokenType: SemanticTokenTypes.enumMember
+                        });
+                        //enum name
                         this.event.semanticTokens.push({
-                            range: token.range,
-                            tokenType: SemanticTokenTypes.namespace
+                            range: tokens.pop().range,
+                            tokenType: SemanticTokenTypes.enum
                         });
+                        //namespace parts
+                        for (const token of tokens) {
+                            this.event.semanticTokens.push({
+                                range: token.range,
+                                tokenType: SemanticTokenTypes.namespace
+                            });
+                        }
                     }
                 }
             }