`SemanticTokensBuilder`: Add `push` overload that accepts objects and support...

`SemanticTokensBuilder`: Add `push` overload that accepts objects and support out-of-order calls (for #93614)

src/vs/vscode.proposed.d.ts

@@ -257,8 +257,9 @@ declare module 'vscode' {
 	}
 
 	export class SemanticTokensBuilder {
-		constructor();
+		constructor(legend?: SemanticTokensLegend);
 		push(line: number, char: number, length: number, tokenType: number, tokenModifiers: number): void;
+		push(range: Range, tokenType: string, tokenModifiers?: string[]): void;
 		build(resultId?: string): SemanticTokens;
 	}
 

src/vs/workbench/api/common/extHostTypes.ts

@@ -2428,24 +2428,126 @@ export class SemanticTokensLegend {
 	}
 }
 
+function isStrArrayOrUndefined(arg: any): arg is string[] | undefined {
+	if (typeof arg === 'undefined') {
+		return true;
+	}
+	if (Array.isArray(arg)) {
+		for (const element of arg) {
+			if (typeof element !== 'string') {
+				return false;
+			}
+		}
+		return true;
+	}
+	return false;
+}
+
 export class SemanticTokensBuilder {
 
 	private _prevLine: number;
 	private _prevChar: number;
+	private _dataIsSortedAndDeltaEncoded: boolean;
 	private _data: number[];
 	private _dataLen: number;
+	private _tokenTypeStrToInt: Map<string, number>;
+	private _tokenModifierStrToInt: Map<string, number>;
+	private _hasLegend: boolean;
 
-	constructor() {
+	constructor(legend?: vscode.SemanticTokensLegend) {
 		this._prevLine = 0;
 		this._prevChar = 0;
+		this._dataIsSortedAndDeltaEncoded = true;
 		this._data = [];
 		this._dataLen = 0;
+		this._tokenTypeStrToInt = new Map<string, number>();
+		this._tokenModifierStrToInt = new Map<string, number>();
+		this._hasLegend = false;
+		if (legend) {
+			this._hasLegend = true;
+			for (let i = 0, len = legend.tokenTypes.length; i < len; i++) {
+				this._tokenTypeStrToInt.set(legend.tokenTypes[i], i);
+			}
+			for (let i = 0, len = legend.tokenModifiers.length; i < len; i++) {
+				this._tokenModifierStrToInt.set(legend.tokenModifiers[i], i);
+			}
+		}
 	}
 
-	public push(line: number, char: number, length: number, tokenType: number, tokenModifiers: number): void {
+	public push(line: number, char: number, length: number, tokenType: number, tokenModifiers: number): void;
+	public push(range: Range, tokenType: string, tokenModifiers?: string[]): void;
+	public push(arg0: any, arg1: any, arg2: any, arg3?: any, arg4?: any): void {
+		if (typeof arg0 === 'number' && typeof arg1 === 'number' && typeof arg2 === 'number' && typeof arg3 === 'number' && typeof arg4 === 'number') {
+			// 1st overload
+			return this._pushEncoded(arg0, arg1, arg2, arg3, arg4);
+		}
+		if (Range.isRange(arg0) && typeof arg1 === 'string' && isStrArrayOrUndefined(arg2)) {
+			// 2nd overload
+			return this._push(arg0, arg1, arg2);
+		}
+		throw illegalArgument();
+	}
+
+	private _push(range: vscode.Range, tokenType: string, tokenModifiers?: string[]): void {
+		if (!this._hasLegend) {
+			throw new Error('Legend must be provided in constructor');
+		}
+		if (range.start.line !== range.end.line) {
+			throw new Error('`range` cannot span multiple lines');
+		}
+		if (!this._tokenTypeStrToInt.has(tokenType)) {
+			throw new Error('`tokenType` is not in the provided legend');
+		}
+		const line = range.start.line;
+		const char = range.start.character;
+		const length = range.end.character - range.start.character;
+		const nTokenType = this._tokenTypeStrToInt.get(tokenType)!;
+		let nTokenModifiers = 0;
+		if (tokenModifiers) {
+			for (const tokenModifier of tokenModifiers) {
+				if (!this._tokenModifierStrToInt.has(tokenModifier)) {
+					throw new Error('`tokenModifier` is not in the provided legend');
+				}
+				const nTokenModifier = this._tokenModifierStrToInt.get(tokenModifier)!;
+				nTokenModifiers |= (1 << nTokenModifier) >>> 0;
+			}
+		}
+		this._pushEncoded(line, char, length, nTokenType, nTokenModifiers);
+	}
+
+	private _pushEncoded(line: number, char: number, length: number, tokenType: number, tokenModifiers: number): void {
+		if (this._dataIsSortedAndDeltaEncoded && (line < this._prevLine || (line === this._prevLine && char < this._prevChar))) {
+			// push calls were ordered and are no longer ordered
+			this._dataIsSortedAndDeltaEncoded = false;
+
+			// Remove delta encoding from data
+			const tokenCount = (this._data.length / 5) | 0;
+			let prevLine = 0;
+			let prevChar = 0;
+			for (let i = 0; i < tokenCount; i++) {
+				let line = this._data[5 * i];
+				let char = this._data[5 * i + 1];
+
+				if (line === 0) {
+					// on the same line as previous token
+					line = prevLine;
+					char += prevChar;
+				} else {
+					// on a different line than previous token
+					line += prevLine;
+				}
+
+				this._data[5 * i] = line;
+				this._data[5 * i + 1] = char;
+
+				prevLine = line;
+				prevChar = char;
+			}
+		}
+
 		let pushLine = line;
 		let pushChar = char;
-		if (this._dataLen > 0) {
+		if (this._dataIsSortedAndDeltaEncoded && this._dataLen > 0) {
 			pushLine -= this._prevLine;
 			if (pushLine === 0) {
 				pushChar -= this._prevChar;
@@ -2462,7 +2564,54 @@ export class SemanticTokensBuilder {
 		this._prevChar = char;
 	}
 
+	private static _sortAndDeltaEncode(data: number[]): Uint32Array {
+		let pos: number[] = [];
+		const tokenCount = (data.length / 5) | 0;
+		for (let i = 0; i < tokenCount; i++) {
+			pos[i] = i;
+		}
+		pos.sort((a, b) => {
+			const aLine = data[5 * a];
+			const bLine = data[5 * b];
+			if (aLine === bLine) {
+				const aChar = data[5 * a + 1];
+				const bChar = data[5 * b + 1];
+				return aChar - bChar;
+			}
+			return aLine - bLine;
+		});
+		const result = new Uint32Array(data.length);
+		let prevLine = 0;
+		let prevChar = 0;
+		for (let i = 0; i < tokenCount; i++) {
+			const srcOffset = 5 * pos[i];
+			const line = data[srcOffset + 0];
+			const char = data[srcOffset + 1];
+			const length = data[srcOffset + 2];
+			const tokenType = data[srcOffset + 3];
+			const tokenModifiers = data[srcOffset + 4];
+
+			const pushLine = line - prevLine;
+			const pushChar = (pushLine === 0 ? char - prevChar : char);
+
+			const dstOffset = 5 * i;
+			result[dstOffset + 0] = pushLine;
+			result[dstOffset + 1] = pushChar;
+			result[dstOffset + 2] = length;
+			result[dstOffset + 3] = tokenType;
+			result[dstOffset + 4] = tokenModifiers;
+
+			prevLine = line;
+			prevChar = char;
+		}
+
+		return result;
+	}
+
 	public build(resultId?: string): SemanticTokens {
+		if (!this._dataIsSortedAndDeltaEncoded) {
+			return new SemanticTokens(SemanticTokensBuilder._sortAndDeltaEncode(this._data), resultId);
+		}
 		return new SemanticTokens(new Uint32Array(this._data), resultId);
 	}
 }

src/vs/workbench/test/browser/api/extHostTypes.test.ts

@@ -566,4 +566,64 @@ suite('ExtHostTypes', function () {
 		assert.ok(!types.CodeActionKind.RefactorExtract.intersects(types.CodeActionKind.Empty.append('other').append('refactor')));
 		assert.ok(!types.CodeActionKind.RefactorExtract.intersects(types.CodeActionKind.Empty.append('refactory')));
 	});
+
+	function toArr(uint32Arr: Uint32Array): number[] {
+		const r = [];
+		for (let i = 0, len = uint32Arr.length; i < len; i++) {
+			r[i] = uint32Arr[i];
+		}
+		return r;
+	}
+
+	test('SemanticTokensBuilder simple', () => {
+		const builder = new types.SemanticTokensBuilder();
+		builder.push(1, 0, 5, 1, 1);
+		builder.push(1, 10, 4, 2, 2);
+		builder.push(2, 2, 3, 2, 2);
+		assert.deepEqual(toArr(builder.build().data), [
+			1, 0, 5, 1, 1,
+			0, 10, 4, 2, 2,
+			1, 2, 3, 2, 2
+		]);
+	});
+
+	test('SemanticTokensBuilder out of order 1', () => {
+		const builder = new types.SemanticTokensBuilder();
+		builder.push(2, 0, 5, 1, 1);
+		builder.push(2, 10, 1, 2, 2);
+		builder.push(2, 15, 2, 3, 3);
+		builder.push(1, 0, 4, 4, 4);
+		assert.deepEqual(toArr(builder.build().data), [
+			1, 0, 4, 4, 4,
+			1, 0, 5, 1, 1,
+			0, 10, 1, 2, 2,
+			0, 5, 2, 3, 3
+		]);
+	});
+
+	test('SemanticTokensBuilder out of order 2', () => {
+		const builder = new types.SemanticTokensBuilder();
+		builder.push(2, 10, 5, 1, 1);
+		builder.push(2, 2, 4, 2, 2);
+		assert.deepEqual(toArr(builder.build().data), [
+			2, 2, 4, 2, 2,
+			0, 8, 5, 1, 1
+		]);
+	});
+
+	test('SemanticTokensBuilder with legend', () => {
+		const legend = new types.SemanticTokensLegend(
+			['aType', 'bType', 'cType', 'dType'],
+			['mod0', 'mod1', 'mod2', 'mod3', 'mod4', 'mod5']
+		);
+		const builder = new types.SemanticTokensBuilder(legend);
+		builder.push(new types.Range(1, 0, 1, 5), 'bType');
+		builder.push(new types.Range(2, 0, 2, 4), 'cType', ['mod0', 'mod5']);
+		builder.push(new types.Range(3, 0, 3, 3), 'dType', ['mod2', 'mod4']);
+		assert.deepEqual(toArr(builder.build().data), [
+			1, 0, 5, 1, 0,
+			1, 0, 4, 2, 1 | (1 << 5),
+			1, 0, 3, 3, (1 << 2) | (1 << 4)
+		]);
+	});
 });