Merge pull request #135602 from rchiodo/dev/rchiodo/multiple_providers

Allow semantic tokens to be provided by more than one provider

src/vs/editor/common/services/getSemanticTokens.ts

@@ -23,30 +23,111 @@ export function isSemanticTokensEdits(v: SemanticTokens | SemanticTokensEdits):
 	return v && Array.isArray((<SemanticTokensEdits>v).edits);
 }
 
-export interface IDocumentSemanticTokensResult {
-	provider: DocumentSemanticTokensProvider;
-	request: Promise<SemanticTokens | SemanticTokensEdits | null | undefined>;
+export class DocumentSemanticTokensResult {
+	constructor(
+		public readonly provider: DocumentSemanticTokensProvider,
+		public readonly tokens: SemanticTokens | SemanticTokensEdits | null,
+	) { }
 }
 
-export function getDocumentSemanticTokens(model: ITextModel, lastResultId: string | null, token: CancellationToken): IDocumentSemanticTokensResult | null {
-	const provider = _getDocumentSemanticTokensProvider(model);
-	if (!provider) {
-		return null;
+export function hasDocumentSemanticTokensProvider(model: ITextModel): boolean {
+	return DocumentSemanticTokensProviderRegistry.has(model);
+}
+
+function getDocumentSemanticTokensProviders(model: ITextModel): DocumentSemanticTokensProvider[] {
+	const groups = DocumentSemanticTokensProviderRegistry.orderedGroups(model);
+	return (groups.length > 0 ? groups[0] : []);
+}
+
+export async function getDocumentSemanticTokens(model: ITextModel, lastProvider: DocumentSemanticTokensProvider | null, lastResultId: string | null, token: CancellationToken): Promise<DocumentSemanticTokensResult | null> {
+	const providers = getDocumentSemanticTokensProviders(model);
+
+	// Get tokens from all providers at the same time.
+	const results = await Promise.all(providers.map(async (provider) => {
+		let result: SemanticTokens | SemanticTokensEdits | null | undefined;
+		try {
+			result = await provider.provideDocumentSemanticTokens(model, (provider === lastProvider ? lastResultId : null), token);
+		} catch (err) {
+			onUnexpectedExternalError(err);
+			result = null;
+		}
+
+		if (!result || (!isSemanticTokens(result) && !isSemanticTokensEdits(result))) {
+			result = null;
+		}
+
+		return new DocumentSemanticTokensResult(provider, result);
+	}));
+
+	// Try to return the first result with actual tokens
+	for (const result of results) {
+		if (result.tokens) {
+			return result;
+		}
+	}
+
+	// Return the first result, even if it doesn't have tokens
+	if (results.length > 0) {
+		return results[0];
 	}
-	return {
-		provider: provider,
-		request: Promise.resolve(provider.provideDocumentSemanticTokens(model, lastResultId, token))
-	};
+
+	return null;
 }
 
-function _getDocumentSemanticTokensProvider(model: ITextModel): DocumentSemanticTokensProvider | null {
-	const result = DocumentSemanticTokensProviderRegistry.ordered(model);
+function _getDocumentSemanticTokensProviderHighestGroup(model: ITextModel): DocumentSemanticTokensProvider[] | null {
+	const result = DocumentSemanticTokensProviderRegistry.orderedGroups(model);
 	return (result.length > 0 ? result[0] : null);
 }
 
-export function getDocumentRangeSemanticTokensProvider(model: ITextModel): DocumentRangeSemanticTokensProvider | null {
-	const result = DocumentRangeSemanticTokensProviderRegistry.ordered(model);
-	return (result.length > 0 ? result[0] : null);
+class DocumentRangeSemanticTokensResult {
+	constructor(
+		public readonly provider: DocumentRangeSemanticTokensProvider,
+		public readonly tokens: SemanticTokens | null,
+	) { }
+}
+
+export function hasDocumentRangeSemanticTokensProvider(model: ITextModel): boolean {
+	return DocumentRangeSemanticTokensProviderRegistry.has(model);
+}
+
+function getDocumentRangeSemanticTokensProviders(model: ITextModel): DocumentRangeSemanticTokensProvider[] {
+	const groups = DocumentRangeSemanticTokensProviderRegistry.orderedGroups(model);
+	return (groups.length > 0 ? groups[0] : []);
+}
+
+export async function getDocumentRangeSemanticTokens(model: ITextModel, range: Range, token: CancellationToken): Promise<DocumentRangeSemanticTokensResult | null> {
+	const providers = getDocumentRangeSemanticTokensProviders(model);
+
+	// Get tokens from all providers at the same time.
+	const results = await Promise.all(providers.map(async (provider) => {
+		let result: SemanticTokens | null | undefined;
+		try {
+			result = await provider.provideDocumentRangeSemanticTokens(model, range, token);
+		} catch (err) {
+			onUnexpectedExternalError(err);
+			result = null;
+		}
+
+		if (!result || !isSemanticTokens(result)) {
+			result = null;
+		}
+
+		return new DocumentRangeSemanticTokensResult(provider, result);
+	}));
+
+	// Try to return the first result with actual tokens
+	for (const result of results) {
+		if (result.tokens) {
+			return result;
+		}
+	}
+
+	// Return the first result, even if it doesn't have tokens
+	if (results.length > 0) {
+		return results[0];
+	}
+
+	return null;
 }
 
 CommandsRegistry.registerCommand('_provideDocumentSemanticTokensLegend', async (accessor, ...args): Promise<SemanticTokensLegend | undefined> => {
@@ -58,13 +139,13 @@ CommandsRegistry.registerCommand('_provideDocumentSemanticTokensLegend', async (
 		return undefined;
 	}
 
-	const provider = _getDocumentSemanticTokensProvider(model);
-	if (!provider) {
+	const providers = _getDocumentSemanticTokensProviderHighestGroup(model);
+	if (!providers) {
 		// there is no provider => fall back to a document range semantic tokens provider
 		return accessor.get(ICommandService).executeCommand('_provideDocumentRangeSemanticTokensLegend', uri);
 	}
 
-	return provider.getLegend();
+	return providers[0].getLegend();
 });
 
 CommandsRegistry.registerCommand('_provideDocumentSemanticTokens', async (accessor, ...args): Promise<VSBuffer | undefined> => {
@@ -76,39 +157,35 @@ CommandsRegistry.registerCommand('_provideDocumentSemanticTokens', async (access
 		return undefined;
 	}
 
-	const r = getDocumentSemanticTokens(model, null, CancellationToken.None);
-	if (!r) {
+	if (!hasDocumentSemanticTokensProvider(model)) {
 		// there is no provider => fall back to a document range semantic tokens provider
 		return accessor.get(ICommandService).executeCommand('_provideDocumentRangeSemanticTokens', uri, model.getFullModelRange());
 	}
 
-	const { provider, request } = r;
-
-	let result: SemanticTokens | SemanticTokensEdits | null | undefined;
-	try {
-		result = await request;
-	} catch (err) {
-		onUnexpectedExternalError(err);
+	const r = await getDocumentSemanticTokens(model, null, null, CancellationToken.None);
+	if (!r) {
 		return undefined;
 	}
 
-	if (!result || !isSemanticTokens(result)) {
+	const { provider, tokens } = r;
+
+	if (!tokens || !isSemanticTokens(tokens)) {
 		return undefined;
 	}
 
 	const buff = encodeSemanticTokensDto({
 		id: 0,
 		type: 'full',
-		data: result.data
+		data: tokens.data
 	});
-	if (result.resultId) {
-		provider.releaseDocumentSemanticTokens(result.resultId);
+	if (tokens.resultId) {
+		provider.releaseDocumentSemanticTokens(tokens.resultId);
 	}
 	return buff;
 });
 
 CommandsRegistry.registerCommand('_provideDocumentRangeSemanticTokensLegend', async (accessor, ...args): Promise<SemanticTokensLegend | undefined> => {
-	const [uri] = args;
+	const [uri, range] = args;
 	assertType(uri instanceof URI);
 
 	const model = accessor.get(IModelService).getModel(uri);
@@ -116,12 +193,31 @@ CommandsRegistry.registerCommand('_provideDocumentRangeSemanticTokensLegend', as
 		return undefined;
 	}
 
-	const provider = getDocumentRangeSemanticTokensProvider(model);
-	if (!provider) {
+	const providers = getDocumentRangeSemanticTokensProviders(model);
+	if (providers.length === 0) {
+		// no providers
+		return undefined;
+	}
+
+	if (providers.length === 1) {
+		// straight forward case, just a single provider
+		return providers[0].getLegend();
+	}
+
+	if (!range || !Range.isIRange(range)) {
+		// if no range is provided, we cannot support multiple providers
+		// as we cannot fall back to the one which would give results
+		// => return the first legend for backwards compatibility and print a warning
+		console.warn(`provideDocumentRangeSemanticTokensLegend might be out-of-sync with provideDocumentRangeSemanticTokens unless a range argument is passed in`);
+		return providers[0].getLegend();
+	}
+
+	const result = await getDocumentRangeSemanticTokens(model, Range.lift(range), CancellationToken.None);
+	if (!result) {
 		return undefined;
 	}
 
-	return provider.getLegend();
+	return result.provider.getLegend();
 });
 
 CommandsRegistry.registerCommand('_provideDocumentRangeSemanticTokens', async (accessor, ...args): Promise<VSBuffer | undefined> => {
@@ -134,27 +230,15 @@ CommandsRegistry.registerCommand('_provideDocumentRangeSemanticTokens', async (a
 		return undefined;
 	}
 
-	const provider = getDocumentRangeSemanticTokensProvider(model);
-	if (!provider) {
-		// there is no provider
-		return undefined;
-	}
-
-	let result: SemanticTokens | null | undefined;
-	try {
-		result = await provider.provideDocumentRangeSemanticTokens(model, Range.lift(range), CancellationToken.None);
-	} catch (err) {
-		onUnexpectedExternalError(err);
-		return undefined;
-	}
-
-	if (!result || !isSemanticTokens(result)) {
+	const result = await getDocumentRangeSemanticTokens(model, Range.lift(range), CancellationToken.None);
+	if (!result || !result.tokens) {
+		// there is no provider or it didn't return tokens
 		return undefined;
 	}
 
 	return encodeSemanticTokensDto({
 		id: 0,
 		type: 'full',
-		data: result.data
+		data: result.tokens.data
 	});
 });

src/vs/editor/common/services/modelServiceImpl.ts

@@ -29,7 +29,7 @@ import { StringSHA1 } from 'vs/base/common/hash';
 import { EditStackElement, isEditStackElement } from 'vs/editor/common/model/editStack';
 import { Schemas } from 'vs/base/common/network';
 import { SemanticTokensProviderStyling, toMultilineTokens2 } from 'vs/editor/common/services/semanticTokensProviderStyling';
-import { getDocumentSemanticTokens, isSemanticTokens, isSemanticTokensEdits } from 'vs/editor/common/services/getSemanticTokens';
+import { getDocumentSemanticTokens, hasDocumentSemanticTokensProvider, isSemanticTokens, isSemanticTokensEdits } from 'vs/editor/common/services/getSemanticTokens';
 import { equals } from 'vs/base/common/objects';
 import { ILanguageConfigurationService } from 'vs/editor/common/modes/languageConfigurationRegistry';
 
@@ -724,13 +724,13 @@ class SemanticStyling extends Disposable {
 
 class SemanticTokensResponse {
 	constructor(
-		private readonly _provider: DocumentSemanticTokensProvider,
+		public readonly provider: DocumentSemanticTokensProvider,
 		public readonly resultId: string | undefined,
 		public readonly data: Uint32Array
 	) { }
 
 	public dispose(): void {
-		this._provider.releaseDocumentSemanticTokens(this.resultId);
+		this.provider.releaseDocumentSemanticTokens(this.resultId);
 	}
 }
 
@@ -820,10 +820,7 @@ export class ModelSemanticColoring extends Disposable {
 			return;
 		}
 
-		const cancellationTokenSource = new CancellationTokenSource();
-		const lastResultId = this._currentDocumentResponse ? this._currentDocumentResponse.resultId || null : null;
-		const r = getDocumentSemanticTokens(this._model, lastResultId, cancellationTokenSource.token);
-		if (!r) {
+		if (!hasDocumentSemanticTokensProvider(this._model)) {
 			// there is no provider
 			if (this._currentDocumentResponse) {
 				// there are semantic tokens set
@@ -832,7 +829,10 @@ export class ModelSemanticColoring extends Disposable {
 			return;
 		}
 
-		const { provider, request } = r;
+		const cancellationTokenSource = new CancellationTokenSource();
+		const lastProvider = this._currentDocumentResponse ? this._currentDocumentResponse.provider : null;
+		const lastResultId = this._currentDocumentResponse ? this._currentDocumentResponse.resultId || null : null;
+		const request = getDocumentSemanticTokens(this._model, lastProvider, lastResultId, cancellationTokenSource.token);
 		this._currentDocumentRequestCancellationTokenSource = cancellationTokenSource;
 
 		const pendingChanges: IModelContentChangedEvent[] = [];
@@ -840,12 +840,17 @@ export class ModelSemanticColoring extends Disposable {
 			pendingChanges.push(e);
 		});
 
-		const styling = this._semanticStyling.get(provider);
-
 		request.then((res) => {
 			this._currentDocumentRequestCancellationTokenSource = null;
 			contentChangeListener.dispose();
-			this._setDocumentSemanticTokens(provider, res || null, styling, pendingChanges);
+
+			if (!res) {
+				this._setDocumentSemanticTokens(null, null, null, pendingChanges);
+			} else {
+				const { provider, tokens } = res;
+				const styling = this._semanticStyling.get(provider);
+				this._setDocumentSemanticTokens(provider, tokens || null, styling, pendingChanges);
+			}
 		}, (err) => {
 			const isExpectedError = err && (errors.isPromiseCanceledError(err) || (typeof err.message === 'string' && err.message.indexOf('busy') !== -1));
 			if (!isExpectedError) {

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

@@ -10,11 +10,11 @@ import { registerEditorContribution } from 'vs/editor/browser/editorExtensions';
 import { Range } from 'vs/editor/common/core/range';
 import { IEditorContribution } from 'vs/editor/common/editorCommon';
 import { ITextModel } from 'vs/editor/common/model';
-import { DocumentRangeSemanticTokensProvider, DocumentRangeSemanticTokensProviderRegistry, SemanticTokens } from 'vs/editor/common/modes';
-import { getDocumentRangeSemanticTokensProvider } from 'vs/editor/common/services/getSemanticTokens';
+import { DocumentRangeSemanticTokensProviderRegistry } from 'vs/editor/common/modes';
+import { getDocumentRangeSemanticTokens, hasDocumentRangeSemanticTokensProvider } from 'vs/editor/common/services/getSemanticTokens';
 import { IModelService } from 'vs/editor/common/services/modelService';
 import { isSemanticColoringEnabled, SEMANTIC_HIGHLIGHTING_SETTING_ID } from 'vs/editor/common/services/modelServiceImpl';
-import { SemanticTokensProviderStyling, toMultilineTokens2 } from 'vs/editor/common/services/semanticTokensProviderStyling';
+import { toMultilineTokens2 } from 'vs/editor/common/services/semanticTokensProviderStyling';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { IThemeService } from 'vs/platform/theme/common/themeService';
 
@@ -28,7 +28,7 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 
 	private readonly _editor: ICodeEditor;
 	private readonly _tokenizeViewport: RunOnceScheduler;
-	private _outstandingRequests: CancelablePromise<SemanticTokens | null | undefined>[];
+	private _outstandingRequests: CancelablePromise<any>[];
 
 	constructor(
 		editor: ICodeEditor,
@@ -74,7 +74,7 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 		this._outstandingRequests = [];
 	}
 
-	private _removeOutstandingRequest(req: CancelablePromise<SemanticTokens | null | undefined>): void {
+	private _removeOutstandingRequest(req: CancelablePromise<any>): void {
 		for (let i = 0, len = this._outstandingRequests.length; i < len; i++) {
 			if (this._outstandingRequests[i] === req) {
 				this._outstandingRequests.splice(i, 1);
@@ -97,27 +97,27 @@ class ViewportSemanticTokensContribution extends Disposable implements IEditorCo
 			}
 			return;
 		}
-		const provider = getDocumentRangeSemanticTokensProvider(model);
-		if (!provider) {
+		if (!hasDocumentRangeSemanticTokensProvider(model)) {
 			if (model.hasSomeSemanticTokens()) {
 				model.setSemanticTokens(null, false);
 			}
 			return;
 		}
-		const styling = this._modelService.getSemanticTokensProviderStyling(provider);
 		const visibleRanges = this._editor.getVisibleRangesPlusViewportAboveBelow();
 
-		this._outstandingRequests = this._outstandingRequests.concat(visibleRanges.map(range => this._requestRange(model, range, provider, styling)));
+		this._outstandingRequests = this._outstandingRequests.concat(visibleRanges.map(range => this._requestRange(model, range)));
 	}
 
-	private _requestRange(model: ITextModel, range: Range, provider: DocumentRangeSemanticTokensProvider, styling: SemanticTokensProviderStyling): CancelablePromise<SemanticTokens | null | undefined> {
+	private _requestRange(model: ITextModel, range: Range): CancelablePromise<any> {
 		const requestVersionId = model.getVersionId();
-		const request = createCancelablePromise(token => Promise.resolve(provider.provideDocumentRangeSemanticTokens(model, range, token)));
+		const request = createCancelablePromise(token => Promise.resolve(getDocumentRangeSemanticTokens(model, range, token)));
 		request.then((r) => {
-			if (!r || model.isDisposed() || model.getVersionId() !== requestVersionId) {
+			if (!r || !r.tokens || model.isDisposed() || model.getVersionId() !== requestVersionId) {
 				return;
 			}
-			model.setPartialSemanticTokens(range, toMultilineTokens2(r, styling, model.getLanguageId()));
+			const { provider, tokens: result } = r;
+			const styling = this._modelService.getSemanticTokensProviderStyling(provider);
+			model.setPartialSemanticTokens(range, toMultilineTokens2(result, styling, model.getLanguageId()));
 		}).then(() => this._removeOutstandingRequest(request), () => this._removeOutstandingRequest(request));
 		return request;
 	}

src/vs/editor/test/common/services/modelService.test.ts

@@ -32,6 +32,7 @@ import { IModelService } from 'vs/editor/common/services/modelService';
 import { IModeService } from 'vs/editor/common/services/modeService';
 import { TestTextResourcePropertiesService } from 'vs/editor/test/common/services/testTextResourcePropertiesService';
 import { TestLanguageConfigurationService } from 'vs/editor/test/common/modes/testLanguageConfigurationService';
+import { getDocumentSemanticTokens, isSemanticTokens } from 'vs/editor/common/services/getSemanticTokens';
 
 const GENERATE_TESTS = false;
 
@@ -486,6 +487,81 @@ suite('ModelSemanticColoring', () => {
 		// assert that it got called twice
 		assert.strictEqual(callCount, 2);
 	});
+
+	test('DocumentSemanticTokens should be pick the token provider with actual items', async () => {
+
+		let callCount = 0;
+		disposables.add(ModesRegistry.registerLanguage({ id: 'testMode2' }));
+		disposables.add(DocumentSemanticTokensProviderRegistry.register('testMode2', new class implements DocumentSemanticTokensProvider {
+			getLegend(): SemanticTokensLegend {
+				return { tokenTypes: ['class1'], tokenModifiers: [] };
+			}
+			async provideDocumentSemanticTokens(model: ITextModel, lastResultId: string | null, token: CancellationToken): Promise<SemanticTokens | SemanticTokensEdits | null> {
+				callCount++;
+				// For a secondary request return a different value
+				if (lastResultId) {
+					return {
+						data: new Uint32Array([2, 1, 1, 1, 1, 0, 2, 1, 1, 1])
+					};
+				}
+				return {
+					resultId: '1',
+					data: new Uint32Array([0, 1, 1, 1, 1, 0, 2, 1, 1, 1])
+				};
+			}
+			releaseDocumentSemanticTokens(resultId: string | undefined): void {
+			}
+		}));
+		disposables.add(DocumentSemanticTokensProviderRegistry.register('testMode2', new class implements DocumentSemanticTokensProvider {
+			getLegend(): SemanticTokensLegend {
+				return { tokenTypes: ['class2'], tokenModifiers: [] };
+			}
+			async provideDocumentSemanticTokens(model: ITextModel, lastResultId: string | null, token: CancellationToken): Promise<SemanticTokens | SemanticTokensEdits | null> {
+				callCount++;
+				return null;
+			}
+			releaseDocumentSemanticTokens(resultId: string | undefined): void {
+			}
+		}));
+
+		function toArr(arr: Uint32Array): number[] {
+			let result: number[] = [];
+			for (let i = 0; i < arr.length; i++) {
+				result[i] = arr[i];
+			}
+			return result;
+		}
+
+		const textModel = modelService.createModel('Hello world 2', modeService.create('testMode2'));
+		try {
+			let result = await getDocumentSemanticTokens(textModel, null, null, CancellationToken.None);
+			assert.ok(result, `We should have tokens (1)`);
+			assert.ok(result.tokens, `Tokens are found from multiple providers (1)`);
+			assert.ok(isSemanticTokens(result.tokens), `Tokens are full (1)`);
+			assert.ok(result.tokens.resultId, `Token result id found from multiple providers (1)`);
+			assert.deepStrictEqual(toArr(result.tokens.data), [0, 1, 1, 1, 1, 0, 2, 1, 1, 1], `Token data returned for multiple providers (1)`);
+			assert.deepStrictEqual(callCount, 2, `Called both token providers (1)`);
+			assert.deepStrictEqual(result.provider.getLegend(), { tokenTypes: ['class1'], tokenModifiers: [] }, `Legend matches the tokens (1)`);
+
+			// Make a second request. Make sure we get the secondary value
+			result = await getDocumentSemanticTokens(textModel, result.provider, result.tokens.resultId, CancellationToken.None);
+			assert.ok(result, `We should have tokens (2)`);
+			assert.ok(result.tokens, `Tokens are found from multiple providers (2)`);
+			assert.ok(isSemanticTokens(result.tokens), `Tokens are full (2)`);
+			assert.ok(!result.tokens.resultId, `Token result id found from multiple providers (2)`);
+			assert.deepStrictEqual(toArr(result.tokens.data), [2, 1, 1, 1, 1, 0, 2, 1, 1, 1], `Token data returned for multiple providers (2)`);
+			assert.deepStrictEqual(callCount, 4, `Called both token providers (2)`);
+			assert.deepStrictEqual(result.provider.getLegend(), { tokenTypes: ['class1'], tokenModifiers: [] }, `Legend matches the tokens (2)`);
+		} finally {
+			disposables.clear();
+
+			// Wait for scheduler to finish
+			await timeout(0);
+
+			// Now dispose the text model
+			textModel.dispose();
+		}
+	});
 });
 
 function assertComputeEdits(lines1: string[], lines2: string[]): void {

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

@@ -216,7 +216,7 @@ const newCommands: ApiCommand[] = [
 	),
 	new ApiCommand(
 		'vscode.provideDocumentRangeSemanticTokensLegend', '_provideDocumentRangeSemanticTokensLegend', 'Provide semantic tokens legend for a document range',
-		[ApiCommandArgument.Uri],
+		[ApiCommandArgument.Uri, ApiCommandArgument.Range.optional()],
 		new ApiCommandResult<modes.SemanticTokensLegend, types.SemanticTokensLegend | undefined>('A promise that resolves to SemanticTokensLegend.', value => {
 			if (!value) {
 				return undefined;