Keep previous partial tokens

src/vs/editor/common/model.ts

@@ -832,7 +832,7 @@ export interface ITextModel {
 	/**
 	 * @internal
 	 */
-	setPartialSemanticTokens(tokens: MultilineTokens2[] | null): void;
+	setPartialSemanticTokens(range: Range, tokens: MultilineTokens2[] | null): void;
 
 	/**
 	 * @internal

src/vs/editor/common/model/textModel.ts

@@ -1807,11 +1807,17 @@ export class TextModel extends Disposable implements model.ITextModel {
 		return this._tokens2.isComplete();
 	}
 
-	public setPartialSemanticTokens(tokens: MultilineTokens2[]): void {
+	public setPartialSemanticTokens(range: Range, tokens: MultilineTokens2[]): void {
 		if (this.hasSemanticTokens()) {
 			return;
 		}
-		this.setSemanticTokens(tokens, false);
+		const changedRange = this._tokens2.setPartial(range, tokens);
+
+		this._emitModelTokensChangedEvent({
+			tokenizationSupportChanged: false,
+			semanticTokensApplied: true,
+			ranges: [{ fromLineNumber: changedRange.startLineNumber, toLineNumber: changedRange.endLineNumber }]
+		});
 	}
 
 	public tokenizeViewport(startLineNumber: number, endLineNumber: number): void {

src/vs/editor/common/model/tokensStore.ts

@@ -6,7 +6,7 @@
 import * as arrays from 'vs/base/common/arrays';
 import { LineTokens } from 'vs/editor/common/core/lineTokens';
 import { Position } from 'vs/editor/common/core/position';
-import { IRange } from 'vs/editor/common/core/range';
+import { IRange, Range } from 'vs/editor/common/core/range';
 import { ColorId, FontStyle, LanguageId, MetadataConsts, StandardTokenType, TokenMetadata } from 'vs/editor/common/modes';
 import { writeUInt32BE, readUInt32BE } from 'vs/base/common/buffer';
 import { CharCode } from 'vs/base/common/charCode';
@@ -124,16 +124,7 @@ export class MultilineTokensBuilder {
 	}
 }
 
-export interface IEncodedTokens {
-	getMaxDeltaLine(): number;
-	getLineTokens(deltaLine: number): LineTokens2 | null;
-
-	clear(): void;
-	acceptDeleteRange(horizontalShiftForFirstLineTokens: number, startDeltaLine: number, startCharacter: number, endDeltaLine: number, endCharacter: number): void;
-	acceptInsertText(deltaLine: number, character: number, eolCount: number, firstLineLength: number, lastLineLength: number, firstCharCode: number): void;
-}
-
-export class SparseEncodedTokens implements IEncodedTokens {
+export class SparseEncodedTokens {
 	/**
 	 * The encoding of tokens is:
 	 *  4*i    deltaLine (from `startLineNumber`)
@@ -157,6 +148,17 @@ export class SparseEncodedTokens implements IEncodedTokens {
 		return this._getDeltaLine(tokenCount - 1);
 	}
 
+	public getRange(): Range | null {
+		const tokenCount = this._getTokenCount();
+		if (tokenCount === 0) {
+			return null;
+		}
+		const startChar = this._getStartCharacter(0);
+		const maxDeltaLine = this._getDeltaLine(tokenCount - 1);
+		const endChar = this._getEndCharacter(tokenCount - 1);
+		return new Range(0, startChar + 1, maxDeltaLine, endChar + 1);
+	}
+
 	private _getTokenCount(): number {
 		return this._tokenCount;
 	}
@@ -165,6 +167,18 @@ export class SparseEncodedTokens implements IEncodedTokens {
 		return this._tokens[4 * tokenIndex];
 	}
 
+	private _getStartCharacter(tokenIndex: number): number {
+		return this._tokens[4 * tokenIndex + 1];
+	}
+
+	private _getEndCharacter(tokenIndex: number): number {
+		return this._tokens[4 * tokenIndex + 2];
+	}
+
+	public isEmpty(): boolean {
+		return (this._getTokenCount() === 0);
+	}
+
 	public getLineTokens(deltaLine: number): LineTokens2 | null {
 		let low = 0;
 		let high = this._getTokenCount() - 1;
@@ -201,6 +215,45 @@ export class SparseEncodedTokens implements IEncodedTokens {
 		this._tokenCount = 0;
 	}
 
+	public removeTokens(startDeltaLine: number, startChar: number, endDeltaLine: number, endChar: number): number {
+		const tokens = this._tokens;
+		const tokenCount = this._tokenCount;
+		let newTokenCount = 0;
+		let hasDeletedTokens = false;
+		let firstDeltaLine = 0;
+		for (let i = 0; i < tokenCount; i++) {
+			const srcOffset = 4 * i;
+			const tokenDeltaLine = tokens[srcOffset];
+			const tokenStartCharacter = tokens[srcOffset + 1];
+			const tokenEndCharacter = tokens[srcOffset + 2];
+			const tokenMetadata = tokens[srcOffset + 3];
+
+			if (
+				(tokenDeltaLine > startDeltaLine || (tokenDeltaLine === startDeltaLine && tokenStartCharacter >= startChar))
+				&& (tokenDeltaLine < endDeltaLine || (tokenDeltaLine === endDeltaLine && tokenEndCharacter <= endChar))
+			) {
+				hasDeletedTokens = true;
+			} else {
+				if (newTokenCount === 0) {
+					firstDeltaLine = tokenDeltaLine;
+				}
+				if (hasDeletedTokens) {
+					// must move the token to the left
+					const destOffset = 4 * newTokenCount;
+					tokens[destOffset] = tokenDeltaLine - firstDeltaLine;
+					tokens[destOffset + 1] = tokenStartCharacter;
+					tokens[destOffset + 2] = tokenEndCharacter;
+					tokens[destOffset + 3] = tokenMetadata;
+				}
+				newTokenCount++;
+			}
+		}
+
+		this._tokenCount = newTokenCount;
+
+		return firstDeltaLine;
+	}
+
 	public acceptDeleteRange(horizontalShiftForFirstLineTokens: number, startDeltaLine: number, startCharacter: number, endDeltaLine: number, endCharacter: number): void {
 		// This is a bit complex, here are the cases I used to think about this:
 		//
@@ -457,9 +510,9 @@ export class MultilineTokens2 {
 
 	public startLineNumber: number;
 	public endLineNumber: number;
-	public tokens: IEncodedTokens;
+	public tokens: SparseEncodedTokens;
 
-	constructor(startLineNumber: number, tokens: IEncodedTokens) {
+	constructor(startLineNumber: number, tokens: SparseEncodedTokens) {
 		this.startLineNumber = startLineNumber;
 		this.tokens = tokens;
 		this.endLineNumber = this.startLineNumber + this.tokens.getMaxDeltaLine();
@@ -469,6 +522,10 @@ export class MultilineTokens2 {
 		this.endLineNumber = this.startLineNumber + this.tokens.getMaxDeltaLine();
 	}
 
+	public isEmpty(): boolean {
+		return this.tokens.isEmpty();
+	}
+
 	public getLineTokens(lineNumber: number): LineTokens2 | null {
 		if (this.startLineNumber <= lineNumber && lineNumber <= this.endLineNumber) {
 			return this.tokens.getLineTokens(lineNumber - this.startLineNumber);
@@ -476,6 +533,22 @@ export class MultilineTokens2 {
 		return null;
 	}
 
+	public getRange(): Range | null {
+		const deltaRange = this.tokens.getRange();
+		if (!deltaRange) {
+			return deltaRange;
+		}
+		return new Range(this.startLineNumber + deltaRange.startLineNumber, deltaRange.startColumn, this.startLineNumber + deltaRange.endLineNumber, deltaRange.endColumn);
+	}
+
+	public removeTokens(range: Range): void {
+		const startLineIndex = range.startLineNumber - this.startLineNumber;
+		const endLineIndex = range.endLineNumber - this.startLineNumber;
+
+		this.startLineNumber += this.tokens.removeTokens(startLineIndex, range.startColumn - 1, endLineIndex, range.endColumn - 1);
+		this._updateEndLineNumber();
+	}
+
 	public applyEdit(range: IRange, text: string): void {
 		const [eolCount, firstLineLength, lastLineLength] = countEOL(text);
 		this.acceptEdit(range, eolCount, firstLineLength, lastLineLength, text.length > 0 ? text.charCodeAt(0) : CharCode.Null);
@@ -749,11 +822,50 @@ export class TokensStore2 {
 		this._isComplete = false;
 	}
 
-	public set(pieces: MultilineTokens2[] | null, isComplete: boolean) {
+	public set(pieces: MultilineTokens2[] | null, isComplete: boolean): void {
 		this._pieces = pieces || [];
 		this._isComplete = isComplete;
 	}
 
+	public setPartial(_range: Range, pieces: MultilineTokens2[]): Range {
+		if (pieces.length === 0) {
+			return _range;
+		}
+		const _firstRange = pieces[0].getRange();
+		const _lastRange = pieces[pieces.length - 1].getRange();
+		if (!_firstRange || !_lastRange) {
+			return _range;
+		}
+		const range = _range.plusRange(_firstRange).plusRange(_lastRange);
+		let insertIndex = this._pieces.length;
+		for (let i = 0, len = this._pieces.length; i < len; i++) {
+			const piece = this._pieces[i];
+			if (piece.endLineNumber < range.startLineNumber) {
+				continue;
+			}
+			if (piece.startLineNumber > range.endLineNumber) {
+				insertIndex = Math.min(i, insertIndex);
+				break;
+			}
+			piece.removeTokens(range);
+
+			if (piece.isEmpty()) {
+				this._pieces.splice(i, 1);
+				i--;
+				len--;
+				insertIndex--;
+				continue;
+			}
+
+			if (piece.startLineNumber >= range.endLineNumber) {
+				insertIndex = Math.min(i, insertIndex);
+			}
+		}
+
+		this._pieces = arrays.arrayInsert(this._pieces, insertIndex, pieces);
+		return range;
+	}
+
 	public isComplete(): boolean {
 		return this._isComplete;
 	}
@@ -766,7 +878,7 @@ export class TokensStore2 {
 		}
 
 		const pieceIndex = TokensStore2._findFirstPieceWithLine(pieces, lineNumber);
-		const bTokens = this._pieces[pieceIndex].getLineTokens(lineNumber);
+		const bTokens = pieces[pieceIndex].getLineTokens(lineNumber);
 
 		if (!bTokens) {
 			return aTokens;

src/vs/editor/contrib/viewportSemanticTokens/viewportSemanticTokens.ts

@@ -3,15 +3,16 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { RunOnceScheduler, createCancelablePromise } from 'vs/base/common/async';
+import { RunOnceScheduler, createCancelablePromise, CancelablePromise } from 'vs/base/common/async';
 import { Disposable } from 'vs/base/common/lifecycle';
+import { Range } from 'vs/editor/common/core/range';
 import { ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { registerEditorContribution } from 'vs/editor/browser/editorExtensions';
 import { IEditorContribution } from 'vs/editor/common/editorCommon';
 import { ITextModel } from 'vs/editor/common/model';
-import { DocumentRangeSemanticTokensProviderRegistry, DocumentRangeSemanticTokensProvider } from 'vs/editor/common/modes';
+import { DocumentRangeSemanticTokensProviderRegistry, DocumentRangeSemanticTokensProvider, SemanticTokens } from 'vs/editor/common/modes';
 import { IModelService } from 'vs/editor/common/services/modelService';
-import { toMultilineTokens2 } from 'vs/editor/common/services/semanticTokensProviderStyling';
+import { toMultilineTokens2, SemanticTokensProviderStyling } from 'vs/editor/common/services/semanticTokensProviderStyling';
 
 class ViewportSemanticTokensContribution extends Disposable implements IEditorContribution {
 
@@ -23,6 +24,7 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 
 	private readonly _editor: ICodeEditor;
 	private readonly _tokenizeViewport: RunOnceScheduler;
+	private _outstandingRequests: CancelablePromise<SemanticTokens | null | undefined>[];
 
 	constructor(
 		editor: ICodeEditor,
@@ -31,13 +33,20 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 		super();
 		this._editor = editor;
 		this._tokenizeViewport = new RunOnceScheduler(() => this._tokenizeViewportNow(), 100);
+		this._outstandingRequests = [];
 		this._register(this._editor.onDidScrollChange(() => {
 			this._tokenizeViewport.schedule();
 		}));
 		this._register(this._editor.onDidChangeModel(() => {
+			this._cancelAll();
+			this._tokenizeViewport.schedule();
+		}));
+		this._register(this._editor.onDidChangeModelContent((e) => {
+			this._cancelAll();
 			this._tokenizeViewport.schedule();
 		}));
 		this._register(DocumentRangeSemanticTokensProviderRegistry.onDidChange(() => {
+			this._cancelAll();
 			this._tokenizeViewport.schedule();
 		}));
 	}
@@ -47,6 +56,22 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 		return (result.length > 0 ? result[0] : null);
 	}
 
+	private _cancelAll(): void {
+		for (const request of this._outstandingRequests) {
+			request.cancel();
+		}
+		this._outstandingRequests = [];
+	}
+
+	private _removeOutstandingRequest(req: CancelablePromise<SemanticTokens | null | undefined>): void {
+		for (let i = 0, len = this._outstandingRequests.length; i < len; i++) {
+			if (this._outstandingRequests[i] === req) {
+				this._outstandingRequests.splice(i, 1);
+				return;
+			}
+		}
+	}
+
 	private _tokenizeViewportNow(): void {
 		if (!this._editor.hasModel()) {
 			return;
@@ -61,14 +86,20 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 		}
 		const styling = this._modelService.getSemanticTokensProviderStyling(provider);
 		const visibleRanges = this._editor.getVisibleRanges();
-		const request = createCancelablePromise(token => Promise.resolve(provider.provideDocumentRangeSemanticTokens(model, visibleRanges[0], token)));
+
+		this._outstandingRequests = this._outstandingRequests.concat(visibleRanges.map(range => this._requestRange(model, range, provider, styling)));
+	}
+
+	private _requestRange(model: ITextModel, range: Range, provider: DocumentRangeSemanticTokensProvider, styling: SemanticTokensProviderStyling): CancelablePromise<SemanticTokens | null | undefined> {
+		const requestVersionId = model.getVersionId();
+		const request = createCancelablePromise(token => Promise.resolve(provider.provideDocumentRangeSemanticTokens(model, range, token)));
 		request.then((r) => {
-			if (!r || model.isDisposed()) {
+			if (!r || model.isDisposed() || model.getVersionId() !== requestVersionId) {
 				return;
 			}
-			const tokens = toMultilineTokens2(r, styling);
-			model.setPartialSemanticTokens(tokens);
-		});
+			model.setPartialSemanticTokens(range, toMultilineTokens2(r, styling));
+		}).then(() => this._removeOutstandingRequest(request), () => this._removeOutstandingRequest(request));
+		return request;
 	}
 }
 

src/vs/editor/test/common/model/tokensStore.test.ts

@@ -4,12 +4,13 @@
  *--------------------------------------------------------------------------------------------*/
 
 import * as assert from 'assert';
-import { MultilineTokens2, SparseEncodedTokens } from 'vs/editor/common/model/tokensStore';
+import { MultilineTokens2, SparseEncodedTokens, TokensStore2 } from 'vs/editor/common/model/tokensStore';
 import { Range } from 'vs/editor/common/core/range';
 import { TextModel } from 'vs/editor/common/model/textModel';
 import { IIdentifiedSingleEditOperation } from 'vs/editor/common/model';
 import { MetadataConsts, TokenMetadata } from 'vs/editor/common/modes';
 import { createTextModel } from 'vs/editor/test/common/editorTestUtils';
+import { LineTokens } from 'vs/editor/common/core/lineTokens';
 
 suite('TokensStore', () => {
 
@@ -212,4 +213,91 @@ suite('TokensStore', () => {
 		model.dispose();
 	});
 
+	test('partial tokens 1', () => {
+		const store = new TokensStore2();
+
+		// setPartial: [1,1 -> 31,2], [(5,5-10),(10,5-10),(15,5-10),(20,5-10),(25,5-10),(30,5-10)]
+		store.setPartial(new Range(1, 1, 31, 2), [
+			new MultilineTokens2(5, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 1,
+				5, 5, 10, 2,
+				10, 5, 10, 3,
+				15, 5, 10, 4,
+				20, 5, 10, 5,
+				25, 5, 10, 6,
+			])))
+		]);
+
+		// setPartial: [18,1 -> 42,1], [(20,5-10),(25,5-10),(30,5-10),(35,5-10),(40,5-10)]
+		store.setPartial(new Range(18, 1, 42, 1), [
+			new MultilineTokens2(20, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 4,
+				5, 5, 10, 5,
+				10, 5, 10, 6,
+				15, 5, 10, 7,
+				20, 5, 10, 8,
+			])))
+		]);
+
+		// setPartial: [1,1 -> 31,2], [(5,5-10),(10,5-10),(15,5-10),(20,5-10),(25,5-10),(30,5-10)]
+		store.setPartial(new Range(1, 1, 31, 2), [
+			new MultilineTokens2(5, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 1,
+				5, 5, 10, 2,
+				10, 5, 10, 3,
+				15, 5, 10, 4,
+				20, 5, 10, 5,
+				25, 5, 10, 6,
+			])))
+		]);
+
+		const lineTokens = store.addSemanticTokens(10, new LineTokens(new Uint32Array([12, 1]), `enum Enum1 {`));
+		assert.equal(lineTokens.getCount(), 3);
+	});
+
+	test('partial tokens 2', () => {
+		const store = new TokensStore2();
+
+		// setPartial: [1,1 -> 31,2], [(5,5-10),(10,5-10),(15,5-10),(20,5-10),(25,5-10),(30,5-10)]
+		store.setPartial(new Range(1, 1, 31, 2), [
+			new MultilineTokens2(5, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 1,
+				5, 5, 10, 2,
+				10, 5, 10, 3,
+				15, 5, 10, 4,
+				20, 5, 10, 5,
+				25, 5, 10, 6,
+			])))
+		]);
+
+		// setPartial: [6,1 -> 36,2], [(10,5-10),(15,5-10),(20,5-10),(25,5-10),(30,5-10),(35,5-10)]
+		store.setPartial(new Range(6, 1, 36, 2), [
+			new MultilineTokens2(10, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 2,
+				5, 5, 10, 3,
+				10, 5, 10, 4,
+				15, 5, 10, 5,
+				20, 5, 10, 6,
+			])))
+		]);
+
+		// setPartial: [17,1 -> 42,1], [(20,5-10),(25,5-10),(30,5-10),(35,5-10),(40,5-10)]
+		store.setPartial(new Range(17, 1, 42, 1), [
+			new MultilineTokens2(20, new SparseEncodedTokens(new Uint32Array([
+				0, 5, 10, 4,
+				5, 5, 10, 5,
+				10, 5, 10, 6,
+				15, 5, 10, 7,
+				20, 5, 10, 8,
+			])))
+		]);
+
+		const lineTokens = store.addSemanticTokens(20, new LineTokens(new Uint32Array([12, 1]), `enum Enum1 {`));
+		assert.equal(lineTokens.getCount(), 3);
+	});
+
+	// tokensStore.ts:878 ==> INSIDE PARTIAL SET: PIECES: [(5,5-10)]
+	// tokensStore.ts:894 ==> AFTER PARTIAL SET: PIECES: [(10,5-10),(15,5-10),(20,5-10),(25,5-10),(30,5-10),(35,5-10)], [(5,5-10)]
+	// tokensStore.ts:878 ==> INSIDE PARTIAL SET: PIECES: [(10,5-10),(15,5-10)], [(5,5-10)]
+	// tokensStore.ts:894 ==> AFTER PARTIAL SET: PIECES: [(20,5-10),(25,5-10),(30,5-10),(35,5-10),(40,5-10)], [(10,5-10),(15,5-10)], [(5,5-10)]
 });