Add event for when tree tokens change (#228593)

Part of #210475
microsoft · Sep 18, 2024 · 3ea0cbb · 3ea0cbb
1 parent a5f5206
commit 3ea0cbb
Show file tree

Hide file tree

Showing 7 changed files with 55 additions and 6 deletions.
diff --git a/src/vs/editor/browser/services/treeSitter/treeSitterParserService.ts b/src/vs/editor/browser/services/treeSitter/treeSitterParserService.ts
@@ -22,6 +22,7 @@ import { IEnvironmentService } from '../../../../platform/environment/common/env
 import { canASAR } from '../../../../base/common/amd.js';
 import { CancellationError, isCancellationError } from '../../../../base/common/errors.js';
 import { PromiseResult } from '../../../../base/common/observable.js';
+import { Range } from '../../../common/core/range.js';
 
 const EDITOR_TREESITTER_TELEMETRY = 'editor.experimental.treeSitterTelemetry';
 const MODULE_LOCATION_SUBPATH = `@vscode/tree-sitter-wasm/wasm`;
@@ -32,6 +33,8 @@ function getModuleLocation(environmentService: IEnvironmentService): AppResource
 }
 
 export class TextModelTreeSitter extends Disposable {
+	private _onDidChangeParseResult: Emitter<Range[]> = this._register(new Emitter<Range[]>());
+	public readonly onDidChangeParseResult: Event<Range[]> = this._onDidChangeParseResult.event;
 	private _parseResult: TreeSitterParseResult | undefined;
 
 	get parseResult(): ITreeSitterParseResult | undefined { return this._parseResult; }
@@ -102,7 +105,13 @@ export class TextModelTreeSitter extends Disposable {
 	}
 
 	private async _onDidChangeContent(treeSitterTree: TreeSitterParseResult, changes: IModelContentChange[]) {
-		return treeSitterTree.onDidChangeContent(this.model, changes);
+		const oldTree = treeSitterTree.tree?.copy();
+		await treeSitterTree.onDidChangeContent(this.model, changes);
+		if (oldTree && treeSitterTree.tree) {
+			const diff = oldTree.getChangedRanges(treeSitterTree.tree);
+			// Tree sitter is 0 based, text model is 1 based
+			this._onDidChangeParseResult.fire(diff.map(r => new Range(r.startPosition.row + 1, r.startPosition.column + 1, r.endPosition.row + 1, r.endPosition.column + 1)));
+		}
 	}
 }
 
@@ -312,15 +321,23 @@ export class TreeSitterImporter {
 	}
 }
 
