diff --git a/src/extension/linkify/test/vscode-node/findWord.test.ts b/src/extension/linkify/test/vscode-node/findWord.test.ts new file mode 100644 index 0000000000..de22235ec5 --- /dev/null +++ b/src/extension/linkify/test/vscode-node/findWord.test.ts @@ -0,0 +1,254 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import assert from 'assert'; +import * as vscode from 'vscode'; +import { afterEach, suite, test } from 'vitest'; +import { TreeSitterExpressionInfo } from '../../../../platform/parser/node/nodes'; +import { IParserService, TreeSitterAST } from '../../../../platform/parser/node/parserService'; +import { WASMLanguage } from '../../../../platform/parser/node/treeSitterLanguages'; +import { CancellationToken } from '../../../../util/vs/base/common/cancellation'; +import { URI } from '../../../../util/vs/base/common/uri'; +import { findSymbolLocationInFile, SymbolFileCache } from '../../vscode-node/findWord'; + +class TestParserService implements Partial { + public parseCount = 0; + public genericSymbolQueryCount = 0; + + constructor( + private readonly symbols: readonly TreeSitterExpressionInfo[] = [], + private readonly classDeclarations: readonly TreeSitterExpressionInfo[] = [], + private readonly functionDefinitions: readonly TreeSitterExpressionInfo[] = [], + private readonly typeDeclarations: readonly TreeSitterExpressionInfo[] = [], + ) { } + + getTreeSitterASTForWASMLanguage(_language: WASMLanguage, _source: string): TreeSitterAST { + this.parseCount++; + const symbols = this.symbols; + const classDeclarations = this.classDeclarations; + const functionDefinitions = this.functionDefinitions; + const typeDeclarations = this.typeDeclarations; + return { + getClassDeclarations: async () => classDeclarations, + getFunctionDefinitions: async () => functionDefinitions, + getTypeDeclarations: async () => typeDeclarations, + getSymbols: async () => { + this.genericSymbolQueryCount++; + return symbols; + }, + } as unknown as TreeSitterAST; + } +} + +interface MutableTestWorkspace { + textDocuments: typeof vscode.workspace.textDocuments; + fs: { + readFile: typeof vscode.workspace.fs.readFile; + }; +} + +interface PartialMutableTestWorkspace { + textDocuments?: MutableTestWorkspace['textDocuments']; + fs?: Partial; +} + +function ensureTestWorkspace(): MutableTestWorkspace { + const testVscode = vscode as unknown as { workspace?: PartialMutableTestWorkspace }; + testVscode.workspace ??= {}; + testVscode.workspace.textDocuments ??= []; + testVscode.workspace.fs ??= {}; + testVscode.workspace.fs.readFile ??= (async () => { throw new Error('workspace.fs.readFile not mocked in test'); }) as typeof vscode.workspace.fs.readFile; + return testVscode.workspace as MutableTestWorkspace; +} + +const testWorkspace = ensureTestWorkspace(); +const originalWorkspaceReadFile = vscode.workspace.fs.readFile; +const originalWorkspaceTextDocuments = vscode.workspace.textDocuments; + +afterEach(() => { + testWorkspace.textDocuments = originalWorkspaceTextDocuments; + testWorkspace.fs.readFile = originalWorkspaceReadFile; +}); + +function setWorkspaceFileContents(contentsByUri: ReadonlyMap): void { + testWorkspace.textDocuments = []; + testWorkspace.fs.readFile = async (uri: vscode.Uri) => { + const contents = contentsByUri.get(uri.toString()); + if (contents === undefined) { + throw new Error(`File not found: ${uri.toString()}`); + } + return new TextEncoder().encode(contents); + }; +} + +function symbol(contents: string, identifier: string): TreeSitterExpressionInfo { + const startIndex = contents.indexOf(identifier); + return { + identifier, + text: identifier, + startIndex, + endIndex: startIndex + identifier.length, + }; +} + +function declaration(contents: string, identifier: string, text: string): TreeSitterExpressionInfo { + const startIndex = contents.indexOf(text); + return { + identifier, + text, + startIndex, + endIndex: startIndex + text.length, + }; +} + +function asParserService(parserService: TestParserService): IParserService { + return parserService as unknown as IParserService; +} + +suite('Find symbol location in file', () => { + + test('Should return the exact symbol location', async () => { + const contents = [ + 'const value = 1;', + '', + 'class Foo {', + '}', + ].join('\n'); + const uri = URI.file('/workspace/src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const location = await findSymbolLocationInFile( + asParserService(new TestParserService([symbol(contents, 'Foo')])), + uri, + 'Foo', + CancellationToken.None, + ); + + assert(location); + assert.strictEqual(location.uri.toString(), uri.toString()); + assert.strictEqual(location.range.start.line, 2); + assert.strictEqual(location.range.start.character, 6); + }); + + test('Should prefer declaration matches over earlier generic symbol references', async () => { + const declarationText = 'class Foo(Base):'; + const contents = [ + 'if isinstance(module, Foo):', + '', + declarationText, + '\tpass', + ].join('\n'); + const uri = URI.file('/workspace/src/file.py'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const parserService = new TestParserService( + [symbol(contents, 'Foo')], + [declaration(contents, 'Foo', declarationText)], + ); + const location = await findSymbolLocationInFile( + asParserService(parserService), + uri, + 'Foo', + CancellationToken.None, + ); + + assert(location); + assert.strictEqual(location.range.start.line, 2); + assert.strictEqual(location.range.start.character, 0); + assert.strictEqual(parserService.genericSymbolQueryCount, 0); + }); + + test('Should prefer declaration fallback over generic symbol references for qualified names', async () => { + const declarationText = 'class Foo:'; + const contents = [ + 'if value.bar:', + '\tpass', + '', + declarationText, + '\tpass', + ].join('\n'); + const uri = URI.file('/workspace/src/file.py'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const parserService = new TestParserService( + [symbol(contents, 'bar')], + [declaration(contents, 'Foo', declarationText)], + ); + const location = await findSymbolLocationInFile( + asParserService(parserService), + uri, + 'Foo.bar', + CancellationToken.None, + ); + + assert(location); + assert.strictEqual(location.range.start.line, 3); + assert.strictEqual(location.range.start.character, 0); + assert.strictEqual(parserService.genericSymbolQueryCount, 0); + }); + + test('Should use the highest-index qualified name part when there is no exact match', async () => { + const contents = [ + 'class Foo {', + '\tmethod() {', + '\t}', + '}', + ].join('\n'); + const uri = URI.file('/workspace/src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const location = await findSymbolLocationInFile( + asParserService(new TestParserService([ + symbol(contents, 'Foo'), + symbol(contents, 'method'), + ])), + uri, + 'Foo.method', + CancellationToken.None, + ); + + assert(location); + assert.strictEqual(location.range.start.line, 1); + assert.strictEqual(location.range.start.character, 1); + }); + + test('Should return undefined for unsupported, missing, or unmatched files', async () => { + const contents = 'class Foo {}'; + const tsUri = URI.file('/workspace/src/file.ts'); + const txtUri = URI.file('/workspace/src/file.txt'); + setWorkspaceFileContents(new Map([[tsUri.toString(), contents]])); + + const parserService = asParserService(new TestParserService([symbol(contents, 'Foo')])); + + assert.strictEqual(await findSymbolLocationInFile(parserService, txtUri, 'Foo', CancellationToken.None), undefined); + assert.strictEqual(await findSymbolLocationInFile(parserService, URI.file('/workspace/src/missing.ts'), 'Foo', CancellationToken.None), undefined); + assert.strictEqual(await findSymbolLocationInFile(parserService, tsUri, 'Missing', CancellationToken.None), undefined); + }); + + test('Should reuse cached file symbols for repeated URI lookups', async () => { + const contents = [ + 'class Foo {', + '\tmethod() {', + '\t}', + '}', + ].join('\n'); + const uri = URI.file('/workspace/src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const parserService = new TestParserService([ + symbol(contents, 'Foo'), + symbol(contents, 'method'), + ]); + const cache: SymbolFileCache = new Map(); + + const classLocation = await findSymbolLocationInFile(asParserService(parserService), uri, 'Foo', CancellationToken.None, cache); + const methodLocation = await findSymbolLocationInFile(asParserService(parserService), uri, 'Foo.method', CancellationToken.None, cache); + + assert(classLocation); + assert(methodLocation); + assert.strictEqual(parserService.parseCount, 1); + assert.strictEqual(parserService.genericSymbolQueryCount, 1); + }); +}); diff --git a/src/extension/linkify/test/vscode-node/symbolLinkifier.test.ts b/src/extension/linkify/test/vscode-node/symbolLinkifier.test.ts index c49bb5293e..b93527b6d7 100644 --- a/src/extension/linkify/test/vscode-node/symbolLinkifier.test.ts +++ b/src/extension/linkify/test/vscode-node/symbolLinkifier.test.ts @@ -3,28 +3,128 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import assert from 'assert'; import * as vscode from 'vscode'; +import { afterEach, suite, test } from 'vitest'; import { NullEnvService } from '../../../../platform/env/common/nullEnvService'; +import { TreeSitterExpressionInfo } from '../../../../platform/parser/node/nodes'; +import { IParserService, TreeSitterAST } from '../../../../platform/parser/node/parserService'; +import { WASMLanguage } from '../../../../platform/parser/node/treeSitterLanguages'; +import { CancellationToken } from '../../../../util/vs/base/common/cancellation'; import { LinkifySymbolAnchor } from '../../common/linkifiedText'; import { ILinkifyService, LinkifyService } from '../../common/linkifyService'; import { SymbolLinkifier } from '../../vscode-node/symbolLinkifier'; import { assertPartsEqual, createMockFsService, createMockWorkspaceService, linkify, workspaceFile } from '../node/util'; -function createTestLinkifierService(...listOfFiles: readonly string[]): ILinkifyService { +class TestParserService implements Partial { + public parseCount = 0; + + constructor( + private readonly symbols: readonly TreeSitterExpressionInfo[] = [], + private readonly classDeclarations: readonly TreeSitterExpressionInfo[] = [], + private readonly functionDefinitions: readonly TreeSitterExpressionInfo[] = [], + private readonly typeDeclarations: readonly TreeSitterExpressionInfo[] = [], + ) { } + + getTreeSitterASTForWASMLanguage(_language: WASMLanguage, _source: string): TreeSitterAST { + this.parseCount++; + const symbols = this.symbols; + const classDeclarations = this.classDeclarations; + const functionDefinitions = this.functionDefinitions; + const typeDeclarations = this.typeDeclarations; + return { + getClassDeclarations: async () => classDeclarations, + getFunctionDefinitions: async () => functionDefinitions, + getTypeDeclarations: async () => typeDeclarations, + getSymbols: async () => symbols, + } as unknown as TreeSitterAST; + } +} + +interface MutableTestWorkspace { + textDocuments: typeof vscode.workspace.textDocuments; + fs: { + readFile: typeof vscode.workspace.fs.readFile; + }; +} + +interface MutableTestCommands { + executeCommand: typeof vscode.commands.executeCommand; +} + +interface PartialMutableTestWorkspace { + textDocuments?: MutableTestWorkspace['textDocuments']; + fs?: Partial; +} + +function ensureTestWorkspace(): MutableTestWorkspace { + const testVscode = vscode as unknown as { workspace?: PartialMutableTestWorkspace }; + testVscode.workspace ??= {}; + testVscode.workspace.textDocuments ??= []; + testVscode.workspace.fs ??= {}; + testVscode.workspace.fs.readFile ??= (async () => { throw new Error('workspace.fs.readFile not mocked in test'); }) as typeof vscode.workspace.fs.readFile; + return testVscode.workspace as MutableTestWorkspace; +} + +function ensureTestCommands(): MutableTestCommands { + const testVscode = vscode as unknown as { commands?: Partial }; + testVscode.commands ??= {}; + testVscode.commands.executeCommand ??= (async () => undefined) as typeof vscode.commands.executeCommand; + return testVscode.commands as MutableTestCommands; +} + +const testWorkspace = ensureTestWorkspace(); +const testCommands = ensureTestCommands(); +const originalWorkspaceReadFile = vscode.workspace.fs.readFile; +const originalWorkspaceTextDocuments = vscode.workspace.textDocuments; +const originalExecuteCommand = vscode.commands.executeCommand; + +afterEach(() => { + testWorkspace.textDocuments = originalWorkspaceTextDocuments; + testWorkspace.fs.readFile = originalWorkspaceReadFile; + testCommands.executeCommand = originalExecuteCommand; +}); + +function setWorkspaceFileContents(contentsByUri: ReadonlyMap): void { + testWorkspace.textDocuments = []; + testWorkspace.fs.readFile = async (uri: vscode.Uri) => { + const contents = contentsByUri.get(uri.toString()); + if (contents === undefined) { + throw new Error(`File not found: ${uri.toString()}`); + } + return new TextEncoder().encode(contents); + }; +} + +function asParserService(parserService: TestParserService): IParserService { + return parserService as unknown as IParserService; +} + +function symbol(contents: string, identifier: string): TreeSitterExpressionInfo { + const startIndex = contents.indexOf(identifier); + return { + identifier, + text: identifier, + startIndex, + endIndex: startIndex + identifier.length, + }; +} + +function createTestLinkifierService(listOfFiles: readonly string[], parserService: IParserService = asParserService(new TestParserService())): ILinkifyService { const fs = createMockFsService(listOfFiles); const workspaceService = createMockWorkspaceService(); const linkifier = new LinkifyService(fs, workspaceService, NullEnvService.Instance); - linkifier.registerGlobalLinkifier({ create: () => new SymbolLinkifier(fs, workspaceService) }); + linkifier.registerGlobalLinkifier({ create: () => new SymbolLinkifier(fs, parserService, workspaceService) }); return linkifier; } suite('Symbol Linkify', () => { test(`Should create symbol links from Markdown links`, async () => { - const linkifier = createTestLinkifierService( + const linkifier = createTestLinkifierService([ 'file.ts', 'src/file.ts', - ); + ]); assertPartsEqual( (await linkify(linkifier, '[`symbol`](file.ts) [`symbol`](src/file.ts)') @@ -48,7 +148,7 @@ suite('Symbol Linkify', () => { }); test(`Should de-linkify symbol links to files that don't exist`, async () => { - const linkifier = createTestLinkifierService(); + const linkifier = createTestLinkifierService([]); assertPartsEqual( (await linkify(linkifier, '[`symbol`](file.ts) [`symbol`](src/file.ts)' @@ -60,10 +160,10 @@ suite('Symbol Linkify', () => { }); test(`Should create symbol links for symbols containing $ or _`, async () => { - const linkifier = createTestLinkifierService( + const linkifier = createTestLinkifierService([ 'file.ts', 'src/file.ts', - ); + ]); assertPartsEqual( (await linkify(linkifier, '[`_symbol`](file.ts) [`$symbol`](src/file.ts)', @@ -87,10 +187,10 @@ suite('Symbol Linkify', () => { }); test(`Should create symbol links for symbols with function call or generic syntax`, async () => { - const linkifier = createTestLinkifierService( + const linkifier = createTestLinkifierService([ 'file.ts', 'src/file.ts', - ); + ]); assertPartsEqual( (await linkify(linkifier, @@ -115,9 +215,9 @@ suite('Symbol Linkify', () => { }); test(`Should support files with spaces`, async () => { - const linkifier = createTestLinkifierService( + const linkifier = createTestLinkifierService([ 'space file.ts', - ); + ]); assertPartsEqual( (await linkify(linkifier, '[`symbol`](space%20file.ts) [`symbol`](space%20file.ts)') @@ -139,4 +239,132 @@ suite('Symbol Linkify', () => { ], ); }); + + test(`Should use tree-sitter for linked-backtick initial symbol locations`, async () => { + const contents = [ + 'const value = 1;', + '', + 'class Foo {', + '}', + ].join('\n'); + const uri = workspaceFile('src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const testParserService = new TestParserService([ + symbol(contents, 'Foo') + ]); + const parserService = asParserService(testParserService); + const linkifier = createTestLinkifierService(['src/file.ts'], parserService); + + const parts = (await linkify(linkifier, '[`Foo`](src/file.ts)')).parts; + + assert.strictEqual(parts.length, 1); + assert(parts[0] instanceof LinkifySymbolAnchor); + assert.strictEqual(parts[0].symbolInformation.location.uri.toString(), uri.toString()); + assert.strictEqual(parts[0].symbolInformation.location.range.start.line, 2); + assert.strictEqual(parts[0].symbolInformation.location.range.start.character, 6); + }); + + test(`Should keep the start-of-file fallback when tree-sitter does not find a symbol`, async () => { + const contents = [ + 'const value = 1;', + '', + 'class Bar {', + '}', + ].join('\n'); + const uri = workspaceFile('src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const testParserService = new TestParserService([ + symbol(contents, 'Bar') + ]); + const parserService = asParserService(testParserService); + const linkifier = createTestLinkifierService(['src/file.ts'], parserService); + + const parts = (await linkify(linkifier, '[`Foo`](src/file.ts)')).parts; + + assert.strictEqual(parts.length, 1); + assert(parts[0] instanceof LinkifySymbolAnchor); + assert.strictEqual(parts[0].symbolInformation.location.uri.toString(), uri.toString()); + assert.strictEqual(parts[0].symbolInformation.location.range.start.line, 0); + assert.strictEqual(parts[0].symbolInformation.location.range.start.character, 0); + }); + + test(`Should reuse the tree-sitter cache for multiple links to the same file`, async () => { + const contents = [ + 'class Foo {', + '}', + '', + 'class Bar {', + '}', + ].join('\n'); + const uri = workspaceFile('src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + + const testParserService = new TestParserService([ + symbol(contents, 'Foo'), + symbol(contents, 'Bar'), + ]); + const fs = createMockFsService(['src/file.ts']); + const linkifier = new SymbolLinkifier(fs, asParserService(testParserService), createMockWorkspaceService()); + + const result = await linkifier.linkify('[`Foo`](src/file.ts) [`Bar`](src/file.ts)', { requestId: undefined, references: [] }, CancellationToken.None); + + assert.strictEqual(testParserService.parseCount, 1); + assert(result); + const parts = result.parts; + assert.strictEqual(parts.length, 3); + assert(parts[0] instanceof LinkifySymbolAnchor); + assert(parts[2] instanceof LinkifySymbolAnchor); + assert.strictEqual(parts[0].symbolInformation.location.range.start.line, 0); + assert.strictEqual(parts[0].symbolInformation.location.range.start.character, 6); + assert.strictEqual(parts[2].symbolInformation.location.range.start.line, 3); + assert.strictEqual(parts[2].symbolInformation.location.range.start.character, 6); + }); + + test(`Should let LSP resolve upgrade symbol kind and location`, async () => { + const contents = [ + 'const value = 1;', + '', + 'class Foo {', + '}', + '', + 'class Foo {', + '}', + ].join('\n'); + const uri = workspaceFile('src/file.ts'); + setWorkspaceFileContents(new Map([[uri.toString(), contents]])); + testCommands.executeCommand = async (command: string, resolvedUri: vscode.Uri): Promise => { + assert.strictEqual(command, 'vscode.executeDocumentSymbolProvider'); + assert.strictEqual(resolvedUri.toString(), uri.toString()); + return [ + { + name: 'Foo', + detail: '', + kind: vscode.SymbolKind.Class, + range: new vscode.Range(5, 0, 6, 1), + selectionRange: new vscode.Range(5, 6, 5, 9), + children: [], + } + ] as T; + }; + + const linkifier = createTestLinkifierService(['src/file.ts'], asParserService(new TestParserService([ + symbol(contents, 'Foo') + ]))); + const parts = (await linkify(linkifier, '[`Foo`](src/file.ts)')).parts; + + assert.strictEqual(parts.length, 1); + assert(parts[0] instanceof LinkifySymbolAnchor); + assert.strictEqual(parts[0].symbolInformation.kind, vscode.SymbolKind.Variable); + assert.strictEqual(parts[0].symbolInformation.location.range.start.line, 2); + assert(parts[0].resolve); + + const resolved = await parts[0].resolve(CancellationToken.None); + + assert.strictEqual(resolved.kind, vscode.SymbolKind.Class); + assert.strictEqual(resolved.location.uri.toString(), uri.toString()); + assert.strictEqual(resolved.location.range.start.line, 5); + assert.strictEqual(resolved.location.range.start.character, 6); + }); }); diff --git a/src/extension/linkify/vscode-node/findSymbol.ts b/src/extension/linkify/vscode-node/findSymbol.ts index 3d517eebff..45c8340e1b 100644 --- a/src/extension/linkify/vscode-node/findSymbol.ts +++ b/src/extension/linkify/vscode-node/findSymbol.ts @@ -86,7 +86,7 @@ export function findBestSymbolByPath( * * We want just the names without any of the extra punctuation because `symbols` does not include these */ -function extractSymbolNamesInCode(inlineCode: string): string[] { +export function extractSymbolNamesInCode(inlineCode: string): string[] { // TODO: this assumes the language is JS like. // It won't handle symbol parts that include spaces or special characters return Array.from(inlineCode.matchAll(/[#\w$][\w\d$]*/g), x => x[0]); diff --git a/src/extension/linkify/vscode-node/findWord.ts b/src/extension/linkify/vscode-node/findWord.ts index 76272a7167..2c9048457a 100644 --- a/src/extension/linkify/vscode-node/findWord.ts +++ b/src/extension/linkify/vscode-node/findWord.ts @@ -15,6 +15,7 @@ import { escapeRegExpCharacters } from '../../../util/vs/base/common/strings'; import { isUriComponents, URI } from '../../../util/vs/base/common/uri'; import { IInstantiationService, ServicesAccessor } from '../../../util/vs/platform/instantiation/common/instantiation'; import { PromptReference } from '../../prompt/common/conversation'; +import { extractSymbolNamesInCode } from './findSymbol'; /** * How the word was resolved. @@ -43,6 +44,149 @@ interface FindWordOptions { readonly maxResultCount?: number; } +export interface FileSymbol { + readonly identifier: string; + readonly location: vscode.Location; +} + +export interface FileSymbols { + readonly declarations: readonly FileSymbol[]; + readonly getGenericSymbols: () => Promise; +} + +export type SymbolFileCache = Map>; + +export async function findSymbolLocationInFile( + parserService: IParserService, + uri: vscode.Uri, + symbolText: string, + token: CancellationToken, + cache?: SymbolFileCache +): Promise { + if (token.isCancellationRequested) { + return; + } + + const symbols = await getCachedFileSymbols(parserService, uri, token, cache); + if (token.isCancellationRequested) { + return; + } + + const exactMatch = findExactSymbol(symbols.declarations, symbolText); + if (exactMatch) { + return exactMatch.location; + } + + const symbolParts = extractSymbolNamesInCode(symbolText); + if (symbolParts.length) { + const declarationMatch = findBestSymbolPart(symbols.declarations, symbolParts); + if (declarationMatch) { + return declarationMatch.location; + } + } + + const genericSymbols = await symbols.getGenericSymbols(); + if (token.isCancellationRequested) { + return; + } + + return findExactSymbol(genericSymbols, symbolText)?.location + ?? (symbolParts.length ? findBestSymbolPart(genericSymbols, symbolParts)?.location : undefined); +} + +function findExactSymbol(symbols: readonly FileSymbol[], symbolText: string): FileSymbol | undefined { + return symbols.find(symbol => symbol.identifier === symbolText); +} + +function findBestSymbolPart(symbols: readonly FileSymbol[], symbolParts: readonly string[]): FileSymbol | undefined { + let bestMatch: { symbol: FileSymbol; matchIndex: number } | undefined; + for (const symbol of symbols) { + const matchIndex = symbolParts.indexOf(symbol.identifier); + if (matchIndex !== -1 && (!bestMatch || matchIndex > bestMatch.matchIndex)) { + bestMatch = { symbol, matchIndex }; + } + } + + return bestMatch?.symbol; +} + +async function getCachedFileSymbols( + parserService: IParserService, + uri: vscode.Uri, + token: CancellationToken, + cache: SymbolFileCache | undefined +): Promise { + const key = uri.toString(); + const existing = cache?.get(key); + if (existing) { + return existing; + } + + const pending = doGetFileSymbols(parserService, uri, token); + cache?.set(key, pending); + return pending; +} + +async function doGetFileSymbols(parserService: IParserService, uri: vscode.Uri, token: CancellationToken): Promise { + const languageId = getLanguageForResource(uri).languageId; + const wasmLanguage = getWasmLanguage(languageId); + if (!wasmLanguage) { + return emptyFileSymbols(); + } + + const doc = await openDocument(uri); + if (!doc || token.isCancellationRequested) { + return emptyFileSymbols(); + } + + try { + const ast = parserService.getTreeSitterASTForWASMLanguage(wasmLanguage, doc.getText()); + const [classDeclarations, functionDefinitions, typeDeclarations] = await Promise.all([ + ast.getClassDeclarations(), + ast.getFunctionDefinitions(), + ast.getTypeDeclarations(), + ]); + let genericSymbols: Promise | undefined; + return { + declarations: toFileSymbols(uri, doc, [ + ...classDeclarations, + ...functionDefinitions, + ...typeDeclarations, + ]), + getGenericSymbols: () => { + if (token.isCancellationRequested) { + return Promise.resolve([]); + } + genericSymbols ??= (async () => { + const fullFileRange = new vscode.Range(0, 0, Number.MAX_SAFE_INTEGER, 0); + const symbols = await ast.getSymbols({ + startIndex: doc.offsetAt(fullFileRange.start), + endIndex: doc.offsetAt(fullFileRange.end), + }); + return toFileSymbols(uri, doc, symbols); + })().catch(() => []); + return genericSymbols; + }, + }; + } catch { + return emptyFileSymbols(); + } +} + +function toFileSymbols(uri: vscode.Uri, doc: SimpleTextDocument, symbols: readonly TreeSitterExpressionInfo[]): readonly FileSymbol[] { + return symbols.map(symbol => ({ + identifier: symbol.identifier, + location: new vscode.Location(uri, doc.positionAt(symbol.startIndex)) + })); +} + +function emptyFileSymbols(): FileSymbols { + return { + declarations: [], + getGenericSymbols: async () => [] + }; +} + export async function findWordInReferences( accessor: ServicesAccessor, references: readonly PromptReference[], @@ -247,7 +391,7 @@ export class ReferencesSymbolResolver { if (!wordMatches.length) { // Extract all symbol parts from the code text // For example: `TextModel.undo()` -> ['TextModel', 'undo'] - const symbolParts = Array.from(codeText.matchAll(/[#\w$][\w\d$]*/g), x => x[0]); + const symbolParts = extractSymbolNamesInCode(codeText); if (symbolParts.length >= 2) { // For qualified names like `Class.method()`, search for both parts together diff --git a/src/extension/linkify/vscode-node/symbolLinkifier.ts b/src/extension/linkify/vscode-node/symbolLinkifier.ts index 3ecaf47239..7ed24c18b7 100644 --- a/src/extension/linkify/vscode-node/symbolLinkifier.ts +++ b/src/extension/linkify/vscode-node/symbolLinkifier.ts @@ -5,6 +5,7 @@ import * as vscode from 'vscode'; import { IFileSystemService } from '../../../platform/filesystem/common/fileSystemService'; +import { IParserService } from '../../../platform/parser/node/parserService'; import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService'; import { collapseRangeToStart } from '../../../util/common/range'; import { CancellationToken } from '../../../util/vs/base/common/cancellation'; @@ -12,6 +13,7 @@ import { SymbolInformation, Uri } from '../../../vscodeTypes'; import { LinkifiedPart, LinkifiedText, LinkifySymbolAnchor } from '../common/linkifiedText'; import { IContributedLinkifier, LinkifierContext } from '../common/linkifyService'; import { findBestSymbolByPath } from './findSymbol'; +import { findSymbolLocationInFile, type SymbolFileCache } from './findWord'; /** * Linkifies symbol paths in responses. For example: @@ -24,6 +26,7 @@ export class SymbolLinkifier implements IContributedLinkifier { constructor( @IFileSystemService private readonly fileSystem: IFileSystemService, + @IParserService private readonly parserService: IParserService, @IWorkspaceService private readonly workspaceService: IWorkspaceService, ) { } @@ -38,6 +41,7 @@ export class SymbolLinkifier implements IContributedLinkifier { } const out: LinkifiedPart[] = []; + const symbolFileCache: SymbolFileCache = new Map(); let endLastMatch = 0; for (const match of text.matchAll(/\[`([^`\[\]]+?)`]\((\S+?\.\w+)\)/g)) { @@ -57,21 +61,25 @@ export class SymbolLinkifier implements IContributedLinkifier { const resolvedUri = await this.resolveInWorkspace(symbolPath, workspaceFolders); if (resolvedUri) { + const initialLocation = await findSymbolLocationInFile(this.parserService, resolvedUri, symbolText, token, symbolFileCache) + .catch(() => undefined); const info: SymbolInformation = { name: symbolText, containerName: '', kind: vscode.SymbolKind.Variable, - location: new vscode.Location(resolvedUri, new vscode.Position(0, 0)) + location: initialLocation ?? new vscode.Location(resolvedUri, new vscode.Position(0, 0)) }; out.push(new LinkifySymbolAnchor(info, async (token) => { let symbols: Array | undefined; try { symbols = await vscode.commands.executeCommand | undefined>('vscode.executeDocumentSymbolProvider', resolvedUri); - } catch (e) { - // Noop + } catch { + // noop } + // Tree-sitter gives a best-effort initial location. Document symbols remain + // the richer source for symbol kind and nested same-name disambiguation. if (symbols?.length) { const matchingSymbol = findBestSymbolByPath(symbols, symbolText); if (matchingSymbol) {