Improve workspace/document symbol handling

src/LanguageServer.spec.ts

@@ -987,7 +987,6 @@ describe('LanguageServer', () => {
             // We run the check twice as the first time is with it not cached and second time is with it cached
             for (let i = 0; i < 2; i++) {
                 const symbols = await server['onWorkspaceSymbol']({} as any);
-                expect(symbols.length).to.equal(4);
                 expect(
                     symbols.map(x => ({
                         name: x.name,

src/bscPlugin/symbols/DocumentSymbolProcessor.spec.ts

@@ -130,7 +130,9 @@ describe('DocumentSymbolProcessor', () => {
             enum alpha
                 name = 1
             end enum
-        `, ['tokens', 'name']);
+        `, ['body', '0', 'tokens', 'name'], {
+            alpha: SymbolKind.Enum
+        });
     });
 
     it('finds functions', () => {

src/bscPlugin/symbols/DocumentSymbolProcessor.ts

@@ -1,7 +1,7 @@
 import { isBrsFile } from '../../astUtils/reflection';
 import type { BrsFile } from '../../files/BrsFile';
 import type { ProvideDocumentSymbolsEvent } from '../../interfaces';
-import { getDocumentSymbolsFromStatement } from './symbolUtils';
+import { getDocumentSymbolsFromBrsFile } from './symbolUtils';
 
 export class DocumentSymbolProcessor {
     public constructor(
@@ -17,12 +17,8 @@ export class DocumentSymbolProcessor {
     }
 
     private getBrsFileDocumentSymbols(file: BrsFile) {
-        for (const statement of file.ast.statements) {
-            const symbol = getDocumentSymbolsFromStatement(statement);
-            if (symbol) {
-                this.event.documentSymbols.push(...symbol);
-            }
-        }
+        const symbols = getDocumentSymbolsFromBrsFile(file);
+        this.event.documentSymbols.push(...symbols);
         return this.event.documentSymbols;
     }
 }

src/bscPlugin/symbols/WorkspaceSymbolProcessor.spec.ts

@@ -57,7 +57,7 @@ describe('WorkspaceSymbolProcessor', () => {
         let result = `${symbol.name}|${SymbolKindMap.get(symbol.kind)}|${symbol.location.uri}`;
         const range = (symbol as any).location.range;
         if (range) {
-            result += util.rangeToString(range);
+            result += '|' + util.rangeToString(range);
         }
         return result;
     }
@@ -163,7 +163,9 @@ describe('WorkspaceSymbolProcessor', () => {
             enum alpha
                 name = 1
             end enum
-        `, ['tokens', 'name']);
+        `, ['body', '0', 'tokens', 'name'], [
+            ['alpha', SymbolKind.Enum]
+        ]);
     });
 
     it('finds functions', () => {
@@ -173,8 +175,8 @@ describe('WorkspaceSymbolProcessor', () => {
             function beta()
             end function
         `], [
-            ['alpha', SymbolKind.Function, 'source/lib0.brs', 1, 12, 2, 24],
-            ['beta', SymbolKind.Function, 'source/lib0.brs', 3, 12, 4, 24]
+            ['alpha', SymbolKind.Function, 'source/lib0.brs', 1, 21, 1, 26],
+            ['beta', SymbolKind.Function, 'source/lib0.brs', 3, 21, 3, 25]
         ]);
     });
 

src/bscPlugin/symbols/WorkspaceSymbolProcessor.ts

@@ -1,8 +1,7 @@
 import { isBrsFile } from '../../astUtils/reflection';
 import type { BrsFile } from '../../files/BrsFile';
 import type { ProvideWorkspaceSymbolsEvent } from '../../interfaces';
-import { getWorkspaceSymbolsFromStatement } from './symbolUtils';
-import util from '../../util';
+import { getWorkspaceSymbolsFromBrsFile } from './symbolUtils';
 
 export class WorkspaceSymbolProcessor {
     public constructor(
@@ -22,12 +21,8 @@ export class WorkspaceSymbolProcessor {
     }
 
     private getBrsFileWorkspaceSymbols(file: BrsFile) {
-        for (const statement of file.ast.statements) {
-            const symbol = getWorkspaceSymbolsFromStatement(statement, util.pathToUri(file.srcPath));
-            if (symbol) {
-                this.event.workspaceSymbols.push(...symbol);
-            }
-        }
+        const symbols = getWorkspaceSymbolsFromBrsFile(file);
+        this.event.workspaceSymbols.push(...symbols);
         return this.event.workspaceSymbols;
     }
 }

src/bscPlugin/symbols/symbolUtils.ts

@@ -1,116 +1,158 @@
 import type { Range } from 'vscode-languageserver-protocol';
 import { WorkspaceSymbol } from 'vscode-languageserver-protocol';
 import { DocumentSymbol, SymbolKind } from 'vscode-languageserver-protocol';
-import type { Statement } from '../../parser/AstNode';
-import { ClassStatement, ConstStatement, EnumMemberStatement, EnumStatement, FieldStatement, FunctionStatement, InterfaceFieldStatement, InterfaceMethodStatement, InterfaceStatement, MethodStatement, NamespaceStatement } from '../../parser/Statement';
+import type { AstNode } from '../../parser/AstNode';
+import util from '../../util';
+import type { BrsFile } from '../../files/BrsFile';
+import { WalkMode, createVisitor } from '../../astUtils/visitors';
 
-export function getDocumentSymbolsFromStatement(statement: Statement) {
-    return getSymbolsFromStatement(statement, (name: string, documenation: string, kind: SymbolKind, range: Range, selectionRange: Range, children: DocumentSymbol[], containerName: string) => {
-        return [DocumentSymbol.create(name, documenation, kind, range, selectionRange, children)];
-    });
-}
+export function getDocumentSymbolsFromBrsFile(file: BrsFile) {
+    let result: DocumentSymbol[] = [];
+    const symbols = getSymbolsFromAstNode(file.ast);
+    for (let symbol of symbols) {
+        result.push(
+            createSymbol(symbol)
+        );
+    }
+    return result;
 
-export function getWorkspaceSymbolsFromStatement(statement: Statement, uri: string) {
-    return getSymbolsFromStatement(statement, (name: string, documenation: string, kind: SymbolKind, range: Range, selectionRange: Range, children: WorkspaceSymbol[], containerName: string) => {
-        const symbol = WorkspaceSymbol.create(name, kind, uri, range);
-        symbol.containerName = containerName;
-        return [symbol, ...(children ?? [])];
-    });
+    function createSymbol(symbol: SymbolInfo): DocumentSymbol {
+        return DocumentSymbol.create(
+            symbol.name,
+            symbol.documentation,
+            symbol.kind,
+            symbol.range,
+            symbol.selectionRange,
+            symbol.children.map(x => createSymbol(x))
+        );
+    }
 }
 
-type SymbolFactory<T> = (name: string, documenation: string, kind: SymbolKind, range: Range, selectionRange: Range, children: T[], containerName: string) => T[];
+export function getWorkspaceSymbolsFromBrsFile(file: BrsFile) {
+    const result: WorkspaceSymbol[] = [];
+    const uri = util.pathToUri(file.srcPath);
+    let symbolsToProcess = getSymbolsFromAstNode(file.ast);
+    while (symbolsToProcess.length > 0) {
+        //get the symbol
+        const symbolInfo = symbolsToProcess.shift();
+        //push any children to be processed later
+        symbolsToProcess.push(...symbolInfo.children);
+        const workspaceSymbol = WorkspaceSymbol.create(
+            symbolInfo.name,
+            symbolInfo.kind,
+            uri,
+            symbolInfo.selectionRange
+        );
+        workspaceSymbol.containerName = symbolInfo.containerName;
+        result.push(workspaceSymbol);
+    }
+    return result;
+}
 
-/**
- * TypeScript won't type narrow within a switch statement, so we use this function to do the type narrowing for us.
- * Hopefully v8 will just inline the function and we won't pay a perf penalty for this. This does not actually do any runtime checking, it just narrows the type for TypeScript's benefit.
- */
-function coerce<T>(value: any): value is T {
-    return true;
+interface SymbolInfo {
+    name: string;
+    documentation: string;
+    kind: SymbolKind;
+    range: Range;
+    selectionRange: Range;
+    containerName: string;
+    children: SymbolInfo[];
 }
 
-// eslint-disable-next-line @typescript-eslint/no-redundant-type-constituents
-function getSymbolsFromStatement<T extends WorkspaceSymbol | DocumentSymbol>(statement: Statement, factory: SymbolFactory<T>, containerName?: string): T[] {
-    switch (statement?.constructor?.name) {
-        case FunctionStatement.name:
-            if (coerce<FunctionStatement>(statement) && statement.name?.text) {
-                return factory(statement.name.text, '', SymbolKind.Function, statement.range, statement.name.range, undefined, containerName);
-            }
-            break;
+function getSymbolsFromAstNode(node: AstNode): SymbolInfo[] {
+    //collection of every symbol, indexed by the node it was based on (this is useful to help attach children to their parents)
+    const result: SymbolInfo[] = [];
+    const lookup = new Map<AstNode, SymbolInfo>();
 
-        case ClassStatement.name:
-            if (coerce<ClassStatement>(statement) && statement.name?.text) {
-                const children = statement.body
-                    .map((x) => getSymbolsFromStatement(x, factory, statement.name.text))
-                    .flat()
-                    .filter(x => !!x);
-                return factory(statement.name.text, '', SymbolKind.Class, statement.range, statement.name.range, children, containerName);
-            }
-            break;
+    function addSymbol(node: AstNode, name: string, kind: SymbolKind, range: Range, selectionRange: Range, documenation?: string) {
+        const symbol = {
+            name: name,
+            documentation: documenation,
+            kind: kind,
+            range: range,
+            selectionRange: selectionRange,
+            containerName: undefined,
+            children: []
+        };
+        lookup.set(node, symbol);
 
-        case FieldStatement.name:
-            if (coerce<FieldStatement>(statement) && statement.name?.text) {
-                return factory(statement.name.text, '', SymbolKind.Field, statement.range, statement.name.range, undefined, containerName);
+        let parent = node.parent;
+        while (parent) {
+            if (lookup.has(parent)) {
+                break;
             }
-            break;
+            parent = parent.parent;
+        }
+        //if we found a parent, add this symbol as a child of the parent
+        if (parent) {
+            const parentSymbol = lookup.get(parent);
+            symbol.containerName = parentSymbol.name;
 
-        case MethodStatement.name:
-            if (coerce<MethodStatement>(statement) && statement.name?.text) {
-                return factory(statement.name.text, '', SymbolKind.Method, statement.range, statement.name.range, undefined, containerName);
-            }
-            break;
+            parentSymbol.children.push(symbol);
+        } else {
+            //there's no parent. add the symbol as a top level result
+            result.push(symbol);
+        }
+    }
 
-        case InterfaceStatement.name:
-            if (coerce<InterfaceStatement>(statement) && statement.tokens.name?.text) {
-                const children = statement.body
-                    .map((x) => getSymbolsFromStatement(x, factory, statement.tokens.name.text))
-                    .flat()
-                    .filter(x => !!x);
-                return factory(statement.tokens.name.text, '', SymbolKind.Interface, statement.range, statement.tokens.name.range, children, containerName);
+    node.walk(createVisitor({
+        FunctionStatement: (statement) => {
+            if (statement.name?.text) {
+                addSymbol(statement, statement.name.text, SymbolKind.Function, statement.range, statement.name.range);
             }
-            break;
-
-        case InterfaceFieldStatement.name:
-            if (coerce<InterfaceFieldStatement>(statement) && statement.tokens.name?.text) {
-                return factory(statement.tokens.name.text, '', SymbolKind.Field, statement.range, statement.tokens.name.range, undefined, containerName);
+        },
+        ClassStatement: (statement, parent) => {
+            if (statement.name?.text) {
+                addSymbol(statement, statement.name.text, SymbolKind.Class, statement.range, statement.name.range);
             }
-            break;
-
-        case InterfaceMethodStatement.name:
-            if (coerce<InterfaceMethodStatement>(statement) && statement.tokens.name?.text) {
-                return factory(statement.tokens.name.text, '', SymbolKind.Method, statement.range, statement.tokens.name.range, undefined, containerName);
+        },
+        FieldStatement: (statement, parent) => {
+            if (statement.name?.text) {
+                addSymbol(statement, statement.name.text, SymbolKind.Field, statement.range, statement.name.range);
             }
-            break;
-
-        case ConstStatement.name:
-            if (coerce<ConstStatement>(statement) && statement.tokens.name?.text) {
-                return factory(statement.tokens.name.text, '', SymbolKind.Constant, statement.range, statement.tokens.name.range, undefined, containerName);
+        },
+        MethodStatement: (statement, parent) => {
+            if (statement.name?.text) {
+                addSymbol(statement, statement.name.text, SymbolKind.Method, statement.range, statement.name.range);
             }
-            break;
-
-        case NamespaceStatement.name:
-            if (coerce<NamespaceStatement>(statement) && statement.nameExpression) {
-                const children = statement.body.statements
-                    .map((x) => getSymbolsFromStatement(x, factory, statement.name))
-                    .flat()
-                    .filter(x => !!x);
-                return factory(statement.nameExpression.getNameParts().pop(), '', SymbolKind.Namespace, statement.range, statement.nameExpression.range, children, containerName);
+        },
+        InterfaceStatement: (statement, parent) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.Interface, statement.range, statement.tokens.name.range);
             }
-            break;
-
-        case EnumStatement.name:
-            if (coerce<EnumStatement>(statement) && statement.tokens.name?.text) {
-                const children = statement.body
-                    .map((x) => getSymbolsFromStatement(x, factory, statement.name))
-                    .flat()
-                    .filter(x => !!x);
-                return factory(statement.tokens.name.text, '', SymbolKind.Enum, statement.range, statement.tokens.name.range, children, containerName);
+        },
+        InterfaceFieldStatement: (statement, parent) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.Field, statement.range, statement.tokens.name.range);
             }
-            break;
-
-        case EnumMemberStatement.name:
-            if (coerce<EnumMemberStatement>(statement)) {
-                return factory(statement.tokens.name.text, '', SymbolKind.EnumMember, statement.range, statement.tokens.name.range, undefined, containerName);
+        },
+        InterfaceMethodStatement: (statement, parent) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.Method, statement.range, statement.tokens.name.range);
             }
-            break;
-    }
+        },
+        ConstStatement: (statement) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.Constant, statement.range, statement.tokens.name.range);
+            }
+        },
+        NamespaceStatement: (statement) => {
+            if (statement.nameExpression) {
+                addSymbol(statement, statement.nameExpression.getNameParts().pop(), SymbolKind.Namespace, statement.range, statement.nameExpression.range);
+            }
+        },
+        EnumStatement: (statement) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.Enum, statement.range, statement.tokens.name.range);
+            }
+        },
+        EnumMemberStatement: (statement) => {
+            if (statement.tokens.name?.text) {
+                addSymbol(statement, statement.tokens.name.text, SymbolKind.EnumMember, statement.range, statement.tokens.name.range);
+            }
+        }
+    }), {
+        walkMode: WalkMode.visitAllRecursive
+    });
+    return result;
 }