Skip to content

Commit

Permalink
Add event for when tree tokens change
Browse files Browse the repository at this point in the history
Part of #210475
  • Loading branch information
alexr00 committed Sep 13, 2024
1 parent 03dc521 commit c0aab60
Show file tree
Hide file tree
Showing 7 changed files with 55 additions and 6 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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`;
Expand All @@ -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; }
Expand Down Expand Up @@ -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)));
}
}
}

Expand Down Expand Up @@ -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,
Expand All @@ -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() {
Expand Down Expand Up @@ -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) {
Expand Down
2 changes: 2 additions & 0 deletions src/vs/editor/common/languages.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 }>;
}

/**
Expand Down
7 changes: 7 additions & 0 deletions src/vs/editor/common/model/treeSitterTokens.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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);
}
});
}
}

Expand Down
2 changes: 2 additions & 0 deletions src/vs/editor/common/services/treeSitterParserService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand All @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
2 changes: 2 additions & 0 deletions src/vs/editor/test/common/services/testTreeSitterService.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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 {
Expand Down

0 comments on commit c0aab60

Please sign in to comment.