@codingame/monaco-vscode-treesitter-service-override 13.1.7 → 14.0.0

package/index.js

@@ -1,10 +1,11 @@
 
 import { SyncDescriptor } from '@codingame/monaco-vscode-api/vscode/vs/platform/instantiation/common/descriptors';
 import { ITreeSitterParserService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService';
-import { TreeSitterTextModelService } from './vscode/src/vs/editor/browser/services/treeSitter/treeSitterParserService.js';
+import { TreeSitterTextModelService } from './vscode/src/vs/editor/common/services/treeSitter/treeSitterParserService.js';
 import { ITreeSitterTokenizationFeature } from '@codingame/monaco-vscode-api/vscode/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.service';
 import { TreeSitterTokenizationFeature } from './vscode/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.js';
 import { registerAssets } from '@codingame/monaco-vscode-api/assets';
+import { ITreeSitterTokenizationStoreService, TreeSitterTokenizationStoreService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/model/treeSitterTokenStoreService';
 import './vscode/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.contribution.js';
 
 registerAssets({
@@ -15,7 +16,8 @@ registerAssets({
 function getServiceOverride() {
     return {
         [ITreeSitterParserService.toString()]: new SyncDescriptor(TreeSitterTextModelService, [], false),
-        [ITreeSitterTokenizationFeature.toString()]: new SyncDescriptor(TreeSitterTokenizationFeature, [], false)
+        [ITreeSitterTokenizationFeature.toString()]: new SyncDescriptor(TreeSitterTokenizationFeature, [], false),
+        [ITreeSitterTokenizationStoreService.toString()]: new SyncDescriptor(TreeSitterTokenizationStoreService, [], false)
     };
 }
 

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@codingame/monaco-vscode-treesitter-service-override",
-  "version": "13.1.7",
+  "version": "14.0.0",
   "private": false,
   "description": "VSCode public API plugged on the monaco editor - treesitter service-override",
   "keywords": [],
@@ -15,9 +15,10 @@
   },
   "type": "module",
   "dependencies": {
-    "@codingame/monaco-vscode-9d0168a3-519b-57f3-9bcc-89efc41f951a-common": "13.1.7",
-    "@codingame/monaco-vscode-api": "13.1.7",
-    "@vscode/tree-sitter-wasm": "0.0.4"
+    "@codingame/monaco-vscode-168b98e5-dc20-5807-b1f9-798f1f92b37f-common": "14.0.0",
+    "@codingame/monaco-vscode-9d0168a3-519b-57f3-9bcc-89efc41f951a-common": "14.0.0",
+    "@codingame/monaco-vscode-api": "14.0.0",
+    "@vscode/tree-sitter-wasm": "0.0.5"
   },
   "main": "index.js",
   "module": "index.js",

package/vscode/src/vs/editor/{browser → common}/services/treeSitter/treeSitterParserService.d.ts

@@ -1,16 +1,15 @@
 import type { Parser } from "@vscode/tree-sitter-wasm";
-import { ITreeSitterParserService, ITreeSitterParseResult, ITextModelTreeSitter } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService";
+import { ITreeSitterParserService, ITreeSitterParseResult, ITextModelTreeSitter, TreeUpdateEvent, TreeParseUpdateEvent } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService";
 import { IModelService } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/services/model";
 import { Disposable, IDisposable } from "@codingame/monaco-vscode-api/vscode/vs/base/common/lifecycle";
 import { ITextModel } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/model";
 import { IFileService } from "@codingame/monaco-vscode-api/vscode/vs/platform/files/common/files.service";
-import { IModelContentChange } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/textModelEvents";
+import { IModelContentChangedEvent } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/textModelEvents";
 import { ITelemetryService } from "@codingame/monaco-vscode-api/vscode/vs/platform/telemetry/common/telemetry.service";
 import { ILogService } from "@codingame/monaco-vscode-api/vscode/vs/platform/log/common/log.service";
 import { IConfigurationService } from "@codingame/monaco-vscode-api/vscode/vs/platform/configuration/common/configuration.service";
 import { Emitter, Event } from "@codingame/monaco-vscode-api/vscode/vs/base/common/event";
 import { IEnvironmentService } from "@codingame/monaco-vscode-api/vscode/vs/platform/environment/common/environment.service";
-import { Range } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/core/range";
 export declare class TextModelTreeSitter extends Disposable implements ITextModelTreeSitter {
     readonly model: ITextModel;
     private readonly _treeSitterLanguages;
@@ -18,11 +17,12 @@ export declare class TextModelTreeSitter extends Disposable implements ITextMode
     private readonly _logService;
     private readonly _telemetryService;
     private _onDidChangeParseResult;
-    readonly onDidChangeParseResult: Event<Range[]>;
+    readonly onDidChangeParseResult: Event<TreeParseUpdateEvent>;
     private _parseResult;
+    private _versionId;
     get parseResult(): ITreeSitterParseResult | undefined;
     constructor(model: ITextModel, _treeSitterLanguages: TreeSitterLanguages, _treeSitterImporter: TreeSitterImporter, _logService: ILogService, _telemetryService: ITelemetryService, parseImmediately?: boolean);
-    private readonly _languageSessionDisposables;
+    private readonly _parseSessionDisposables;
     private _onDidChangeLanguage;
     parse(languageId?: string): Promise<ITreeSitterParseResult | undefined>;
     private _getLanguage;
@@ -34,15 +34,22 @@ export declare class TreeSitterParseResult implements IDisposable, ITreeSitterPa
     private readonly _logService;
     private readonly _telemetryService;
     private _tree;
+    private _lastFullyParsed;
+    private _lastFullyParsedWithEdits;
+    private readonly _onDidUpdate;
+    readonly onDidUpdate: Event<TreeParseUpdateEvent>;
+    private _versionId;
+    private _editVersion;
+    get versionId(): number;
     private _isDisposed;
     constructor(parser: Parser, language: Parser.Language, _logService: ILogService, _telemetryService: ITelemetryService);
     dispose(): void;
     get tree(): Parser.Tree | undefined;
-    private set tree(value);
     get isDisposed(): boolean;
+    private findChangedNodes;
+    private calculateRangeChange;
     private _onDidChangeContentQueue;
-    onDidChangeContent(model: ITextModel, changes: IModelContentChange[]): Promise<Parser.Range[] | undefined>;
-    private _newEdits;
+    onDidChangeContent(model: ITextModel, changes: IModelContentChangedEvent | undefined): void;
     private _applyEdits;
     private _parseAndUpdateTree;
     private _parse;
@@ -94,21 +101,20 @@ export declare class TreeSitterTextModelService extends Disposable implements IT
         language: Parser.Language;
     }>;
     private _onDidUpdateTree;
-    readonly onDidUpdateTree: Event<{
-        textModel: ITextModel;
-        ranges: Range[];
-    }>;
+    readonly onDidUpdateTree: Event<TreeUpdateEvent>;
+    isTest: boolean;
     constructor(_modelService: IModelService, fileService: IFileService, _telemetryService: ITelemetryService, _logService: ILogService, _configurationService: IConfigurationService, _environmentService: IEnvironmentService);
     getOrInitLanguage(languageId: string): Parser.Language | undefined;
     getParseResult(textModel: ITextModel): ITreeSitterParseResult | undefined;
     getTree(content: string, languageId: string): Promise<Parser.Tree | undefined>;
+    getLanguage(languageId: string): Promise<Parser.Language | undefined>;
     private _doInitParser;
     private _hasInit;
     private _initParser;
     private _supportedLanguagesChanged;
     private _getSetting;
     private _registerModelServiceListeners;
-    getTextModelTreeSitter(model: ITextModel): ITextModelTreeSitter;
+    getTextModelTreeSitter(model: ITextModel, parseImmediately?: boolean): Promise<ITextModelTreeSitter>;
     private _createTextModelTreeSitter;
     private _addGrammar;
     private _removeGrammar;

package/vscode/src/vs/editor/{browser → common}/services/treeSitter/treeSitterParserService.js

@@ -16,9 +16,12 @@ import { isCancellationError, CancellationError } from '@codingame/monaco-vscode
 import '@codingame/monaco-vscode-api/vscode/vs/base/common/arrays';
 import '@codingame/monaco-vscode-api/vscode/vs/base/common/observableInternal/autorun';
 import '@codingame/monaco-vscode-api/vscode/vs/base/common/observableInternal/derived';
-import { PromiseResult } from '@codingame/monaco-vscode-api/vscode/vs/base/common/observableInternal/promise';
+import { PromiseResult } from '@codingame/monaco-vscode-168b98e5-dc20-5807-b1f9-798f1f92b37f-common/vscode/vs/base/common/observableInternal/promise';
 import '@codingame/monaco-vscode-api/vscode/vs/base/common/observableInternal/utils';
 import { Range } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/core/range';
+import { Position } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/core/position';
+import { LimitedQueue } from '@codingame/monaco-vscode-api/vscode/vs/base/common/async';
+import { TextLength } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/core/textLength';
 
 const EDITOR_TREESITTER_TELEMETRY = 'editor.experimental.treeSitterTelemetry';
 const MODULE_LOCATION_SUBPATH = `@vscode/tree-sitter-wasm/wasm`;
@@ -37,7 +40,8 @@ class TextModelTreeSitter extends Disposable {
         this._telemetryService = _telemetryService;
         this._onDidChangeParseResult = this._register(( new Emitter()));
         this.onDidChangeParseResult = this._onDidChangeParseResult.event;
-        this._languageSessionDisposables = this._register(( new DisposableStore()));
+        this._versionId = 0;
+        this._parseSessionDisposables = this._register(( new DisposableStore()));
         if (parseImmediately) {
             this._register(Event.runAndSubscribe(this.model.onDidChangeLanguage, (e => this._onDidChangeLanguage(e ? e.newLanguage : this.model.getLanguageId()))));
         }
@@ -49,9 +53,9 @@ class TextModelTreeSitter extends Disposable {
         this.parse(languageId);
     }
     async parse(languageId = this.model.getLanguageId()) {
-        this._languageSessionDisposables.clear();
+        this._parseSessionDisposables.clear();
         this._parseResult = undefined;
-        const token = cancelOnDispose(this._languageSessionDisposables);
+        const token = cancelOnDispose(this._parseSessionDisposables);
         let language;
         try {
             language = await this._getLanguage(languageId, token);
@@ -66,13 +70,19 @@ class TextModelTreeSitter extends Disposable {
         if (token.isCancellationRequested) {
             return;
         }
-        const treeSitterTree = this._languageSessionDisposables.add(( new TreeSitterParseResult(( new Parser()), language, this._logService, this._telemetryService)));
-        this._languageSessionDisposables.add(this.model.onDidChangeContent(e => this._onDidChangeContent(treeSitterTree, e.changes)));
-        await this._onDidChangeContent(treeSitterTree, []);
+        const treeSitterTree = this._parseSessionDisposables.add(( new TreeSitterParseResult(( new Parser()), language, this._logService, this._telemetryService)));
+        this._parseResult = treeSitterTree;
+        this._parseSessionDisposables.add(treeSitterTree.onDidUpdate(e => {
+            if (e.ranges && (e.versionId > this._versionId)) {
+                this._versionId = e.versionId;
+                this._onDidChangeParseResult.fire({ ranges: e.ranges, versionId: e.versionId });
+            }
+        }));
+        this._parseSessionDisposables.add(this.model.onDidChangeContent(e => this._onDidChangeContent(treeSitterTree, e)));
+        this._onDidChangeContent(treeSitterTree, undefined);
         if (token.isCancellationRequested) {
             return;
         }
-        this._parseResult = treeSitterTree;
         return this._parseResult;
     }
     _getLanguage(languageId, token) {
@@ -94,17 +104,8 @@ class TextModelTreeSitter extends Disposable {
             }, undefined, disposables);
         }));
     }
-    async _onDidChangeContent(treeSitterTree, changes) {
-        const diff = await treeSitterTree.onDidChangeContent(this.model, changes);
-        if (!diff || diff.length > 0) {
-            const ranges = diff ? ( diff.map(r => ( new Range(
-                r.startPosition.row + 1,
-                r.startPosition.column + 1,
-                r.endPosition.row + 1,
-                r.endPosition.column + 1
-            )))) : [this.model.getFullModelRange()];
-            this._onDidChangeParseResult.fire(ranges);
-        }
+    _onDidChangeContent(treeSitterTree, change) {
+        return treeSitterTree.onDidChangeContent(this.model, change);
     }
 }
 var TelemetryParseType;
@@ -113,63 +114,246 @@ var TelemetryParseType;
     TelemetryParseType["Incremental"] = "incrementalParse";
 })(TelemetryParseType || (TelemetryParseType = {}));
 class TreeSitterParseResult {
+    get versionId() {
+        return this._versionId;
+    }
     constructor(parser, language, _logService, _telemetryService) {
         this.parser = parser;
         this.language = language;
         this._logService = _logService;
         this._telemetryService = _telemetryService;
+        this._onDidUpdate = ( new Emitter());
+        this.onDidUpdate = this._onDidUpdate.event;
+        this._versionId = 0;
+        this._editVersion = 0;
         this._isDisposed = false;
-        this._onDidChangeContentQueue = Promise.resolve();
-        this._newEdits = true;
+        this._onDidChangeContentQueue = ( new LimitedQueue());
         this.parser.setTimeoutMicros(50 * 1000);
         this.parser.setLanguage(language);
     }
     dispose() {
         this._isDisposed = true;
+        this._onDidUpdate.dispose();
         this._tree?.delete();
+        this._lastFullyParsed?.delete();
+        this._lastFullyParsedWithEdits?.delete();
         this.parser?.delete();
     }
-    get tree() { return this._tree; }
-    set tree(newTree) {
-        this._tree?.delete();
-        this._tree = newTree;
-    }
+    get tree() { return this._lastFullyParsed; }
     get isDisposed() { return this._isDisposed; }
-    async onDidChangeContent(model, changes) {
-        const oldTree = this.tree?.copy();
-        this._applyEdits(model, changes);
-        return ( new Promise(resolve => {
-            this._onDidChangeContentQueue = this._onDidChangeContentQueue.then(async () => {
-                if (this.isDisposed) {
-                    return;
+    findChangedNodes(newTree, oldTree) {
+        const newCursor = newTree.walk();
+        const oldCursor = oldTree.walk();
+        const gotoNextSibling = () => {
+            const n = newCursor.gotoNextSibling();
+            const o = oldCursor.gotoNextSibling();
+            if (n !== o) {
+                throw ( new Error('Trees are out of sync'));
+            }
+            return n && o;
+        };
+        const gotoParent = () => {
+            const n = newCursor.gotoParent();
+            const o = oldCursor.gotoParent();
+            if (n !== o) {
+                throw ( new Error('Trees are out of sync'));
+            }
+            return n && o;
+        };
+        const gotoNthChild = (index) => {
+            const n = newCursor.gotoFirstChild();
+            const o = oldCursor.gotoFirstChild();
+            if (n !== o) {
+                throw ( new Error('Trees are out of sync'));
+            }
+            if (index === 0) {
+                return n && o;
+            }
+            for (let i = 1; i <= index; i++) {
+                const nn = newCursor.gotoNextSibling();
+                const oo = oldCursor.gotoNextSibling();
+                if (nn !== oo) {
+                    throw ( new Error('Trees are out of sync'));
                 }
-                await this._parseAndUpdateTree(model);
-                resolve((this.tree && oldTree) ? oldTree.getChangedRanges(this.tree) : undefined);
-            }).catch((e) => {
-                this._logService.error('Error parsing tree-sitter tree', e);
-            });
-        }));
+                if (!nn || !oo) {
+                    return false;
+                }
+            }
+            return n && o;
+        };
+        const changedRanges = [];
+        let next = true;
+        const nextSiblingOrParentSibling = () => {
+            do {
+                if (newCursor.currentNode.nextSibling) {
+                    return gotoNextSibling();
+                }
+                if (newCursor.currentNode.parent) {
+                    gotoParent();
+                }
+            } while (newCursor.currentNode.nextSibling || newCursor.currentNode.parent);
+            return false;
+        };
+        const getClosestPreviousNodes = () => {
+            const newFindPrev = newTree.walk();
+            newFindPrev.resetTo(newCursor);
+            const oldFindPrev = oldTree.walk();
+            oldFindPrev.resetTo(oldCursor);
+            const startingNode = newCursor.currentNode;
+            do {
+                if (newFindPrev.currentNode.previousSibling && ((newFindPrev.currentNode.endIndex - newFindPrev.currentNode.startIndex) !== 0)) {
+                    newFindPrev.gotoPreviousSibling();
+                    oldFindPrev.gotoPreviousSibling();
+                }
+                else {
+                    while (!newFindPrev.currentNode.previousSibling && newFindPrev.currentNode.parent) {
+                        newFindPrev.gotoParent();
+                        oldFindPrev.gotoParent();
+                    }
+                    newFindPrev.gotoPreviousSibling();
+                    oldFindPrev.gotoPreviousSibling();
+                }
+            } while ((newFindPrev.currentNode.endIndex > startingNode.startIndex)
+                && (newFindPrev.currentNode.parent || newFindPrev.currentNode.previousSibling)
+                && (newFindPrev.currentNode.id !== startingNode.id));
+            if ((newFindPrev.currentNode.id !== startingNode.id) && newFindPrev.currentNode.endIndex <= startingNode.startIndex) {
+                return { old: oldFindPrev.currentNode, new: newFindPrev.currentNode };
+            }
+            else {
+                return undefined;
+            }
+        };
+        do {
+            if (newCursor.currentNode.hasChanges) {
+                const newChildren = newCursor.currentNode.children;
+                const indexChangedChildren = [];
+                const changedChildren = newChildren.filter((c, index) => {
+                    if (c.hasChanges) {
+                        indexChangedChildren.push(index);
+                    }
+                    return c.hasChanges;
+                });
+                if (changedChildren.length >= 1) {
+                    next = gotoNthChild(indexChangedChildren[0]);
+                }
+                else if (changedChildren.length === 0) {
+                    while (newCursor.currentNode.parent && !newCursor.currentNode.isNamed && next) {
+                        next = gotoParent();
+                    }
+                    const newNode = newCursor.currentNode;
+                    const oldNode = oldCursor.currentNode;
+                    const newEndPosition = ( new Position(newNode.endPosition.row + 1, newNode.endPosition.column + 1));
+                    const oldEndIndex = oldNode.endIndex;
+                    const closestPrev = getClosestPreviousNodes();
+                    const newStartPosition = ( new Position(
+                        closestPrev ? closestPrev.new.endPosition.row + 1 : newNode.startPosition.row + 1,
+                        closestPrev ? closestPrev.new.endPosition.column + 1 : newNode.startPosition.column + 1
+                    ));
+                    const newStartIndex = closestPrev ? closestPrev.new.endIndex : newNode.startIndex;
+                    const oldStartIndex = closestPrev ? closestPrev.old.endIndex : oldNode.startIndex;
+                    changedRanges.push({ newStartPosition, newEndPosition, oldStartIndex, oldEndIndex, newNodeId: newNode.id, newStartIndex, newEndIndex: newNode.endIndex });
+                    next = nextSiblingOrParentSibling();
+                }
+            }
+            else {
+                next = nextSiblingOrParentSibling();
+            }
+        } while (next);
+        if (changedRanges.length === 0 && newTree.rootNode.hasChanges) {
+            return [{ newStartPosition: ( new Position(
+                newTree.rootNode.startPosition.row + 1,
+                newTree.rootNode.startPosition.column + 1
+            )), newEndPosition: ( new Position(
+                newTree.rootNode.endPosition.row + 1,
+                newTree.rootNode.endPosition.column + 1
+            )), oldStartIndex: oldTree.rootNode.startIndex, oldEndIndex: oldTree.rootNode.endIndex, newStartIndex: newTree.rootNode.startIndex, newEndIndex: newTree.rootNode.endIndex, newNodeId: newTree.rootNode.id }];
+        }
+        else {
+            return changedRanges;
+        }
     }
-    _applyEdits(model, changes) {
+    calculateRangeChange(changedNodes) {
+        if (!changedNodes) {
+            return undefined;
+        }
+        const ranges = [];
+        for (let i = 0; i < changedNodes.length; i++) {
+            const node = changedNodes[i];
+            const prevNode = changedNodes[i - 1];
+            if ((i > 0) && prevNode.newEndPosition.equals(node.newStartPosition)) {
+                const prevRangeChange = ranges[ranges.length - 1];
+                prevRangeChange.newRange = ( new Range(
+                    prevRangeChange.newRange.startLineNumber,
+                    prevRangeChange.newRange.startColumn,
+                    node.newEndPosition.lineNumber,
+                    node.newEndPosition.column
+                ));
+                prevRangeChange.oldRangeLength = node.oldEndIndex - prevNode.oldStartIndex;
+                prevRangeChange.newRangeEndOffset = node.newEndIndex;
+            }
+            else {
+                ranges.push({ newRange: Range.fromPositions(node.newStartPosition, node.newEndPosition), oldRangeLength: node.oldEndIndex - node.oldStartIndex, newRangeStartOffset: node.newStartIndex, newRangeEndOffset: node.newEndIndex });
+            }
+        }
+        return ranges;
+    }
+    onDidChangeContent(model, changes) {
+        const version = model.getVersionId();
+        if (version === this._editVersion) {
+            return;
+        }
+        this._applyEdits(changes?.changes ?? [], version);
+        this._onDidChangeContentQueue.queue(async () => {
+            if (this.isDisposed) {
+                return;
+            }
+            let ranges;
+            if (this._lastFullyParsedWithEdits && this._lastFullyParsed) {
+                ranges = this.calculateRangeChange(this.findChangedNodes(this._lastFullyParsedWithEdits, this._lastFullyParsed));
+            }
+            const completed = await this._parseAndUpdateTree(model, version);
+            if (completed) {
+                if (!ranges) {
+                    ranges = [{ newRange: model.getFullModelRange(), oldRangeLength: model.getValueLength(), newRangeStartOffset: 0, newRangeEndOffset: model.getValueLength() }];
+                }
+                this._onDidUpdate.fire({ ranges, versionId: version });
+            }
+        });
+    }
+    _applyEdits(changes, version) {
         for (const change of changes) {
-            const newEndOffset = change.rangeOffset + change.text.length;
-            const newEndPosition = model.getPositionAt(newEndOffset);
-            this.tree?.edit({
+            const originalTextLength = TextLength.ofRange(Range.lift(change.range));
+            const newTextLength = TextLength.ofText(change.text);
+            const summedTextLengths = change.text.length === 0 ? newTextLength : originalTextLength.add(newTextLength);
+            const edit = {
                 startIndex: change.rangeOffset,
                 oldEndIndex: change.rangeOffset + change.rangeLength,
                 newEndIndex: change.rangeOffset + change.text.length,
                 startPosition: { row: change.range.startLineNumber - 1, column: change.range.startColumn - 1 },
                 oldEndPosition: { row: change.range.endLineNumber - 1, column: change.range.endColumn - 1 },
-                newEndPosition: { row: newEndPosition.lineNumber - 1, column: newEndPosition.column - 1 }
-            });
-            this._newEdits = true;
+                newEndPosition: { row: change.range.startLineNumber + summedTextLengths.lineCount - 1, column: summedTextLengths.lineCount ? summedTextLengths.columnCount : (change.range.endColumn + summedTextLengths.columnCount) }
+            };
+            this._tree?.edit(edit);
+            this._lastFullyParsedWithEdits?.edit(edit);
         }
+        this._editVersion = version;
     }
-    async _parseAndUpdateTree(model) {
+    async _parseAndUpdateTree(model, version) {
         const tree = await this._parse(model);
-        if (!this._newEdits) {
-            this.tree = tree;
+        if (tree) {
+            this._tree?.delete();
+            this._tree = tree;
+            this._lastFullyParsed?.delete();
+            this._lastFullyParsed = tree.copy();
+            this._lastFullyParsedWithEdits?.delete();
+            this._lastFullyParsedWithEdits = tree.copy();
+            this._versionId = version;
+            return tree;
+        }
+        else if (!this._tree) {
+            this.parser.reset();
         }
+        return undefined;
     }
     _parse(model) {
         let parseType = TelemetryParseType.Full;
@@ -180,14 +364,14 @@ class TreeSitterParseResult {
     }
     async _parseAndYield(model, parseType) {
         const language = model.getLanguageId();
-        let tree;
         let time = 0;
         let passes = 0;
-        this._newEdits = false;
+        const inProgressVersion = this._editVersion;
+        let newTree;
         do {
             const timer = performance.now();
             try {
-                tree = this.parser.parse((index, position) => this._parseCallback(model, index), this.tree);
+                newTree = this.parser.parse((index, position) => this._parseCallback(model, index), this._tree);
             }
             catch (e) {
             }
@@ -196,12 +380,9 @@ class TreeSitterParseResult {
                 passes++;
             }
             await ( new Promise(resolve => setTimeout0(resolve)));
-            if (model.isDisposed() || this.isDisposed) {
-                return;
-            }
-        } while (!tree && !this._newEdits);
+        } while (!model.isDisposed() && !this.isDisposed && !newTree && inProgressVersion === model.getVersionId());
         this.sendParseTimeTelemetry(parseType, language, time, passes);
-        return tree;
+        return (newTree && (inProgressVersion === model.getVersionId())) ? newTree : undefined;
     }
     _parseCallback(textModel, index) {
         try {
@@ -308,6 +489,7 @@ let TreeSitterTextModelService = class TreeSitterTextModelService extends Dispos
         this._treeSitterImporter = ( new TreeSitterImporter());
         this._onDidUpdateTree = this._register(( new Emitter()));
         this.onDidUpdateTree = this._onDidUpdateTree.event;
+        this.isTest = false;
         this._hasInit = false;
         this._treeSitterLanguages = this._register(( new TreeSitterLanguages(
             this._treeSitterImporter,
@@ -331,8 +513,7 @@ let TreeSitterTextModelService = class TreeSitterTextModelService extends Dispos
         return textModelTreeSitter?.textModelTreeSitter.parseResult;
     }
     async getTree(content, languageId) {
-        await this._init;
-        const language = await this._treeSitterLanguages.getLanguage(languageId);
+        const language = await this.getLanguage(languageId);
         const Parser = await this._treeSitterImporter.getParserClass();
         if (language) {
             const parser = ( new Parser());
@@ -341,11 +522,22 @@ let TreeSitterTextModelService = class TreeSitterTextModelService extends Dispos
         }
         return undefined;
     }
+    async getLanguage(languageId) {
+        await this._init;
+        return this._treeSitterLanguages.getLanguage(languageId);
+    }
     async _doInitParser() {
         const Parser = await this._treeSitterImporter.getParserClass();
+        const isTest = this.isTest;
         await Parser.init({
             locateFile(_file, _folder) {
-                return ( ( FileAccess.asBrowserUri(`${getModuleLocation()}/${FILENAME_TREESITTER_WASM}`)).toString(true));
+                const location = `${getModuleLocation()}/${FILENAME_TREESITTER_WASM}`;
+                if (isTest) {
+                    return ( ( FileAccess.asFileUri(location)).toString(true));
+                }
+                else {
+                    return ( ( FileAccess.asBrowserUri(location)).toString(true));
+                }
             }
         });
         return true;
@@ -400,32 +592,28 @@ let TreeSitterTextModelService = class TreeSitterTextModelService extends Dispos
         }));
         this._modelService.getModels().forEach(model => this._createTextModelTreeSitter(model));
     }
-    getTextModelTreeSitter(model) {
-        return ( new TextModelTreeSitter(
-            model,
-            this._treeSitterLanguages,
-            this._treeSitterImporter,
-            this._logService,
-            this._telemetryService,
-            false
-        ));
+    async getTextModelTreeSitter(model, parseImmediately = false) {
+        await this.getLanguage(model.getLanguageId());
+        return this._createTextModelTreeSitter(model, parseImmediately);
     }
-    _createTextModelTreeSitter(model) {
+    _createTextModelTreeSitter(model, parseImmediately = true) {
         const textModelTreeSitter = ( new TextModelTreeSitter(
             model,
             this._treeSitterLanguages,
             this._treeSitterImporter,
             this._logService,
-            this._telemetryService
+            this._telemetryService,
+            parseImmediately
         ));
         const disposables = ( new DisposableStore());
         disposables.add(textModelTreeSitter);
-        disposables.add(textModelTreeSitter.onDidChangeParseResult((ranges) => this._onDidUpdateTree.fire({ textModel: model, ranges })));
+        disposables.add(textModelTreeSitter.onDidChangeParseResult(change => this._onDidUpdateTree.fire({ textModel: model, ranges: change.ranges ?? [], versionId: change.versionId })));
         this._textModelTreeSitters.set(model, {
             textModelTreeSitter,
             disposables,
             dispose: disposables.dispose.bind(disposables)
         });
+        return textModelTreeSitter;
     }
     _addGrammar(languageId, grammarName) {
         if (!( this._registeredLanguages.has(languageId))) {

package/vscode/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.contribution.js

@@ -2,7 +2,7 @@
 import { __decorate, __param } from '@codingame/monaco-vscode-api/external/tslib/tslib.es6';
 import '@codingame/monaco-vscode-api/vscode/vs/platform/instantiation/common/extensions';
 import { registerWorkbenchContribution2, WorkbenchPhase } from '@codingame/monaco-vscode-api/vscode/vs/workbench/common/contributions';
-import '../../../../editor/browser/services/treeSitter/treeSitterParserService.js';
+import '../../../../editor/common/services/treeSitter/treeSitterParserService.js';
 import { ITreeSitterParserService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService';
 import { ITreeSitterTokenizationFeature } from '@codingame/monaco-vscode-api/vscode/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.service';
 import { CommandsRegistry } from '@codingame/monaco-vscode-api/vscode/vs/platform/commands/common/commands';
@@ -30,7 +30,7 @@ CommandsRegistry.registerCommand('_workbench.colorizeTreeSitterTokens', async (a
     if (!tokenizer) {
         throw ( new Error(`Cannot resolve tokenizer for language ${textModel.getLanguageId()}`));
     }
-    const textModelTreeSitter = treeSitterParserService.getTextModelTreeSitter(textModel);
+    const textModelTreeSitter = await treeSitterParserService.getTextModelTreeSitter(textModel);
     if (!textModelTreeSitter) {
         throw ( new Error(
             `Cannot resolve tree sitter parser for language ${textModel.getLanguageId()}`

package/vscode/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.d.ts

@@ -1,9 +1,21 @@
+import type { Parser } from "@vscode/tree-sitter-wasm";
+import { Event } from "@codingame/monaco-vscode-api/vscode/vs/base/common/event";
 import { Disposable } from "@codingame/monaco-vscode-api/vscode/vs/base/common/lifecycle";
+import { ILanguageIdCodec, ITreeSitterTokenizationSupport, QueryCapture } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/languages";
+import { ITextModel } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/model";
+import { ITreeSitterParserService, ITreeSitterParseResult } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService";
+import { IModelTokensChangedEvent } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/textModelEvents";
 import { IConfigurationService } from "@codingame/monaco-vscode-api/vscode/vs/platform/configuration/common/configuration.service";
 import { IFileService } from "@codingame/monaco-vscode-api/vscode/vs/platform/files/common/files.service";
 import { IInstantiationService } from "@codingame/monaco-vscode-api/vscode/vs/platform/instantiation/common/instantiation";
+import { IThemeService } from "@codingame/monaco-vscode-api/vscode/vs/platform/theme/common/themeService.service";
 import { ILanguageService } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/languages/language";
+import { ITreeSitterTokenizationStoreService } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/model/treeSitterTokenStoreService";
+import { TokenUpdate } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/model/tokenStore";
+import { Range } from "@codingame/monaco-vscode-api/vscode/vs/editor/common/core/range";
+import { ICodeEditorService } from "@codingame/monaco-vscode-api/vscode/vs/editor/browser/services/codeEditorService";
 import { ITreeSitterTokenizationFeature } from "@codingame/monaco-vscode-api/vscode/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.service";
+type TreeSitterQueries = string;
 export declare class TreeSitterTokenizationFeature extends Disposable implements ITreeSitterTokenizationFeature {
     private readonly _languageService;
     private readonly _configurationService;
@@ -17,3 +29,52 @@ export declare class TreeSitterTokenizationFeature extends Disposable implements
     private _fetchQueries;
     private _createTokenizationSupport;
 }
+export declare class TreeSitterTokenizationSupport extends Disposable implements ITreeSitterTokenizationSupport {
+    private readonly _queries;
+    private readonly _languageId;
+    private readonly _languageIdCodec;
+    private readonly _treeSitterService;
+    private readonly _themeService;
+    private readonly _tokenizationStoreService;
+    private readonly _codeEditorService;
+    private _query;
+    private readonly _onDidChangeTokens;
+    readonly onDidChangeTokens: Event<{
+        textModel: ITextModel;
+        changes: IModelTokensChangedEvent;
+    }>;
+    private _colorThemeData;
+    private _languageAddedListener;
+    constructor(_queries: TreeSitterQueries, _languageId: string, _languageIdCodec: ILanguageIdCodec, _treeSitterService: ITreeSitterParserService, _themeService: IThemeService, _tokenizationStoreService: ITreeSitterTokenizationStoreService, _codeEditorService: ICodeEditorService);
+    private _createEmptyTokens;
+    private _firstTreeUpdate;
+    private _setViewPortTokens;
+    private _handleTreeUpdate;
+    private _updateTreeForRanges;
+    private _refreshNeedsRefresh;
+    private _rangeTokensAsUpdates;
+    getTokensInRange(textModel: ITextModel, range: Range, rangeStartOffset: number, rangeEndOffset: number, captures?: {
+        tree: ITreeSitterParseResult | undefined;
+        captures: QueryCapture[];
+    }): TokenUpdate[] | undefined;
+    private _getTree;
+    private _ensureQuery;
+    private reset;
+    captureAtPosition(lineNumber: number, column: number, textModel: ITextModel): QueryCapture[];
+    captureAtPositionTree(lineNumber: number, column: number, tree: Parser.Tree): QueryCapture[];
+    private _captureAtRange;
+    tokenizeEncoded(lineNumber: number, textModel: ITextModel): Uint32Array | undefined;
+    tokenizeEncodedInstrumented(lineNumber: number, textModel: ITextModel): {
+        result: Uint32Array;
+        captureTime: number;
+        metadataTime: number;
+    } | undefined;
+    private _getTreeAndCaptures;
+    private _tokenize;
+    private _createTokensFromCaptures;
+    private _tokenizeCapturesWithMetadata;
+    private _emptyToken;
+    private _tokenizeEncoded;
+    dispose(): void;
+}
+export {};

package/vscode/src/vs/workbench/services/treeSitter/browser/treeSitterTokenizationFeature.js

@@ -5,7 +5,6 @@ import { Disposable, DisposableMap, DisposableStore } from '@codingame/monaco-vs
 import { FileAccess } from '@codingame/monaco-vscode-api/vscode/vs/base/common/network';
 import { LazyTokenizationSupport, TreeSitterTokenizationRegistry } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/languages';
 import { EDITOR_EXPERIMENTAL_PREFER_TREESITTER, ITreeSitterParserService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/services/treeSitterParserService';
-import { ColumnRange } from '@codingame/monaco-vscode-api/vscode/vs/editor/contrib/inlineCompletions/browser/utils';
 import { IConfigurationService } from '@codingame/monaco-vscode-api/vscode/vs/platform/configuration/common/configuration.service';
 import { IFileService } from '@codingame/monaco-vscode-api/vscode/vs/platform/files/common/files.service';
 import '@codingame/monaco-vscode-api/vscode/vs/platform/instantiation/common/extensions';
@@ -14,6 +13,10 @@ import { IThemeService } from '@codingame/monaco-vscode-api/vscode/vs/platform/t
 import { findMetadata } from '@codingame/monaco-vscode-9d0168a3-519b-57f3-9bcc-89efc41f951a-common/vscode/vs/workbench/services/themes/common/colorThemeData';
 import { ILanguageService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/languages/language';
 import { StopWatch } from '@codingame/monaco-vscode-api/vscode/vs/base/common/stopwatch';
+import { ITreeSitterTokenizationStoreService } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/model/treeSitterTokenStoreService';
+import { Range } from '@codingame/monaco-vscode-api/vscode/vs/editor/common/core/range';
+import { ICodeEditorService } from '@codingame/monaco-vscode-api/vscode/vs/editor/browser/services/codeEditorService';
+import { setTimeout0 } from '@codingame/monaco-vscode-api/vscode/vs/base/common/platform';
 
 const ALLOWED_SUPPORT = ['typescript'];
 let TreeSitterTokenizationFeature = class TreeSitterTokenizationFeature extends Disposable {
@@ -23,7 +26,7 @@ let TreeSitterTokenizationFeature = class TreeSitterTokenizationFeature extends
         this._configurationService = _configurationService;
         this._instantiationService = _instantiationService;
         this._fileService = _fileService;
-        this._tokenizersRegistrations = ( new DisposableMap());
+        this._tokenizersRegistrations = this._register(( new DisposableMap()));
         this._handleGrammarsExtPoint();
         this._register(this._configurationService.onDidChangeConfiguration(e => {
             if (e.affectsConfiguration(EDITOR_EXPERIMENTAL_PREFER_TREESITTER)) {
@@ -68,28 +71,183 @@ TreeSitterTokenizationFeature = ( __decorate([
     ( __param(3, IFileService))
 ], TreeSitterTokenizationFeature));
 let TreeSitterTokenizationSupport = class TreeSitterTokenizationSupport extends Disposable {
-    constructor(_queries, _languageId, _languageIdCodec, _treeSitterService, _themeService) {
+    constructor(_queries, _languageId, _languageIdCodec, _treeSitterService, _themeService, _tokenizationStoreService, _codeEditorService) {
         super();
         this._queries = _queries;
         this._languageId = _languageId;
         this._languageIdCodec = _languageIdCodec;
         this._treeSitterService = _treeSitterService;
         this._themeService = _themeService;
+        this._tokenizationStoreService = _tokenizationStoreService;
+        this._codeEditorService = _codeEditorService;
         this._onDidChangeTokens = ( new Emitter());
         this.onDidChangeTokens = this._onDidChangeTokens.event;
         this._register(Event.runAndSubscribe(this._themeService.onDidColorThemeChange, () => this.reset()));
         this._register(this._treeSitterService.onDidUpdateTree((e) => {
-            const maxLine = e.textModel.getLineCount();
+            if (this._tokenizationStoreService.hasTokens(e.textModel)) {
+                for (const range of e.ranges) {
+                    this._tokenizationStoreService.markForRefresh(e.textModel, range.newRange);
+                }
+            }
+            if (e.versionId !== e.textModel.getVersionId()) {
+                return;
+            }
+            if (!this._tokenizationStoreService.hasTokens(e.textModel)) {
+                this._firstTreeUpdate(e.textModel, e.versionId);
+            }
+            else {
+                this._handleTreeUpdate(e);
+            }
+        }));
+    }
+    _createEmptyTokens(textModel) {
+        const languageId = this._languageIdCodec.encodeLanguageId(this._languageId);
+        const emptyToken = this._emptyToken(languageId);
+        const modelEndOffset = textModel.getValueLength();
+        const emptyTokens = [{ token: emptyToken, length: modelEndOffset, startOffsetInclusive: 0 }];
+        return emptyTokens;
+    }
+    _firstTreeUpdate(textModel, versionId) {
+        const tokens = this._createEmptyTokens(textModel);
+        this._tokenizationStoreService.setTokens(textModel, tokens);
+        this._setViewPortTokens(textModel, versionId);
+    }
+    _setViewPortTokens(textModel, versionId) {
+        const maxLine = textModel.getLineCount();
+        let rangeChanges;
+        const editor = this._codeEditorService.listCodeEditors().find(editor => editor.getModel() === textModel);
+        if (editor) {
+            const viewPort = editor.getVisibleRangesPlusViewportAboveBelow();
+            const ranges = ( new Array(viewPort.length));
+            rangeChanges = ( new Array(viewPort.length));
+            for (let i = 0; i < viewPort.length; i++) {
+                const range = viewPort[i];
+                ranges[i] = { fromLineNumber: range.startLineNumber, toLineNumber: range.endLineNumber < maxLine ? range.endLineNumber : maxLine };
+                const newRangeStartOffset = textModel.getOffsetAt(range.getStartPosition());
+                const newRangeEndOffset = textModel.getOffsetAt(range.getEndPosition());
+                rangeChanges[i] = {
+                    newRange: range,
+                    newRangeStartOffset,
+                    newRangeEndOffset,
+                    oldRangeLength: newRangeEndOffset - newRangeStartOffset
+                };
+            }
+        }
+        else {
+            const valueLength = textModel.getValueLength();
+            rangeChanges = [{ newRange: ( new Range(1, 1, maxLine, textModel.getLineMaxColumn(maxLine))), newRangeStartOffset: 0, newRangeEndOffset: valueLength, oldRangeLength: valueLength }];
+        }
+        this._handleTreeUpdate({ ranges: rangeChanges, textModel, versionId });
+    }
+    _handleTreeUpdate(e) {
+        let rangeChanges = [];
+        const chunkSize = 10000;
+        for (let i = 0; i < e.ranges.length; i++) {
+            const rangeLength = e.ranges[i].newRangeEndOffset - e.ranges[i].newRangeStartOffset;
+            if (e.ranges[i].oldRangeLength === rangeLength) {
+                if (rangeLength > chunkSize) {
+                    const fullRangeEndOffset = e.ranges[i].newRangeEndOffset;
+                    let chunkStart = e.ranges[i].newRangeStartOffset;
+                    let chunkEnd = chunkStart + chunkSize;
+                    let chunkStartingPosition = e.ranges[i].newRange.getStartPosition();
+                    do {
+                        const chunkEndPosition = e.textModel.getPositionAt(chunkEnd);
+                        const chunkRange = Range.fromPositions(chunkStartingPosition, chunkEndPosition);
+                        rangeChanges.push({
+                            newRange: chunkRange,
+                            newRangeStartOffset: chunkStart,
+                            newRangeEndOffset: chunkEnd,
+                            oldRangeLength: chunkEnd - chunkStart
+                        });
+                        chunkStart = chunkEnd;
+                        if (chunkEnd < fullRangeEndOffset && chunkEnd + chunkSize > fullRangeEndOffset) {
+                            chunkEnd = fullRangeEndOffset;
+                        }
+                        else {
+                            chunkEnd = chunkEnd + chunkSize;
+                        }
+                        chunkStartingPosition = chunkEndPosition;
+                    } while (chunkEnd <= fullRangeEndOffset);
+                }
+                else {
+                    rangeChanges.push(e.ranges[i]);
+                }
+            }
+            else {
+                rangeChanges = e.ranges;
+                break;
+            }
+        }
+        const captures = ( rangeChanges.map(range => this._getTreeAndCaptures(range.newRange, e.textModel)));
+        this._updateTreeForRanges(e.textModel, rangeChanges, e.versionId, captures).then(() => {
+            const tree = this._getTree(e.textModel);
+            if (!e.textModel.isDisposed() && (tree?.versionId === e.textModel.getVersionId())) {
+                this._refreshNeedsRefresh(e.textModel);
+            }
+        });
+    }
+    async _updateTreeForRanges(textModel, rangeChanges, versionId, captures) {
+        let tokenUpdate;
+        for (let i = 0; i < rangeChanges.length; i++) {
+            if (versionId !== textModel.getVersionId()) {
+                break;
+            }
+            const capture = captures[i];
+            const range = rangeChanges[i];
+            const updates = this.getTokensInRange(textModel, range.newRange, range.newRangeStartOffset, range.newRangeEndOffset, capture);
+            if (updates) {
+                tokenUpdate = { oldRangeLength: range.oldRangeLength, newTokens: updates };
+            }
+            else {
+                tokenUpdate = { oldRangeLength: range.oldRangeLength, newTokens: [] };
+            }
+            this._tokenizationStoreService.updateTokens(textModel, versionId, [tokenUpdate]);
             this._onDidChangeTokens.fire({
-                textModel: e.textModel,
+                textModel: textModel,
                 changes: {
                     semanticTokensApplied: false,
-                    ranges: ( e.ranges.map(
-                        range => ({ fromLineNumber: range.startLineNumber, toLineNumber: range.endLineNumber < maxLine ? range.endLineNumber : maxLine })
-                    )),
+                    ranges: [{ fromLineNumber: range.newRange.getStartPosition().lineNumber, toLineNumber: range.newRange.getEndPosition().lineNumber }]
                 }
             });
-        }));
+            await ( new Promise(resolve => setTimeout0(resolve)));
+        }
+    }
+    _refreshNeedsRefresh(textModel) {
+        const rangesToRefresh = this._tokenizationStoreService.getNeedsRefresh(textModel);
+        if (rangesToRefresh.length === 0) {
+            return;
+        }
+        const rangeChanges = ( new Array(rangesToRefresh.length));
+        for (let i = 0; i < rangesToRefresh.length; i++) {
+            const range = rangesToRefresh[i];
+            rangeChanges[i] = {
+                newRange: range.range,
+                newRangeStartOffset: range.startOffset,
+                newRangeEndOffset: range.endOffset,
+                oldRangeLength: range.endOffset - range.startOffset
+            };
+        }
+        this._handleTreeUpdate({ ranges: rangeChanges, textModel, versionId: textModel.getVersionId() });
+    }
+    _rangeTokensAsUpdates(rangeOffset, endOffsetToken) {
+        const updates = [];
+        let lastEnd = 0;
+        for (const token of endOffsetToken) {
+            if (token.endOffset <= lastEnd) {
+                continue;
+            }
+            updates.push({ startOffsetInclusive: rangeOffset + lastEnd, length: token.endOffset - lastEnd, token: token.metadata });
+            lastEnd = token.endOffset;
+        }
+        return updates;
+    }
+    getTokensInRange(textModel, range, rangeStartOffset, rangeEndOffset, captures) {
+        const languageId = this._languageIdCodec.encodeLanguageId(this._languageId);
+        const tokens = captures ? this._tokenizeCapturesWithMetadata(captures.tree, captures.captures, languageId, rangeStartOffset, rangeEndOffset) : this._tokenize(languageId, range, rangeStartOffset, rangeEndOffset, textModel);
+        if (tokens?.endOffsetsAndMetadata) {
+            return this._rangeTokensAsUpdates(rangeStartOffset, tokens.endOffsetsAndMetadata);
+        }
+        return undefined;
     }
     _getTree(textModel) {
         return this._treeSitterService.getParseResult(textModel);
@@ -114,19 +272,26 @@ let TreeSitterTokenizationSupport = class TreeSitterTokenizationSupport extends
     }
     captureAtPosition(lineNumber, column, textModel) {
         const tree = this._getTree(textModel);
-        const captures = this._captureAtRange(lineNumber, ( new ColumnRange(column, column)), tree?.tree);
+        const captures = this._captureAtRange(( new Range(lineNumber, column, lineNumber, column + 1)), tree?.tree);
         return captures;
     }
     captureAtPositionTree(lineNumber, column, tree) {
-        const captures = this._captureAtRange(lineNumber, ( new ColumnRange(column, column)), tree);
+        const captures = this._captureAtRange(( new Range(lineNumber, column, lineNumber, column + 1)), tree);
         return captures;
     }
-    _captureAtRange(lineNumber, columnRange, tree) {
+    _captureAtRange(range, tree) {
         const query = this._ensureQuery();
         if (!tree || !query) {
             return [];
         }
-        return query.captures(tree.rootNode, { startPosition: { row: lineNumber - 1, column: columnRange.startColumn - 1 }, endPosition: { row: lineNumber - 1, column: columnRange.endColumnExclusive } });
+        return ( query.captures(tree.rootNode, { startPosition: { row: range.startLineNumber - 1, column: range.startColumn - 1 }, endPosition: { row: range.endLineNumber - 1, column: range.endColumn - 1 } }).map(capture => ({
+            name: capture.name,
+            text: capture.node.text,
+            node: {
+                startIndex: capture.node.startIndex,
+                endIndex: capture.node.endIndex
+            }
+        })));
     }
     tokenizeEncoded(lineNumber, textModel) {
         return this._tokenizeEncoded(lineNumber, textModel)?.result;
@@ -134,34 +299,44 @@ let TreeSitterTokenizationSupport = class TreeSitterTokenizationSupport extends
     tokenizeEncodedInstrumented(lineNumber, textModel) {
         return this._tokenizeEncoded(lineNumber, textModel);
     }
-    _tokenizeEncoded(lineNumber, textModel) {
-        const stopwatch = StopWatch.create();
-        const lineLength = textModel.getLineMaxColumn(lineNumber);
+    _getTreeAndCaptures(range, textModel) {
         const tree = this._getTree(textModel);
-        const captures = this._captureAtRange(lineNumber, ( new ColumnRange(1, lineLength)), tree?.tree);
+        const captures = this._captureAtRange(range, tree?.tree);
+        return { tree, captures };
+    }
+    _tokenize(encodedLanguageId, range, rangeStartOffset, rangeEndOffset, textModel) {
+        const { tree, captures } = this._getTreeAndCaptures(range, textModel);
+        return this._tokenizeCapturesWithMetadata(tree, captures, encodedLanguageId, rangeStartOffset, rangeEndOffset);
+    }
+    _createTokensFromCaptures(tree, captures, rangeStartOffset, rangeEndOffset) {
+        const stopwatch = StopWatch.create();
+        const rangeLength = rangeEndOffset - rangeStartOffset;
         if (captures.length === 0) {
+            if (tree) {
+                stopwatch.stop();
+                const endOffsetsAndMetadata = [{ endOffset: rangeLength, scopes: [] }];
+                return { endOffsets: endOffsetsAndMetadata, captureTime: stopwatch.elapsed() };
+            }
             return undefined;
         }
         const endOffsetsAndScopes = Array(captures.length);
         endOffsetsAndScopes.fill({ endOffset: 0, scopes: [] });
         let tokenIndex = 0;
-        const lineStartOffset = textModel.getOffsetAt({ lineNumber: lineNumber, column: 1 });
         const increaseSizeOfTokensByOneToken = () => {
             endOffsetsAndScopes.push({ endOffset: 0, scopes: [] });
         };
-        const encodedLanguageId = this._languageIdCodec.encodeLanguageId(this._languageId);
         for (let captureIndex = 0; captureIndex < captures.length; captureIndex++) {
             const capture = captures[captureIndex];
-            const tokenEndIndex = capture.node.endIndex < lineStartOffset + lineLength ? capture.node.endIndex : lineStartOffset + lineLength;
-            const tokenStartIndex = capture.node.startIndex < lineStartOffset ? lineStartOffset : capture.node.startIndex;
-            const lineRelativeOffset = tokenEndIndex - lineStartOffset;
+            const tokenEndIndex = capture.node.endIndex < rangeEndOffset ? ((capture.node.endIndex < rangeStartOffset) ? rangeStartOffset : capture.node.endIndex) : rangeEndOffset;
+            const tokenStartIndex = capture.node.startIndex < rangeStartOffset ? rangeStartOffset : ((capture.node.startIndex > tokenEndIndex) ? tokenEndIndex : capture.node.startIndex);
+            const lineRelativeOffset = tokenEndIndex - rangeStartOffset;
             let previousTokenEnd;
             const currentTokenLength = tokenEndIndex - tokenStartIndex;
             if (captureIndex > 0) {
                 previousTokenEnd = endOffsetsAndScopes[(tokenIndex - 1)].endOffset;
             }
             else {
-                previousTokenEnd = tokenStartIndex - lineStartOffset - 1;
+                previousTokenEnd = tokenStartIndex - rangeStartOffset - 1;
             }
             const intermediateTokenOffset = lineRelativeOffset - currentTokenLength;
             if ((previousTokenEnd >= 0) && (previousTokenEnd < intermediateTokenOffset)) {
@@ -174,42 +349,89 @@ let TreeSitterTokenizationSupport = class TreeSitterTokenizationSupport extends
                 tokenIndex++;
             };
             if (previousTokenEnd >= lineRelativeOffset) {
-                const previousTokenStartOffset = ((tokenIndex >= 2) ? endOffsetsAndScopes[tokenIndex - 2].endOffset : 0);
                 const originalPreviousTokenEndOffset = endOffsetsAndScopes[tokenIndex - 1].endOffset;
-                if ((previousTokenStartOffset + currentTokenLength) === originalPreviousTokenEndOffset) {
+                const previousTokenStartOffset = ((tokenIndex >= 2) ? endOffsetsAndScopes[tokenIndex - 2].endOffset : 0);
+                const loopOriginalPreviousTokenEndOffset = endOffsetsAndScopes[tokenIndex - 1].endOffset;
+                const previousPreviousTokenEndOffset = (tokenIndex >= 2) ? endOffsetsAndScopes[tokenIndex - 2].endOffset : 0;
+                if ((previousTokenStartOffset + currentTokenLength) === loopOriginalPreviousTokenEndOffset) {
                     endOffsetsAndScopes[tokenIndex - 1].scopes[endOffsetsAndScopes[tokenIndex - 1].scopes.length - 1] = capture.name;
                 }
-                else {
-                    endOffsetsAndScopes[tokenIndex - 1].endOffset = intermediateTokenOffset;
-                    addCurrentTokenToArray();
-                    increaseSizeOfTokensByOneToken();
-                    endOffsetsAndScopes[tokenIndex].endOffset = originalPreviousTokenEndOffset;
-                    endOffsetsAndScopes[tokenIndex].scopes = endOffsetsAndScopes[tokenIndex - 2].scopes;
-                    tokenIndex++;
+                else if (previousPreviousTokenEndOffset <= intermediateTokenOffset) {
+                    let originalPreviousTokenScopes;
+                    if (previousPreviousTokenEndOffset !== intermediateTokenOffset) {
+                        endOffsetsAndScopes[tokenIndex - 1] = { endOffset: intermediateTokenOffset, scopes: endOffsetsAndScopes[tokenIndex - 1].scopes };
+                        addCurrentTokenToArray();
+                        originalPreviousTokenScopes = [...endOffsetsAndScopes[tokenIndex - 2].scopes];
+                    }
+                    else {
+                        originalPreviousTokenScopes = [...endOffsetsAndScopes[tokenIndex - 1].scopes];
+                        endOffsetsAndScopes[tokenIndex - 1] = { endOffset: lineRelativeOffset, scopes: [capture.name] };
+                    }
+                    if (originalPreviousTokenEndOffset !== lineRelativeOffset) {
+                        increaseSizeOfTokensByOneToken();
+                        endOffsetsAndScopes[tokenIndex] = { endOffset: originalPreviousTokenEndOffset, scopes: originalPreviousTokenScopes };
+                        tokenIndex++;
+                    }
+                    else {
+                        endOffsetsAndScopes[tokenIndex - 1].scopes.unshift(...originalPreviousTokenScopes);
+                    }
                 }
             }
             else {
                 addCurrentTokenToArray();
             }
         }
-        if (captures[captures.length - 1].node.endPosition.column + 1 < lineLength) {
-            increaseSizeOfTokensByOneToken();
-            endOffsetsAndScopes[tokenIndex].endOffset = lineLength - 1;
-            tokenIndex++;
+        if ((endOffsetsAndScopes[tokenIndex - 1].endOffset < rangeLength)) {
+            if (rangeLength - endOffsetsAndScopes[tokenIndex - 1].endOffset > 0) {
+                increaseSizeOfTokensByOneToken();
+                endOffsetsAndScopes[tokenIndex] = { endOffset: rangeLength, scopes: endOffsetsAndScopes[tokenIndex].scopes };
+                tokenIndex++;
+            }
         }
-        const captureTime = stopwatch.elapsed();
-        stopwatch.reset();
-        const tokens = ( new Uint32Array((tokenIndex) * 2));
-        for (let i = 0; i < tokenIndex; i++) {
+        for (let i = 0; i < endOffsetsAndScopes.length; i++) {
             const token = endOffsetsAndScopes[i];
-            if (token.endOffset === 0 && token.scopes.length === 0) {
+            if (token.endOffset === 0 && token.scopes.length === 0 && i !== 0) {
+                endOffsetsAndScopes.splice(i, endOffsetsAndScopes.length - i);
                 break;
             }
-            tokens[i * 2] = token.endOffset;
-            tokens[i * 2 + 1] = findMetadata(this._colorThemeData, token.scopes, encodedLanguageId);
+        }
+        const captureTime = stopwatch.elapsed();
+        return { endOffsets: endOffsetsAndScopes, captureTime };
+    }
+    _tokenizeCapturesWithMetadata(tree, captures, encodedLanguageId, rangeStartOffset, rangeEndOffset) {
+        const stopwatch = StopWatch.create();
+        const emptyTokens = this._createTokensFromCaptures(tree, captures, rangeStartOffset, rangeEndOffset);
+        if (!emptyTokens) {
+            return undefined;
+        }
+        const endOffsetsAndScopes = emptyTokens.endOffsets;
+        for (let i = 0; i < endOffsetsAndScopes.length; i++) {
+            const token = endOffsetsAndScopes[i];
+            token.metadata = findMetadata(this._colorThemeData, token.scopes, encodedLanguageId);
         }
         const metadataTime = stopwatch.elapsed();
-        return { result: tokens, captureTime, metadataTime };
+        return { endOffsetsAndMetadata: endOffsetsAndScopes, captureTime: emptyTokens.captureTime, metadataTime };
+    }
+    _emptyToken(encodedLanguageId) {
+        return findMetadata(this._colorThemeData, [], encodedLanguageId);
+    }
+    _tokenizeEncoded(lineNumber, textModel) {
+        const encodedLanguageId = this._languageIdCodec.encodeLanguageId(this._languageId);
+        const lineOffset = textModel.getOffsetAt({ lineNumber: lineNumber, column: 1 });
+        const maxLine = textModel.getLineCount();
+        const lineEndOffset = (lineNumber + 1 <= maxLine) ? textModel.getOffsetAt({ lineNumber: lineNumber + 1, column: 1 }) : textModel.getValueLength();
+        const lineLength = lineEndOffset - lineOffset;
+        const result = this._tokenize(encodedLanguageId, ( new Range(lineNumber, 1, lineNumber, lineLength)), lineOffset, lineEndOffset, textModel);
+        if (!result) {
+            return undefined;
+        }
+        const tokens = ( new Uint32Array((result.endOffsetsAndMetadata.length) * 2));
+        for (let i = 0; i < result.endOffsetsAndMetadata.length; i++) {
+            const token = result.endOffsetsAndMetadata[i];
+            tokens[i * 2] = token.endOffset;
+            tokens[i * 2 + 1] = token.metadata;
+        }
+        return { result: tokens, captureTime: result.captureTime, metadataTime: result.metadataTime };
     }
     dispose() {
         super.dispose();
@@ -219,7 +441,9 @@ let TreeSitterTokenizationSupport = class TreeSitterTokenizationSupport extends
 };
 TreeSitterTokenizationSupport = ( __decorate([
     ( __param(3, ITreeSitterParserService)),
-    ( __param(4, IThemeService))
+    ( __param(4, IThemeService)),
+    ( __param(5, ITreeSitterTokenizationStoreService)),
+    ( __param(6, ICodeEditorService))
 ], TreeSitterTokenizationSupport));
 
-export { TreeSitterTokenizationFeature };
+export { TreeSitterTokenizationFeature, TreeSitterTokenizationSupport };