+interface TextModelTreeSitterItem {
+	dispose(): void;
+	textModelTreeSitter: TextModelTreeSitter;
+	disposables: DisposableStore;
+}
+
 export class TreeSitterTextModelService extends Disposable implements ITreeSitterParserService {
 	readonly _serviceBrand: undefined;
 	private _init!: Promise<boolean>;
-	private _textModelTreeSitters: DisposableMap<ITextModel, TextModelTreeSitter> = this._register(new DisposableMap());
+	private _textModelTreeSitters: DisposableMap<ITextModel, TextModelTreeSitterItem> = this._register(new DisposableMap());
 	private readonly _registeredLanguages: Map<string, string> = new Map();
 	private readonly _treeSitterImporter: TreeSitterImporter = new TreeSitterImporter();
 	private readonly _treeSitterLanguages: TreeSitterLanguages;
 
 	public readonly onDidAddLanguage: Event<{ id: string; language: Parser.Language }>;
+	private _onDidUpdateTree: Emitter<{ textModel: ITextModel; ranges: Range[] }> = this._register(new Emitter());
+	public readonly onDidUpdateTree: Event<{ textModel: ITextModel; ranges: Range[] }> = this._onDidUpdateTree.event;
 
 	constructor(@IModelService private readonly _modelService: IModelService,
 		@IFileService fileService: IFileService,
@@ -346,7 +363,7 @@ export class TreeSitterTextModelService extends Disposable implements ITreeSitte
 
 	getParseResult(textModel: ITextModel): ITreeSitterParseResult | undefined {
 		const textModelTreeSitter = this._textModelTreeSitters.get(textModel);
-		return textModelTreeSitter?.parseResult;
+		return textModelTreeSitter?.textModelTreeSitter.parseResult;
 	}
 
 	private async _doInitParser() {
@@ -421,7 +438,14 @@ export class TreeSitterTextModelService extends Disposable implements ITreeSitte
 
 	private _createTextModelTreeSitter(model: ITextModel) {
 		const textModelTreeSitter = new TextModelTreeSitter(model, this._treeSitterLanguages, this._treeSitterImporter, this._logService, this._telemetryService);
-		this._textModelTreeSitters.set(model, textModelTreeSitter);
+		const disposables = new DisposableStore();
+		disposables.add(textModelTreeSitter);
+		disposables.add(textModelTreeSitter.onDidChangeParseResult((ranges) => this._onDidUpdateTree.fire({ textModel: model, ranges })));
+		this._textModelTreeSitters.set(model, {
+			textModelTreeSitter,
+			disposables,
+			dispose: disposables.dispose
+		});
 	}
 
 	private _addGrammar(languageId: string, grammarName: string) {

diff --git a/src/vs/editor/common/languages.ts b/src/vs/editor/common/languages.ts
@@ -26,6 +26,7 @@ import { ContiguousMultilineTokens } from './tokens/contiguousMultilineTokens.js
 import { localize } from '../../nls.js';
 import { ExtensionIdentifier } from '../../platform/extensions/common/extensions.js';
 import { IMarkerData } from '../../platform/markers/common/markers.js';
+import { IModelTokensChangedEvent } from './textModelEvents.js';
 
 /**
  * @internal
@@ -89,6 +90,7 @@ export class EncodedTokenizationResult {
 export interface ITreeSitterTokenizationSupport {
 	tokenizeEncoded(lineNumber: number, textModel: model.ITextModel): Uint32Array | undefined;
 	captureAtPosition(lineNumber: number, column: number, textModel: model.ITextModel): any;
+	onDidChangeTokens: Event<{ textModel: model.ITextModel; changes: IModelTokensChangedEvent }>;
 }
 
 /**

diff --git a/src/vs/editor/common/model/treeSitterTokens.ts b/src/vs/editor/common/model/treeSitterTokens.ts
@@ -11,10 +11,12 @@ import { ITreeSitterParserService } from '../services/treeSitterParserService.js
 import { IModelContentChangedEvent } from '../textModelEvents.js';
 import { AbstractTokens } from './tokens.js';
 import { ITokenizeLineWithEditResult, LineEditWithAdditionalLines } from '../tokenizationTextModelPart.js';
+import { IDisposable, MutableDisposable } from '../../../base/common/lifecycle.js';
 
 export class TreeSitterTokens extends AbstractTokens {
 	private _tokenizationSupport: ITreeSitterTokenizationSupport | null = null;
 	private _lastLanguageId: string | undefined;
+	private readonly _tokensChangedListener: MutableDisposable<IDisposable> = this._register(new MutableDisposable());
 
 	constructor(private readonly _treeSitterService: ITreeSitterParserService,
 		languageIdCodec: ILanguageIdCodec,
@@ -30,6 +32,11 @@ export class TreeSitterTokens extends AbstractTokens {
 		if (!this._tokenizationSupport || this._lastLanguageId !== newLanguage) {
 			this._lastLanguageId = newLanguage;
 			this._tokenizationSupport = TreeSitterTokenizationRegistry.get(newLanguage);
+			this._tokensChangedListener.value = this._tokenizationSupport?.onDidChangeTokens((e) => {
+				if (e.textModel === this._textModel) {
+					this._onDidChangeTokens.fire(e.changes);
+				}
+			});
 		}
 	}
 

diff --git a/src/vs/editor/common/services/treeSitterParserService.ts b/src/vs/editor/common/services/treeSitterParserService.ts
@@ -7,6 +7,7 @@ import type { Parser } from '@vscode/tree-sitter-wasm';
 import { Event } from '../../../base/common/event.js';
 import { ITextModel } from '../model.js';
 import { createDecorator } from '../../../platform/instantiation/common/instantiation.js';
+import { Range } from '../core/range.js';
 
 export const EDITOR_EXPERIMENTAL_PREFER_TREESITTER = 'editor.experimental.preferTreeSitter';
 
@@ -17,6 +18,7 @@ export interface ITreeSitterParserService {
 	onDidAddLanguage: Event<{ id: string; language: Parser.Language }>;
 	getOrInitLanguage(languageId: string): Parser.Language | undefined;
 	getParseResult(textModel: ITextModel): ITreeSitterParseResult | undefined;
+	onDidUpdateTree: Event<{ textModel: ITextModel; ranges: Range[] }>;
 }
 
 export interface ITreeSitterParseResult {

diff --git a/src/vs/editor/standalone/browser/standaloneTreeSitterService.ts b/src/vs/editor/standalone/browser/standaloneTreeSitterService.ts
@@ -7,12 +7,14 @@ import type { Parser } from '@vscode/tree-sitter-wasm';
 import { Event } from '../../../base/common/event.js';
 import { ITextModel } from '../../common/model.js';
 import { ITreeSitterParseResult, ITreeSitterParserService } from '../../common/services/treeSitterParserService.js';
+import { Range } from '../../common/core/range.js';
 
 /**
  * The monaco build doesn't like the dynamic import of tree sitter in the real service.
  * We use a dummy sertive here to make the build happy.
  */
 export class StandaloneTreeSitterParserService implements ITreeSitterParserService {
+	onDidUpdateTree: Event<{ textModel: ITextModel; ranges: Range[] }> = Event.None;
 	readonly _serviceBrand: undefined;
 	onDidAddLanguage: Event<{ id: string; language: Parser.Language }> = Event.None;
 

diff --git a/src/vs/editor/test/common/services/testTreeSitterService.ts b/src/vs/editor/test/common/services/testTreeSitterService.ts
@@ -7,8 +7,10 @@ import type { Parser } from '@vscode/tree-sitter-wasm';
 import { Event } from '../../../../base/common/event.js';
 import { ITextModel } from '../../../common/model.js';
 import { ITreeSitterParserService, ITreeSitterParseResult } from '../../../common/services/treeSitterParserService.js';
+import { Range } from '../../../common/core/range.js';
 
 export class TestTreeSitterParserService implements ITreeSitterParserService {
+	onDidUpdateTree: Event<{ textModel: ITextModel; ranges: Range[] }> = Event.None;
 	onDidAddLanguage: Event<{ id: string; language: Parser.Language }> = Event.None;
 	_serviceBrand: undefined;
 	getOrInitLanguage(languageId: string): Parser.Language | undefined {

diff --git a/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.ts b/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.ts
@@ -87,8 +87,8 @@ class TreeSitterTokenizationFeature extends Disposable implements ITreeSitterTok
 
 class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTokenizationSupport {
 	private _query: Parser.Query | undefined;
-	private readonly _onDidChangeTokens: Emitter<IModelTokensChangedEvent> = new Emitter();
-	public readonly onDidChangeTokens: Event<IModelTokensChangedEvent> = this._onDidChangeTokens.event;
+	private readonly _onDidChangeTokens: Emitter<{ textModel: ITextModel; changes: IModelTokensChangedEvent }> = new Emitter();
+	public readonly onDidChangeTokens: Event<{ textModel: ITextModel; changes: IModelTokensChangedEvent }> = this._onDidChangeTokens.event;
 	private _colorThemeData!: ColorThemeData;
 	private _languageAddedListener: IDisposable | undefined;
 
@@ -100,6 +100,16 @@ class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTok
 	) {
 		super();
 		this._register(Event.runAndSubscribe(this._themeService.onDidColorThemeChange, () => this.reset()));
+		this._register(this._treeSitterService.onDidUpdateTree((e) => {
+			const maxLine = e.textModel.getLineCount();
+			this._onDidChangeTokens.fire({
+				textModel: e.textModel,
+				changes: {
+					semanticTokensApplied: false,
+					ranges: e.ranges.map(range => ({ fromLineNumber: range.startLineNumber, toLineNumber: range.endLineNumber < maxLine ? range.endLineNumber : maxLine })),
+				}
+			});
+		}));
 	}
 
 	private _getTree(textModel: ITextModel): ITreeSitterParseResult | undefined {