borrow notebook editors from a notebook widget service. let the service manage...

borrow notebook editors from a notebook widget service. let the service manage editor move and deal with group changes

src/vs/workbench/contrib/notebook/browser/notebook.contribution.ts

@@ -32,7 +32,7 @@ import { INotebookService } from 'vs/workbench/contrib/notebook/common/notebookS
 import { NotebookService } from 'vs/workbench/contrib/notebook/browser/notebookServiceImpl';
 import { CellKind, CellUri, NotebookDocumentBackupData, NotebookEditorPriority } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { NotebookProviderInfo } from 'vs/workbench/contrib/notebook/common/notebookProvider';
-import { IEditorGroup, OpenEditorContext, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
+import { IEditorGroup } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { IEditorService, IOpenEditorOverride } from 'vs/workbench/services/editor/common/editorService';
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { CustomEditorsAssociations, customEditorsAssociationsSettingId } from 'vs/workbench/services/editor/common/editorAssociationsSetting';
@@ -40,7 +40,6 @@ import { CustomEditorInfo } from 'vs/workbench/contrib/customEditor/common/custo
 import { NotebookEditorOptions } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
 import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
-import { NotebookRegistry } from 'vs/workbench/contrib/notebook/browser/notebookRegistry';
 import { INotebookEditorModelResolverService, NotebookModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
 
 // Editor Contribution
@@ -146,7 +145,6 @@ export class NotebookContribution extends Disposable implements IWorkbenchContri
 		@IInstantiationService private readonly instantiationService: IInstantiationService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
 		@IUndoRedoService undoRedoService: IUndoRedoService,
-		@IEditorGroupsService private readonly editorGroupsService: IEditorGroupsService,
 	) {
 		super();
 
@@ -189,24 +187,6 @@ export class NotebookContribution extends Disposable implements IWorkbenchContri
 			}
 		}));
 
