notebooks: better memorization of preferred renderers

Fixes #125876
microsoft · Oct 22, 2021 · 24af01a · 24af01a
1 parent baa380d
commit 24af01a
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 79 deletions.
diff --git a/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts b/src/vs/workbench/contrib/notebook/browser/notebookServiceImpl.ts
@@ -11,6 +11,7 @@ import { Lazy } from 'vs/base/common/lazy';
 import { Disposable, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { ResourceMap } from 'vs/base/common/map';
 import { Schemas } from 'vs/base/common/network';
+import { isDefined } from 'vs/base/common/types';
 import { URI } from 'vs/base/common/uri';
 import { ICodeEditorService } from 'vs/editor/browser/services/codeEditorService';
 import { IEditorOptions } from 'vs/editor/common/config/editorOptions';
@@ -264,12 +265,16 @@ export class NotebookProviderInfoStore extends Disposable {
 }
 
 export class NotebookOutputRendererInfoStore {
-	private readonly contributedRenderers = new Map<string, NotebookOutputRendererInfo>();
+	private readonly contributedRenderers = new Map</* rendererId */ string, NotebookOutputRendererInfo>();
 	private readonly preferredMimetypeMemento: Memento;
-	private readonly preferredMimetype = new Lazy(() => this.preferredMimetypeMemento.getMemento(StorageScope.WORKSPACE, StorageTarget.USER));
+	private readonly preferredMimetype = new Lazy<{ [notebookType: string]: { [mimeType: string]: /* rendererId */ string } }>(
+		() => this.preferredMimetypeMemento.getMemento(StorageScope.WORKSPACE, StorageTarget.USER));
 
-	constructor(@IStorageService storageService: IStorageService) {
-		this.preferredMimetypeMemento = new Memento('workbench.editor.notebook.preferredRenderer', storageService);
+	constructor(
+		@IStorageService storageService: IStorageService,
+		@IWorkspaceTrustManagementService private readonly workspaceTrustManagementService: IWorkspaceTrustManagementService,
+	) {
+		this.preferredMimetypeMemento = new Memento('workbench.editor.notebook.preferredRenderer2', storageService);
 	}
 
 	clear() {
@@ -292,25 +297,65 @@ export class NotebookOutputRendererInfoStore {
 	}
 
 	/** Update and remember the preferred renderer for the given mimetype in this workspace */
-	setPreferred(mimeType: string, rendererId: string) {
-		this.preferredMimetype.getValue()[mimeType] = rendererId;
+	setPreferred(notebookProviderInfo: NotebookProviderInfo, mimeType: string, rendererId: string) {
+		const mementoObj = this.preferredMimetype.getValue();
+		const forNotebook = mementoObj[notebookProviderInfo.id];
+		if (forNotebook) {
+			forNotebook[mimeType] = rendererId;
+		} else {
+			mementoObj[notebookProviderInfo.id] = { [mimeType]: rendererId };
+		}
+
 		this.preferredMimetypeMemento.saveMemento();
 	}
 
-	getContributedRenderer(mimeType: string, kernelProvides: readonly string[] | undefined): NotebookOutputRendererInfo[] {
-		const preferred = this.preferredMimetype.getValue()[mimeType];
-		const possible = Array.from(this.contributedRenderers.values())
-			.map(renderer => ({
-				renderer,
-				score: kernelProvides === undefined
+	findBestRenderers(notebookProviderInfo: NotebookProviderInfo | undefined, mimeType: string, kernelProvides: readonly string[] | undefined): IOrderedMimeType[] {
+
+		const enum ReuseOrder {
+			PreviouslySelected = 1 << 8,
+			SameExtensionAsNotebook = 2 << 8,
+			BuiltIn = 3 << 8,
+			OtherRenderer = 4 << 8
+		}
+
+		const preferred = notebookProviderInfo && this.preferredMimetype.getValue()[notebookProviderInfo.id]?.[mimeType];
+		const renderers: { ordered: IOrderedMimeType, score: number }[] = Array.from(this.contributedRenderers.values())
+			.map(renderer => {
+				const ownScore = kernelProvides === undefined
 					? renderer.matchesWithoutKernel(mimeType)
-					: renderer.matches(mimeType, kernelProvides),
-			}))
-			.sort((a, b) => a.score - b.score)
-			.filter(r => r.score !== NotebookRendererMatch.Never)
-			.map(r => r.renderer);
+					: renderer.matches(mimeType, kernelProvides);
+
+				if (ownScore === NotebookRendererMatch.Never) {
+					return undefined;
+				}
+
+				const reuseScore = preferred === renderer.id
+					? ReuseOrder.PreviouslySelected
+					: renderer.extensionId.value === notebookProviderInfo?.extension?.value
+						? ReuseOrder.SameExtensionAsNotebook
+						: ReuseOrder.BuiltIn;
+				return {
+					ordered: { mimeType, rendererId: renderer.id, isTrusted: true },
+					score: reuseScore | ownScore,
+				};
+			}).filter(isDefined);
+
+		if (mimeTypeSupportedByCore(mimeType)) {
+			renderers.push({
+				score: ReuseOrder.BuiltIn << 8,
+				ordered: {
+					mimeType,
+					rendererId: BUILTIN_RENDERER_ID,
+					isTrusted: mimeTypeIsAlwaysSecure(mimeType) || this.workspaceTrustManagementService.isWorkspaceTrusted()
+				}
+			});
+		}
+
+		if (renderers.length === 0) {
+			return [{ mimeType, rendererId: RENDERER_NOT_AVAILABLE, isTrusted: true }];
+		}
 
-		return preferred ? possible.sort((a, b) => (a.id === preferred ? -1 : 0) + (b.id === preferred ? 1 : 0)) : possible;
+		return renderers.sort((a, b) => a.score - b.score).map(r => r.ordered);
 	}
 }
 
@@ -376,7 +421,6 @@ export class NotebookService extends Disposable implements INotebookService {
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@ICodeEditorService private readonly _codeEditorService: ICodeEditorService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
-		@IWorkspaceTrustManagementService private readonly workspaceTrustManagementService: IWorkspaceTrustManagementService,
 	) {
 		super();
 
@@ -569,8 +613,11 @@ export class NotebookService extends Disposable implements INotebookService {
 		return this._notebookRenderersInfoStore.get(rendererId);
 	}
 
-	updateMimePreferredRenderer(mimeType: string, rendererId: string): void {
-		this._notebookRenderersInfoStore.setPreferred(mimeType, rendererId);
+	updateMimePreferredRenderer(viewType: string, mimeType: string, rendererId: string): void {
+		const info = this.notebookProviderInfoStore.get(viewType);
+		if (info) {
+			this._notebookRenderersInfoStore.setPreferred(info, mimeType, rendererId);
+		}
 	}
 
 	getRenderers(): INotebookRendererInfo[] {
@@ -620,58 +667,9 @@ export class NotebookService extends Disposable implements INotebookService {
 
 		const coreDisplayOrder = this._displayOrder;
 		const sorted = sortMimeTypes(mimeTypes, coreDisplayOrder?.userOrder ?? [], coreDisplayOrder?.defaultOrder ?? []);
+		const notebookProviderInfo = this.notebookProviderInfoStore.get(textModel.viewType);
 
-		const orderMimeTypes: IOrderedMimeType[] = [];
-
-		sorted.forEach(mimeType => {
-			const handlers = this._findBestMatchedRenderer(mimeType, kernelProvides);
-
-			if (handlers.length) {
-				const handler = handlers[0];
-
-				orderMimeTypes.push({
-					mimeType: mimeType,
-					rendererId: handler.id,
-					isTrusted: true
-				});
-
-				for (let i = 1; i < handlers.length; i++) {
-					orderMimeTypes.push({
-						mimeType: mimeType,
-						rendererId: handlers[i].id,
-						isTrusted: true
-					});
-				}
-
-				if (mimeTypeSupportedByCore(mimeType)) {
-					orderMimeTypes.push({
-						mimeType: mimeType,
-						rendererId: BUILTIN_RENDERER_ID,
-						isTrusted: mimeTypeIsAlwaysSecure(mimeType) || this.workspaceTrustManagementService.isWorkspaceTrusted()
-					});
-				}
-			} else {
-				if (mimeTypeSupportedByCore(mimeType)) {
-					orderMimeTypes.push({
-						mimeType: mimeType,
-						rendererId: BUILTIN_RENDERER_ID,
-						isTrusted: mimeTypeIsAlwaysSecure(mimeType) || this.workspaceTrustManagementService.isWorkspaceTrusted()
-					});
-				} else {
-					orderMimeTypes.push({
-						mimeType: mimeType,
-						rendererId: RENDERER_NOT_AVAILABLE,
-						isTrusted: true
-					});
-				}
-			}
-		});
-
-		return orderMimeTypes;
-	}
-
-	private _findBestMatchedRenderer(mimeType: string, kernelProvides: readonly string[] | undefined): readonly NotebookOutputRendererInfo[] {
-		return this._notebookRenderersInfoStore.getContributedRenderer(mimeType, kernelProvides);
+		return sorted.flatMap(mimeType => this._notebookRenderersInfoStore.findBestRenderers(notebookProviderInfo, mimeType, kernelProvides));
 	}
 
 	getContributedNotebookTypes(resource?: URI): readonly NotebookProviderInfo[] {

diff --git a/src/vs/workbench/contrib/notebook/browser/view/renderers/cellOutput.ts b/src/vs/workbench/contrib/notebook/browser/view/renderers/cellOutput.ts
@@ -487,7 +487,7 @@ export class CellOutputElement extends Disposable {
 		this.viewCell.updateOutputMinHeight(this.viewCell.layoutInfo.outputTotalHeight);
 
 		const { mimeType, rendererId } = mimeTypes[pick.index];
-		this.notebookService.updateMimePreferredRenderer(mimeType, rendererId);
+		this.notebookService.updateMimePreferredRenderer(notebookTextModel.viewType, mimeType, rendererId);
 		this.render(nextElement as HTMLElement);
 		this._validateFinalOutputHeight(false);
 		this._relayoutCell();

diff --git a/src/vs/workbench/contrib/notebook/common/notebookCommon.ts b/src/vs/workbench/contrib/notebook/common/notebookCommon.ts
@@ -109,8 +109,6 @@ export interface TransientOptions {
 	transientDocumentMetadata: TransientDocumentMetadata;
 }
 
-
-
 /** Note: enum values are used for sorting */
 export const enum NotebookRendererMatch {
 	/** Renderer has a hard dependency on an available kernel */

diff --git a/src/vs/workbench/contrib/notebook/common/notebookOutputRenderer.ts b/src/vs/workbench/contrib/notebook/common/notebookOutputRenderer.ts
@@ -81,11 +81,11 @@ export class NotebookOutputRendererInfo implements INotebookRendererInfo {
 		this.messaging = descriptor.requiresMessaging ?? RendererMessagingSpec.Never;
 	}
 
-	get dependencies(): string[] {
+	public get dependencies(): string[] {
 		return this.hardDependencies.values();
 	}
 
-	matchesWithoutKernel(mimeType: string) {
+	public matchesWithoutKernel(mimeType: string) {
 		if (!this.matchesMimeTypeOnly(mimeType)) {
 			return NotebookRendererMatch.Never;
 		}
@@ -101,7 +101,7 @@ export class NotebookOutputRendererInfo implements INotebookRendererInfo {
 		return NotebookRendererMatch.Pure;
 	}
 
-	matches(mimeType: string, kernelProvides: ReadonlyArray<string>) {
+	public matches(mimeType: string, kernelProvides: ReadonlyArray<string>) {
 		if (!this.matchesMimeTypeOnly(mimeType)) {
 			return NotebookRendererMatch.Never;
 		}

diff --git a/src/vs/workbench/contrib/notebook/common/notebookService.ts b/src/vs/workbench/contrib/notebook/common/notebookService.ts
@@ -76,7 +76,7 @@ export interface INotebookService {
 	getRenderers(): INotebookRendererInfo[];
 
 	/** Updates the preferred renderer for the given mimetype in the workspace. */
-	updateMimePreferredRenderer(mimeType: string, rendererId: string): void;
+	updateMimePreferredRenderer(viewType: string, mimeType: string, rendererId: string): void;
 
 	createNotebookTextModel(viewType: string, uri: URI, data: NotebookData, transientOptions: TransientOptions): NotebookTextModel;
 	getNotebookTextModel(uri: URI): NotebookTextModel | undefined;