Plugin support for "get all references"

src/LanguageServer.ts

@@ -1237,16 +1237,25 @@
     private async onReferences(params: ReferenceParams) {
         await this.waitAllProjectFirstRuns();
 
-        const position = params.position;
-        const srcPath = util.uriToPath(params.textDocument.uri);
-
-        const results = util.flatMap(
+        const references = util.flatMap(
             await Promise.all(this.getProjects().map(project => {
-                return project.builder.program.getReferences(srcPath, position);
+                return project.builder.program.getReferences({
+                    srcPath: util.uriToPath(params.textDocument.uri),
+                    position: params.position,
+                    includeDeclaration: params?.context?.includeDeclaration ?? true
+                });
             })),
             c => c ?? []
         );
-        return results.filter((r) => r);
+        //TODO filter to distinct locations
+        const result = references
+            //discard null results
+            .filter((r) => r)
+            .map(r => ({
+                uri: URI.file(r.srcPath).toString(),
+                range: r.range
+            }));
+        return result;
     }
 
     private onValidateSettled() {

src/Program.ts

@@ -8,7 +8,7 @@ import { Scope } from './Scope';
 import { DiagnosticMessages } from './DiagnosticMessages';
 import { BrsFile } from './files/BrsFile';
 import { XmlFile } from './files/XmlFile';
-import type { BsDiagnostic, File, FileReference, FileObj, BscFile, SemanticToken, AfterFileTranspileEvent, FileLink, ProvideHoverEvent, ProvideCompletionsEvent, Hover } from './interfaces';
+import type { BsDiagnostic, File, FileReference, FileObj, BscFile, SemanticToken, AfterFileTranspileEvent, FileLink, ProvideHoverEvent, ProvideCompletionsEvent, Hover, Reference, ProvideReferencesEvent } from './interfaces';
 import { standardizePath as s, util } from './util';
 import { XmlScope } from './XmlScope';
 import { DiagnosticFilterer } from './DiagnosticFilterer';
@@ -1016,14 +1016,44 @@ export class Program {
         return signatureHelpUtil.getSignatureHelpItems(callExpressionInfo);
     }
 
-    public getReferences(srcPath: string, position: Position) {
+    /**
+     * @deprecated use the object param type instead
+     */
+    public getReferences(srcPath: string, position: Position);
+    public getReferences(options: GetReferencesOptions);
+    /**
+     * Get all references for the given file and position
+     */
+    public getReferences(...params: any[]) {
+        let options: GetReferencesOptions;
+        //support the old parameter style for backwards compatibility
+        if (typeof params === 'string') {
+            options = {
+                srcPath: params[0],
+                position: params[1]
+            };
+        } else {
+            options = params[0];
+        }
         //find the file
-        let file = this.getFile(srcPath);
-        if (!file) {
-            return null;
+        let file = this.getFile(options.srcPath);
+        let result: Reference[];
+        if (file) {
+            const event = {
+                program: this,
+                file: file,
+                position: options.position,
+                includeDeclaration: options.includeDeclaration ?? true,
+                scopes: this.getScopesForFile(file),
+                references: []
+            } as ProvideReferencesEvent;
+            this.plugins.emit('beforeProvideReferences', event);
+            this.plugins.emit('provideReferences', event);
+            this.plugins.emit('afterProvideReferences', event);
+            result = event.references;
         }
 
-        return file.getReferences(position);
+        return result ?? [];
     }
 
     /**
@@ -1438,3 +1468,9 @@ export interface FileTranspileResult {
     map: SourceMapGenerator;
     typedef: string;
 }
+
+export interface GetReferencesOptions {
+    srcPath: string;
+    position: Position;
+    includeDeclaration?: boolean;
+}

src/bscPlugin/BscPlugin.ts

@@ -1,9 +1,10 @@
 import { isBrsFile, isXmlFile } from '../astUtils/reflection';
-import type { BeforeFileTranspileEvent, CompilerPlugin, OnFileValidateEvent, OnGetCodeActionsEvent, ProvideHoverEvent, OnGetSemanticTokensEvent, OnScopeValidateEvent, ProvideCompletionsEvent } from '../interfaces';
+import type { BeforeFileTranspileEvent, CompilerPlugin, OnFileValidateEvent, OnGetCodeActionsEvent, ProvideHoverEvent, OnGetSemanticTokensEvent, OnScopeValidateEvent, ProvideCompletionsEvent, ProvideReferencesEvent } from '../interfaces';
 import type { Program } from '../Program';
 import { CodeActionsProcessor } from './codeActions/CodeActionsProcessor';
 import { CompletionsProcessor } from './completions/CompletionsProcessor';
 import { HoverProcessor } from './hover/HoverProcessor';
+import { ReferencesProcessor } from './references/ReferencesProcessor';
 import { BrsFileSemanticTokensProcessor } from './semanticTokens/BrsFileSemanticTokensProcessor';
 import { BrsFilePreTranspileProcessor } from './transpile/BrsFilePreTranspileProcessor';
 import { BrsFileValidator } from './validation/BrsFileValidator';
@@ -26,6 +27,13 @@ export class BscPlugin implements CompilerPlugin {
         new CompletionsProcessor(event).process();
     }
 
+    /**
+     * Handle the "find all references" event
+     */
+    public provideReferences(event: ProvideReferencesEvent) {
+        new ReferencesProcessor(event).process();
+    }
+
     public onGetSemanticTokens(event: OnGetSemanticTokensEvent) {
         if (isBrsFile(event.file)) {
             return new BrsFileSemanticTokensProcessor(event as any).process();

src/bscPlugin/references/ReferencesProcessor.spec.ts

@@ -0,0 +1,39 @@
+import { expect } from '../../chai-config.spec';
+import { Program } from '../../Program';
+import { util } from '../../util';
+import { createSandbox } from 'sinon';
+import { rootDir } from '../../testHelpers.spec';
+let sinon = createSandbox();
+
+describe('ReferencesProcessor', () => {
+    let program: Program;
+    beforeEach(() => {
+        program = new Program({ rootDir: rootDir, sourceMap: true });
+    });
+    afterEach(() => {
+        sinon.restore();
+        program.dispose();
+    });
+
+    it('provides references', () => {
+        const file = program.setFile('source/main.bs', `
+            sub main()
+                hello = 1
+                print hello
+            end sub
+        `);
+        const references = program.getReferences({
+            srcPath: file.srcPath,
+            position: util.createPosition(3, 25)
+        });
+        expect(
+            references
+        ).to.eql([{
+            srcPath: file.srcPath,
+            range: util.createRange(2, 16, 2, 21)
+        }, {
+            srcPath: file.srcPath,
+            range: util.createRange(3, 22, 3, 27)
+        }]);
+    });
+});

src/bscPlugin/references/ReferencesProcessor.ts

@@ -0,0 +1,59 @@
+import { isBrsFile, isXmlFile } from '../../astUtils/reflection';
+import { createVisitor, WalkMode } from '../../astUtils/visitors';
+import type { BrsFile } from '../../files/BrsFile';
+import type { ProvideReferencesEvent, Reference } from '../../interfaces';
+
+export class ReferencesProcessor {
+    public constructor(
+        private event: ProvideReferencesEvent
+    ) {
+
+    }
+
+    public process() {
+        if (isBrsFile(this.event.file)) {
+            this.event.references.push(
+                ...this.findVariableReferences(this.event.file)
+            );
+        }
+    }
+
+    private findVariableReferences(file: BrsFile) {
+        const callSiteToken = file.getTokenAt(this.event.position);
+
+        let locations = [] as Reference[];
+
+        const searchFor = callSiteToken.text.toLowerCase();
+
+        for (const scope of this.event.scopes) {
+            const processedFiles = new Set<BrsFile>();
+            for (const file of scope.getAllFiles()) {
+                if (isXmlFile(file) || processedFiles.has(file)) {
+                    continue;
+                }
+                processedFiles.add(file);
+                file.ast.walk(createVisitor({
+                    VariableExpression: (e) => {
+                        if (e.name.text.toLowerCase() === searchFor) {
+                            locations.push({
+                                srcPath: file.srcPath,
+                                range: e.range
+                            });
+                        }
+                    },
+                    AssignmentStatement: (e) => {
+                        if (e.name.text.toLowerCase() === searchFor) {
+                            locations.push({
+                                srcPath: file.srcPath,
+                                range: e.name.range
+                            });
+                        }
+                    }
+                }), {
+                    walkMode: WalkMode.visitAllRecursive
+                });
+            }
+        }
+        return locations;
+    }
+}

src/files/BrsFile.ts

@@ -1637,35 +1637,11 @@ export class BrsFile {
         return statement;
     }
 
+    /**
+     * @deprecated this is handled in the BscPlugin now
+     */
     public getReferences(position: Position) {
-
-        const callSiteToken = this.getTokenAt(position);
-
-        let locations = [] as Location[];
-
-        const searchFor = callSiteToken.text.toLowerCase();
-
-        const scopes = this.program.getScopesForFile(this);
-
-        for (const scope of scopes) {
-            const processedFiles = new Set<BrsFile>();
-            for (const file of scope.getAllFiles()) {
-                if (isXmlFile(file) || processedFiles.has(file)) {
-                    continue;
-                }
-                processedFiles.add(file);
-                file.ast.walk(createVisitor({
-                    VariableExpression: (e) => {
-                        if (e.name.text.toLowerCase() === searchFor) {
-                            locations.push(util.createLocation(util.pathToUri(file.srcPath), e.range));
-                        }
-                    }
-                }), {
-                    walkMode: WalkMode.visitExpressionsRecursive
-                });
-            }
-        }
-        return locations;
+        return [];
     }
 
     /**

src/interfaces.ts

@@ -216,6 +216,21 @@ export interface CompilerPlugin {
      */
     afterProvideCompletions?: PluginHandler<AfterProvideCompletionsEvent>;
 
+
+    /**
+     * Emitted before the program starts collecting references
+     */
+    beforeProvideReferences?: PluginHandler<BeforeProvideReferencesEvent>;
+    /**
+     * Use this event to contribute references
+     */
+    provideReferences?: PluginHandler<ProvideReferencesEvent>;
+    /**
+     * Emitted after the program has finished collecting references, but before they are sent to the client
+     */
+    afterProvideReferences?: PluginHandler<AfterProvideReferencesEvent>;
+
+
     /**
      * Called before the `provideHover` hook. Use this if you need to prepare any of the in-memory objects before the `provideHover` gets called
      */
@@ -278,6 +293,46 @@ export interface ProvideCompletionsEvent<TFile extends BscFile = BscFile> {
 export type BeforeProvideCompletionsEvent<TFile extends BscFile = BscFile> = ProvideCompletionsEvent<TFile>;
 export type AfterProvideCompletionsEvent<TFile extends BscFile = BscFile> = ProvideCompletionsEvent<TFile>;
 
+
+export interface ProvideReferencesEvent<TFile extends BscFile = BscFile> {
+    program: Program;
+    /**
+     * The file where the event was triggered
+     */
+    file: TFile;
+    /**
+     * The scopes that this file is currently loaded in.
+     * Some types of `references` only apply to current scopes, so this might be helpful.
+     * For others, these scopes can be ignored and `program.getScopes()` can be used instead
+     */
+    scopes: Scope[];
+    /**
+     * The position in the file where the cursor was located
+     */
+    position: Position;
+    /**
+     * Should the symbol's declaration be included in the results?
+     */
+    includeDeclaration: boolean;
+    /**
+     * The list of all locations that reference the symbol at the specified position
+     */
+    references: Reference[];
+}
+export type BeforeProvideReferencesEvent<TFile extends BscFile = BscFile> = ProvideReferencesEvent<TFile>;
+export type AfterProvideReferencesEvent<TFile extends BscFile = BscFile> = ProvideReferencesEvent<TFile>;
+export interface Reference {
+    /**
+     * The srcPath for the location being referenced
+     */
+    srcPath: string;
+    /**
+     * The range that the reference points to
+     */
+    range: Range;
+}
+
+
 export interface ProvideHoverEvent {
     program: Program;
     file: BscFile;