-		// HACK
-		// we use the open override to spy on tab movements because that's the only
-		// way to do that...
-		this._register(this.editorService.overrideOpenEditor({
-			open: (input, _options, group, context) => {
-				if (input instanceof NotebookEditorInput && context === OpenEditorContext.MOVE_EDITOR) {
-					// when moving a notebook editor we release it from its current tab and we
-					// "place" it into its future slot so that the editor can pick it up from there
-					const widgetRef = NotebookRegistry.getNotebookEditorWidget(input.resource, this.editorGroupsService.activeGroup);
-					if (widgetRef) {
-						NotebookRegistry.releaseNotebookEditorWidget(input.resource, this.editorGroupsService.activeGroup);
-						NotebookRegistry.claimNotebookEditorWidget(input.resource, group, widgetRef);
-					}
-				}
-				return undefined;
-			}
-		}));
-
 		this._register(this.editorService.onDidVisibleEditorsChange(() => {
 			const visibleNotebookEditors = editorService.visibleEditorPanes
 				.filter(pane => (pane as any).isNotebookEditor)

src/vs/workbench/contrib/notebook/browser/notebookEditor.ts

@@ -17,11 +17,12 @@ import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/browser/noteb
 import { INotebookEditorViewState, NotebookViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
 import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
-import { NotebookRegistry } from 'vs/workbench/contrib/notebook/browser/notebookRegistry';
 import { EditorPart } from 'vs/workbench/browser/parts/editor/editorPart';
 import { INotificationService, Severity } from 'vs/platform/notification/common/notification';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IEditorOptions, ITextEditorOptions } from 'vs/platform/editor/common/editor';
+import { INotebookEditorWidgetService, IBorrowValue } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidgetService';
+import { localize } from 'vs/nls';
 
 const NOTEBOOK_EDITOR_VIEW_STATE_PREFERENCE_KEY = 'NotebookEditorViewState';
 
@@ -29,7 +30,7 @@ export class NotebookEditor extends BaseEditor {
 	static readonly ID: string = 'workbench.editor.notebook';
 	private editorMemento: IEditorMemento<INotebookEditorViewState>;
 	private readonly groupListener = this._register(new MutableDisposable());
-	private _widget?: NotebookEditorWidget;
+	private _widget: IBorrowValue<NotebookEditorWidget> = { value: undefined };
 	private _rootElement!: HTMLElement;
 	private dimension: DOM.Dimension | null = null;
 	private _widgetDisposableStore: DisposableStore = new DisposableStore();
@@ -43,7 +44,8 @@ export class NotebookEditor extends BaseEditor {
 		@IStorageService storageService: IStorageService,
 		@IEditorService private readonly editorService: IEditorService,
 		@IEditorGroupsService private readonly editorGroupService: IEditorGroupsService,
-		@INotificationService private readonly notificationService: INotificationService
+		@INotificationService private readonly notificationService: INotificationService,
+		@INotebookEditorWidgetService private readonly notebookWidgetService: INotebookEditorWidgetService,
 	) {
 		super(NotebookEditor.ID, telemetryService, themeService, storageService);
 		this.editorMemento = this.getEditorMemento<INotebookEditorViewState>(editorGroupService, NOTEBOOK_EDITOR_VIEW_STATE_PREFERENCE_KEY);
@@ -54,14 +56,14 @@ export class NotebookEditor extends BaseEditor {
 
 
 	set viewModel(newModel: NotebookViewModel | undefined) {
-		if (this._widget) {
-			this._widget.viewModel = newModel;
+		if (this._widget.value) {
+			this._widget.value.viewModel = newModel;
 			this._onDidChangeModel.fire();
 		}
 	}
 
 	get viewModel() {
-		return this._widget?.viewModel;
+		return this._widget.value?.viewModel;
 	}
 
 	get minimumWidth(): number { return 375; }
@@ -83,28 +85,26 @@ export class NotebookEditor extends BaseEditor {
 		this._rootElement = DOM.append(parent, DOM.$('.notebook-editor'));
 
 		// this._widget.createEditor();
-		this._register(this.onDidFocus(() => this._widget?.updateEditorFocus()));
-		this._register(this.onDidBlur(() => this._widget?.updateEditorFocus()));
+		this._register(this.onDidFocus(() => this._widget.value?.updateEditorFocus()));
+		this._register(this.onDidBlur(() => this._widget.value?.updateEditorFocus()));
 	}
 
 	getDomNode() {
 		return this._rootElement;
 	}
 
-	getControl() {
-		return this._widget;
+	getControl(): NotebookEditorWidget | undefined {
+		return this._widget.value;
 	}
 
 	onWillHide() {
-		if (this.input && this.input instanceof NotebookEditorInput && !this.input.isDisposed()) {
+		if (this.input instanceof NotebookEditorInput) {
 			this.saveEditorViewState(this.input);
 		}
-
-		if (this.input && NotebookRegistry.getNotebookEditorWidget(this.input.resource!, this.group!) === this._widget) {
+		if (this.input && this._widget.value) {
 			// the widget is not transfered to other editor inputs
-			this._widget?.onWillHide();
+			this._widget.value.onWillHide();
 		}
-
 		super.onWillHide();
 	}
 
@@ -126,7 +126,7 @@ export class NotebookEditor extends BaseEditor {
 
 	focus() {
 		super.focus();
-		this._widget?.focus();
+		this._widget.value?.focus();
 	}
 
 	async setInput(input: NotebookEditorInput, options: EditorOptions | undefined, token: CancellationToken): Promise<void> {
@@ -134,63 +134,34 @@ export class NotebookEditor extends BaseEditor {
 		const group = this.group!;
 
 		if (this.input instanceof NotebookEditorInput) {
-			if (!this.input.isDisposed()) {
-				// set a new input, let's hide previous input
-				this.saveEditorViewState(this.input as NotebookEditorInput);
-				this._widget?.onWillHide();
-			}
+			// set a new input, let's hide previous input
+			this.saveEditorViewState(this.input as NotebookEditorInput);
 		}
 
 		await super.setInput(input, options, token);
 
-
-		// input attached
-		Event.once(input.onDispose)(() => {
-			// todo@jrieken
-			// there is no more input for a resource and that is used as signal to clean up
-			// all widget that might still be around
-			for (let widget of NotebookRegistry.releaseAllNotebookEditorWidgets(input.resource)) {
-				widget.onWillHide();
-				widget.getDomNode().remove();
-				widget.dispose();
-			}
-		});
-
 		this._widgetDisposableStore.clear();
 
-		const existingEditorWidgetForInput = NotebookRegistry.getNotebookEditorWidget(input.resource, group);
-		if (existingEditorWidgetForInput) {
-			// hide previous widget
-			if (NotebookRegistry.getNotebookEditorWidget(this.input!.resource!, this.group!) === this._widget) {
-				// the widet is not transfered to other editor inputs
-				this._widget?.onWillHide();
-			}
-
-			// previous widget is then detached
-			// set the new one
-			this._widget = existingEditorWidgetForInput;
-			NotebookRegistry.claimNotebookEditorWidget(input.resource, group, this._widget);
-		} else {
-			// hide current widget
-			this._widget?.onWillHide();
-			// create a new widget
-			this._widget = this.instantiationService.createInstance(NotebookEditorWidget);
-			this._widget.createEditor();
-			NotebookRegistry.claimNotebookEditorWidget(input.resource, group, this._widget);
+		// there currently is a widget which we still own so
+		// we need to hide it before getting a new widget
+		if (this._widget.value) {
+			this._widget.value.onWillHide();
 		}
 
+		this._widget = this.instantiationService.invokeFunction(this.notebookWidgetService.retrieveWidget, group, input);
+
 		if (this.dimension) {
-			this._widget.layout(this.dimension, this._rootElement);
+			this._widget.value!.layout(this.dimension, this._rootElement);
 		}
 
-		const model = await input.resolve(this._widget!.getId());
+		const model = await input.resolve(this._widget.value!.getId());
 
 		if (model === null) {
 			this.notificationService.prompt(
 				Severity.Error,
-				`Cannot open resource with notebook editor type '${input.viewType}', please check if you have the right extension installed or enabled.`,
+				localize('fail.noEditor', "Cannot open resource with notebook editor type '${input.viewType}', please check if you have the right extension installed or enabled."),
 				[{
-					label: 'Reopen file with VS Code standard text editor',
+					label: localize('fail.reOpen', "Reopen file with VS Code standard text editor"),
 					run: async () => {
 						const fileEditorInput = this.editorService.createEditorInput({ resource: input.resource, forceFile: true });
 						const textOptions: IEditorOptions | ITextEditorOptions = options ? { ...options, override: false } : { override: false };
@@ -203,26 +174,26 @@ export class NotebookEditor extends BaseEditor {
 
 		const viewState = this.loadTextEditorViewState(input);
 
-		await this._widget.setModel(model.notebook, viewState, options);
-		this._widgetDisposableStore.add(this._widget.onDidFocus(() => this._onDidFocusWidget.fire()));
+		await this._widget.value!.setModel(model.notebook, viewState, options);
+		this._widgetDisposableStore.add(this._widget.value!.onDidFocus(() => this._onDidFocusWidget.fire()));
 
 		if (this.editorGroupService instanceof EditorPart) {
-			this._widgetDisposableStore.add(this.editorGroupService.createEditorDropTarget(this._widget.getDomNode(), {
+			this._widgetDisposableStore.add(this.editorGroupService.createEditorDropTarget(this._widget.value!.getDomNode(), {
 				groupContainsPredicate: (group) => this.group?.id === group.group.id
 			}));
 		}
 	}
 
 	clearInput(): void {
-		const existingEditorWidgetForInput = NotebookRegistry.getNotebookEditorWidget(this.input!.resource!, this.group!);
-		existingEditorWidgetForInput?.onWillHide();
-		this._widget = undefined;
+		if (this._widget.value) {
+			this._widget.value.onWillHide();
+		}
 		super.clearInput();
 	}
 
 	private saveEditorViewState(input: NotebookEditorInput): void {
-		if (this.group && this._widget) {
-			const state = this._widget.getEditorViewState();
+		if (this.group && this._widget.value) {
+			const state = this._widget.value.getEditorViewState();
 			this.editorMemento.saveEditorState(this.group, input.resource, state);
 		}
 	}
@@ -240,11 +211,11 @@ export class NotebookEditor extends BaseEditor {
 		DOM.toggleClass(this._rootElement, 'narrow-width', dimension.width < 600);
 		this.dimension = dimension;
 
-		if (this._input === undefined || this._widget === undefined) {
+		if (!this._widget.value || !(this._input instanceof NotebookEditorInput)) {
 			return;
 		}
 
-		if (this._input.resource?.toString() !== this._widget?.viewModel?.uri.toString() && this._widget?.viewModel) {
+		if (this._input.resource.toString() !== this._widget.value.viewModel?.uri.toString() && this._widget.value?.viewModel) {
 			// input and widget mismatch
 			// this happens when
 			// 1. open document A, pin the document
@@ -254,7 +225,7 @@ export class NotebookEditor extends BaseEditor {
 			return;
 		}
 
-		this._widget?.layout(this.dimension, this._rootElement);
+		this._widget.value.layout(this.dimension, this._rootElement);
 	}
 
 	protected saveState(): void {

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -98,6 +98,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 	public get onDidFocus(): Event<any> { return this._onDidFocusWidget.event; }
 	private _cellContextKeyManager: CellContextKeyManager | null = null;
 	private _isVisible = false;
+	private readonly _uuid = generateUuid();
 
 	get isDisposed() {
 		return this._isDisposed;
@@ -130,7 +131,6 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditor
 		this.notebookService.addNotebookEditor(this);
 	}
 
-	private _uuid = generateUuid();
 	public getId(): string {
 		return this._uuid;
 	}

src/vs/workbench/contrib/notebook/browser/notebookEditorWidgetService.ts

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ResourceMap } from 'vs/base/common/map';
+import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
+import { DisposableStore, IDisposable } from 'vs/base/common/lifecycle';
+import { IEditorGroupsService, IEditorGroup, GroupChangeKind, OpenEditorContext } from 'vs/workbench/services/editor/common/editorGroupsService';
+import { IInstantiationService, createDecorator, ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
+import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
+import { registerSingleton } from 'vs/platform/instantiation/common/extensions';
+import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+
+export const INotebookEditorWidgetService = createDecorator<INotebookEditorWidgetService>('INotebookEditorWidgetService');
+
+export interface IBorrowValue<T> {
+	readonly value: T | undefined;
+}
+
+export interface INotebookEditorWidgetService {
+	_serviceBrand: undefined;
+	retrieveWidget(accessor: ServicesAccessor, group: IEditorGroup, input: NotebookEditorInput): IBorrowValue<NotebookEditorWidget>;
+}
+
+class NotebookEditorWidgetService implements INotebookEditorWidgetService {
+
+	readonly _serviceBrand: undefined;
+
+	private _tokenPool = 1;
+
+	private readonly _notebookWidgets = new Map<number, ResourceMap<{ widget: NotebookEditorWidget, token: number | undefined }>>();
+	private readonly _disposables = new DisposableStore();
+
+	constructor(
+		@IEditorGroupsService editorGroupService: IEditorGroupsService,
+		@IEditorService editorService: IEditorService,
+	) {
+
+		const groupListener = new Map<number, IDisposable>();
+		const onNewGroup = (group: IEditorGroup) => {
+			const { id } = group;
+			const listener = group.onDidGroupChange(e => {
+				const widgets = this._notebookWidgets.get(group.id);
+				if (!widgets || e.kind !== GroupChangeKind.EDITOR_CLOSE || !(e.editor instanceof NotebookEditorInput)) {
+					return;
+				}
+				const value = widgets.get(e.editor.resource);
+				if (!value) {
+					return;
+				}
+				value.token = undefined;
+				this._disposeWidget(value.widget);
+				widgets.delete(e.editor.resource);
+			});
+			groupListener.set(id, listener);
+		};
+		this._disposables.add(editorGroupService.onDidAddGroup(onNewGroup));
+		editorGroupService.groups.forEach(onNewGroup);
+
+		// group removed -> clean up listeners, clean up widgets
+		this._disposables.add(editorGroupService.onDidRemoveGroup(group => {
+			const listener = groupListener.get(group.id);
+			if (listener) {
+				listener.dispose();
+				groupListener.delete(group.id);
+			}
+			const widgets = this._notebookWidgets.get(group.id);
+			this._notebookWidgets.delete(group.id);
+			if (widgets) {
+				for (let value of widgets.values()) {
+					value.token = undefined;
+					this._disposeWidget(value.widget);
+				}
+			}
+		}));
+
+		// HACK
+		// we use the open override to spy on tab movements because that's the only
+		// way to do that...
+		this._disposables.add(editorService.overrideOpenEditor({
+			open: (input, _options, group, context) => {
+				if (input instanceof NotebookEditorInput && context === OpenEditorContext.MOVE_EDITOR) {
+					// when moving a notebook editor we release it from its current tab and we
+					// "place" it into its future slot so that the editor can pick it up from there
+					this._freeWidget(input, editorGroupService.activeGroup, group);
+				}
+				return undefined;
+			}
+		}));
+	}
+
+	private _disposeWidget(widget: NotebookEditorWidget): void {
+		widget.onWillHide();
+		widget.getDomNode().remove();
+		widget.dispose();
+	}
+
+	private _freeWidget(input: NotebookEditorInput, source: IEditorGroup, target: IEditorGroup): void {
+		const targetWidget = this._notebookWidgets.get(target.id)?.get(input.resource);
+		if (targetWidget) {
+			// not needed
+			return;
+		}
+
+		const widget = this._notebookWidgets.get(source.id)?.get(input.resource);
+		if (!widget) {
+			throw new Error('no widget at source group');
+		}
+		this._notebookWidgets.get(source.id)?.delete(input.resource);
+		widget.token = undefined;
+
+		let targetMap = this._notebookWidgets.get(target.id);
+		if (!targetMap) {
+			targetMap = new ResourceMap();
+			this._notebookWidgets.set(target.id, targetMap);
+		}
+		targetMap.set(input.resource, widget);
+	}
+
+	retrieveWidget(accessor: ServicesAccessor, group: IEditorGroup, input: NotebookEditorInput): IBorrowValue<NotebookEditorWidget> {
+
+		let value = this._notebookWidgets.get(group.id)?.get(input.resource);
+
+		if (!value) {
+			// NEW widget
+			const instantiationService = accessor.get(IInstantiationService);
+			const widget = instantiationService.createInstance(NotebookEditorWidget);
+			widget.createEditor();
+			const token = this._tokenPool++;
+			value = { widget, token };
+
+			let map = this._notebookWidgets.get(group.id);
+			if (!map) {
+				map = new ResourceMap();
+				this._notebookWidgets.set(group.id, map);
+			}
+			map.set(input.resource, value);
+
+		} else {
+			// reuse a widget which was either free'ed before or which
+			// is simply being reused...
+			value.token = this._tokenPool++;
+		}
+
+		return this._createBorrowValue(value.token!, value);
+	}
+
+	private _createBorrowValue(myToken: number, widget: { widget: NotebookEditorWidget, token: number | undefined }): IBorrowValue<NotebookEditorWidget> {
+		return {
+			get value() {
+				return widget.token === myToken ? widget.widget : undefined;
+			}
+		};
+	}
+}
+
+registerSingleton(INotebookEditorWidgetService, NotebookEditorWidgetService, true);

src/vs/workbench/contrib/notebook/browser/notebookRegistry.ts

@@ -6,10 +6,6 @@
 import { CellOutputKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { BrandedService, IConstructorSignature1 } from 'vs/platform/instantiation/common/instantiation';
 import { INotebookEditor, IOutputTransformContribution } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
-import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
-import { URI } from 'vs/base/common/uri';
-import { IEditorGroup } from 'vs/workbench/services/editor/common/editorGroupsService';
-import { ResourceMap } from 'vs/base/common/map';
 
 export type IOutputTransformCtor = IConstructorSignature1<INotebookEditor, IOutputTransformContribution>;
 
@@ -23,7 +19,6 @@ export interface IOutputTransformDescription {
 export const NotebookRegistry = new class NotebookRegistryImpl {
 
 	readonly outputTransforms: IOutputTransformDescription[] = [];
-	readonly notebookEditorWidgetOwnership = new ResourceMap<Map<number, NotebookEditorWidget>>();
 
 	registerOutputTransform<Services extends BrandedService[]>(id: string, kind: CellOutputKind, ctor: { new(editor: INotebookEditor, ...services: Services): IOutputTransformContribution }): void {
 		this.outputTransforms.push({ id: id, kind: kind, ctor: ctor as IOutputTransformCtor });
@@ -32,34 +27,4 @@ export const NotebookRegistry = new class NotebookRegistryImpl {
 	getOutputTransformContributions(): IOutputTransformDescription[] {
 		return this.outputTransforms.slice(0);
 	}
-
-	claimNotebookEditorWidget(notebook: URI, group: IEditorGroup, widget: NotebookEditorWidget) {
-		let map = this.notebookEditorWidgetOwnership.get(notebook);
-		if (!map) {
-			map = new Map();
-			this.notebookEditorWidgetOwnership.set(notebook, map);
-		}
-		map.set(group.id, widget);
-	}
-
-	releaseNotebookEditorWidget(notebook: URI, group: IEditorGroup) {
-		const map = this.notebookEditorWidgetOwnership.get(notebook);
-		if (!map) {
-			return;
-		}
-		map.delete(group.id);
-		if (map.size === 0) {
-			this.notebookEditorWidgetOwnership.delete(notebook);
-		}
-	}
-
-	getNotebookEditorWidget(notebook: URI, group: IEditorGroup) {
-		return this.notebookEditorWidgetOwnership.get(notebook)?.get(group.id);
-	}
-
-	releaseAllNotebookEditorWidgets(notebook: URI) {
-		let values = [...this.notebookEditorWidgetOwnership.get(notebook)?.values() ?? []];
-		this.notebookEditorWidgetOwnership.delete(notebook);
-		return values;
-	}
 };