Diffing editors/documents/visibleEditors/activeEditor.

src/vs/workbench/api/browser/mainThreadNotebook.ts

@@ -60,6 +60,22 @@ export class MainThreadNotebookDocument extends Disposable {
 }
 
 class DocumentAndEditorState {
+	static ofSets<T>(before: Set<T>, after: Set<T>): { removed: T[], added: T[] } {
+		const removed: T[] = [];
+		const added: T[] = [];
+		before.forEach(element => {
+			if (!after.has(element)) {
+				removed.push(element);
+			}
+		});
+		after.forEach(element => {
+			if (!before.has(element)) {
+				added.push(element);
+			}
+		});
+		return { removed, added };
+	}
+
 	static ofMaps<K, V>(before: Map<K, V>, after: Map<K, V>): { removed: V[], added: V[] } {
 		const removed: V[] = [];
 		const added: V[] = [];
@@ -86,15 +102,16 @@ class DocumentAndEditorState {
 
 			return {
 				addedDocuments: [],
-				addedEditors: apiEditors
+				addedEditors: apiEditors,
+				visibleEditors: [...after.visibleEditors].map(editor => editor[0])
 			};
 		}
-		// const documentDelta = delta.ofSets(before.documents, after.documents);
+		const documentDelta = DocumentAndEditorState.ofSets(before.documents, after.documents);
 		const editorDelta = DocumentAndEditorState.ofMaps(before.textEditors, after.textEditors);
 		const addedAPIEditors = editorDelta.added.map(add => ({
 			id: add.getId(),
 			documentUri: add.uri!,
-			selections: add.textModel!.selections
+			selections: add.textModel!.selections || []
 		}));
 
 		const removedAPIEditors = editorDelta.removed.map(removed => removed.getId());
@@ -102,17 +119,46 @@ class DocumentAndEditorState {
 		// const oldActiveEditor = before.activeEditor !== after.activeEditor ? before.activeEditor : undefined;
 		const newActiveEditor = before.activeEditor !== after.activeEditor ? after.activeEditor : undefined;
 
+		const visibleEditorDelta = DocumentAndEditorState.ofMaps(before.visibleEditors, after.visibleEditors);
+
 		return {
+			addedDocuments: documentDelta.added.map(e => {
+				return {
+					viewType: e.viewType,
+					handle: e.handle,
+					uri: e.uri,
+					metadata: e.metadata,
+					versionId: e.versionId,
+					cells: e.cells.map(cell => ({
+						handle: cell.handle,
+						uri: cell.uri,
+						source: cell.textBuffer.getLinesContent(),
+						language: cell.language,
+						cellKind: cell.cellKind,
+						outputs: cell.outputs,
+						metadata: cell.metadata
+					})),
+					// attachedEditor: editorId ? {
+					// 	id: editorId,
+					// 	selections: document.textModel.selections
+					// } : undefined
+				};
+			}),
+			removedDocuments: documentDelta.removed.map(e => e.uri),
 			addedEditors: addedAPIEditors,
 			removedEditors: removedAPIEditors,
-			newActiveEditor: newActiveEditor
+			newActiveEditor: newActiveEditor,
+			visibleEditors: visibleEditorDelta.added.length === 0 && visibleEditorDelta.removed.length === 0
+				? undefined
+				: [...after.visibleEditors].map(editor => editor[0])
 		};
 	}
 
 	constructor(
-		readonly documents: Set<URI>,
+		readonly documents: Set<NotebookTextModel>,
 		readonly textEditors: Map<string, IEditor>,
 		readonly activeEditor: string | null | undefined,
+		readonly visibleEditors: Map<string, IEditor>
 	) {
 		//
 	}
@@ -150,27 +196,38 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		return false;
 	}
 
+	private _emitDelta(delta: INotebookDocumentsAndEditorsDelta) {
+		this._proxy.$acceptDocumentAndEditorsDelta(delta);
+	}
+
 	registerListeners() {
 		this._notebookService.listNotebookEditors().forEach((e) => {
 			this._addNotebookEditor(e);
 		});
 
 		this._register(this._notebookService.onDidChangeActiveEditor(e => {
-			this._proxy.$acceptDocumentAndEditorsDelta({
-				newActiveEditor: e
-			});
+			this._updateState();
 		}));
 
 		this._register(this._notebookService.onDidChangeVisibleEditors(e => {
-			this._proxy.$acceptDocumentAndEditorsDelta({
-				visibleEditors: e
-			});
+			if (this._notebookProviders.size > 0) {
+				if (!this._currentState) {
+					// no current state means we didn't even create editors in ext host yet.
+					return;
+				}
+
+				// we can't simply update visibleEditors as we need to check if we should create editors first.
+				this._updateState();
+			}
 		}));
 
 		this._register(this._notebookService.onNotebookEditorAdd(editor => {
 			this._addNotebookEditor(editor);
 		}));
 
+		this._register(this._notebookService.onNotebookEditorsRemove(editors => {
+			this._removeNotebookEditor(editors);
+		}));
 
 		this._register(this._notebookService.onNotebookDocumentRemove(() => {
 			this._updateState();
@@ -218,39 +275,60 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 		this._updateState(notebookEditor);
 	}
 
-	private _removeNotebookEditor(e: IEditor) {
-		const sub = this._toDisposeOnEditorRemove.get(e.getId());
-		if (sub) {
-			this._toDisposeOnEditorRemove.delete(e.getId());
-			sub.dispose();
-			this._updateState();
-		}
+	private _removeNotebookEditor(editors: IEditor[]) {
+		editors.forEach(e => {
+			const sub = this._toDisposeOnEditorRemove.get(e.getId());
+			if (sub) {
+				this._toDisposeOnEditorRemove.delete(e.getId());
+				sub.dispose();
+			}
+		});
+
+		this._updateState();
 	}
 
 	private async _updateState(focusedNotebookEditor?: IEditor) {
-		const documents = new Set<URI>();
-		this._notebookService.listNotebookDocuments().forEach(document => {
-			documents.add(document.uri);
-		});
-
-		const editors = new Map<string, IEditor>();
 		let activeEditor: string | null = null;
 
-		for (const editor of this._notebookService.listNotebookEditors()) {
+		const activeEditorPane = this.editorService.activeEditorPane as any | undefined;
+		if (activeEditorPane?.isNotebookEditor) {
+			const notebookEditor = (activeEditorPane.getControl() as INotebookEditor);
+			activeEditor = notebookEditor && notebookEditor.hasModel() ? notebookEditor!.getId() : null;
+		}
+
+		const documentEditorsMap = new Map<string, IEditor>();
+
+		const editors = new Map<string, IEditor>();
+		this._notebookService.listNotebookEditors().forEach(editor => {
 			if (editor.hasModel()) {
 				editors.set(editor.getId(), editor);
-				if (editor.hasFocus()) {
-					activeEditor = editor.getId();
+				documentEditorsMap.set(editor.textModel!.uri.toString(), editor);
+			}
+		});
+
+		const visibleEditorsMap = new Map<string, IEditor>();
+		this.editorService.visibleEditorPanes.forEach(editor => {
+			if ((editor as any).isNotebookEditor) {
+				const nbEditorWidget = (editor as any).getControl() as INotebookEditor;
+				if (nbEditorWidget && editors.has(nbEditorWidget.getId())) {
+					visibleEditorsMap.set(nbEditorWidget.getId(), nbEditorWidget);
 				}
 			}
-		}
+		});
+
+		const documents = new Set<NotebookTextModel>();
+		this._notebookService.listNotebookDocuments().forEach(document => {
+			if (documentEditorsMap.has(document.uri.toString())) {
+				documents.add(document);
+			}
+		});
 
-		if (!activeEditor && focusedNotebookEditor) {
+		if (!activeEditor && focusedNotebookEditor && focusedNotebookEditor.hasModel()) {
 			activeEditor = focusedNotebookEditor.getId();
 		}
 
 		// editors always have view model attached, which means there is already a document in exthost.
-		const newState = new DocumentAndEditorState(documents, editors, activeEditor);
+		const newState = new DocumentAndEditorState(documents, editors, activeEditor, visibleEditorsMap);
 		const delta = DocumentAndEditorState.compute(this._currentState, newState);
 		// const isEmptyChange = (!delta.addedDocuments || delta.addedDocuments.length === 0)
 		// 	&& (!delta.removedDocuments || delta.removedDocuments.length === 0)
@@ -260,7 +338,7 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 
 		// if (!isEmptyChange) {
 		this._currentState = newState;
-		await this._proxy.$acceptDocumentAndEditorsDelta(delta);
+		await this._emitDelta(delta);
 		// }
 	}
 
@@ -502,6 +580,7 @@ export class MainThreadNotebookController implements IMainNotebookController {
 			return;
 		}
 
+		// TODO@rebornix, remove cell should use emitDelta as well to ensure document/editor events are sent together
 		await this._proxy.$acceptDocumentAndEditorsDelta({ removedDocuments: [notebook.uri] });
 		document.dispose();
 		this._mapping.delete(URI.from(notebook.uri).toString());

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

@@ -103,16 +103,18 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 	private readonly _notebookKernels = new Map<string, INotebookKernelInfo>();
 	notebookProviderInfoStore: NotebookProviderInfoStore = new NotebookProviderInfoStore();
 	notebookRenderersInfoStore: NotebookOutputRendererInfoStore = new NotebookOutputRendererInfoStore();
-	private readonly _models: { [modelId: string]: ModelData; };
+	private readonly _models = new Map<string, ModelData>();
 	private _onDidChangeActiveEditor = new Emitter<string | null>();
 	onDidChangeActiveEditor: Event<string | null> = this._onDidChangeActiveEditor.event;
 	private _onDidChangeVisibleEditors = new Emitter<string[]>();
 	onDidChangeVisibleEditors: Event<string[]> = this._onDidChangeVisibleEditors.event;
 	private readonly _onNotebookEditorAdd: Emitter<INotebookEditor> = this._register(new Emitter<INotebookEditor>());
 	public readonly onNotebookEditorAdd: Event<INotebookEditor> = this._onNotebookEditorAdd.event;
-	private readonly _onNotebookEditorRemove: Emitter<INotebookEditor> = this._register(new Emitter<INotebookEditor>());
-	public readonly onNotebookEditorRemove: Event<INotebookEditor> = this._onNotebookEditorRemove.event;
-	private readonly _notebookEditors: { [editorId: string]: INotebookEditor; };
+	private readonly _onNotebookEditorsRemove: Emitter<INotebookEditor[]> = this._register(new Emitter<INotebookEditor[]>());
+	public readonly onNotebookEditorsRemove: Event<INotebookEditor[]> = this._onNotebookEditorsRemove.event;
+	private readonly _onNotebookDocumentRemove: Emitter<URI[]> = this._register(new Emitter<URI[]>());
+	public readonly onNotebookDocumentRemove: Event<URI[]> = this._onNotebookDocumentRemove.event;
+	private readonly _notebookEditors = new Map<string, INotebookEditor>();
 
 	private readonly _onDidChangeViewTypes = new Emitter<void>();
 	onDidChangeViewTypes: Event<void> = this._onDidChangeViewTypes.event;
@@ -133,8 +135,6 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 	) {
 		super();
 
-		this._models = {};
-		this._notebookEditors = Object.create(null);
 		this.modelManager = this.instantiationService.createInstance(NotebookEditorModelManager);
 
 		notebookProviderExtensionPoint.setHandler((extensions) => {
@@ -300,9 +300,9 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		const modelId = MODEL_ID(uri);
 		const modelData = new ModelData(
 			notebookModel,
-			(model) => this._onWillDispose(model),
+			(model) => this._onWillDisposeDocument(model),
 		);
-		this._models[modelId] = modelData;
+		this._models.set(modelId, modelData);
 		return modelData.model;
 	}
 
@@ -319,10 +319,10 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		const modelId = MODEL_ID(uri);
 		const modelData = new ModelData(
 			notebookModel!,
-			(model) => this._onWillDispose(model),
+			(model) => this._onWillDisposeDocument(model),
 		);
 
-		this._models[modelId] = modelData;
+		this._models.set(modelId, modelData);
 		return modelData.model;
 	}
 
@@ -610,30 +610,37 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 	}
 
 	removeNotebookEditor(editor: INotebookEditor) {
-		if (delete this._notebookEditors[editor.getId()]) {
-			this._onNotebookEditorRemove.fire(editor);
+		let editorCache = this._notebookEditors.get(editor.getId());
+
+		if (editorCache) {
+			this._notebookEditors.delete(editor.getId());
+			this._onNotebookEditorsRemove.fire([editor]);
 		}
 	}
 
 	addNotebookEditor(editor: INotebookEditor) {
-		this._notebookEditors[editor.getId()] = editor;
+		this._notebookEditors.set(editor.getId(), editor);
 		this._onNotebookEditorAdd.fire(editor);
 	}
 
 	listNotebookEditors(): INotebookEditor[] {
-		return Object.keys(this._notebookEditors).map(id => this._notebookEditors[id]);
+		return [...this._notebookEditors].map(e => e[1]);
+	}
+
+	listVisibleNotebookEditors(): INotebookEditor[] {
+		return this.editorService.visibleEditorPanes
+			.filter(pane => (pane as any).isNotebookEditor)
+			.map(pane => pane.getControl() as INotebookEditor)
+			.filter(editor => !!editor)
+			.filter(editor => this._notebookEditors.has(editor.getId()));
 	}
 
 	listNotebookDocuments(): NotebookTextModel[] {
-		return Object.keys(this._models).map(id => this._models[id].model);
+		return [...this._models].map(e => e[1].model);
 	}
 
 	destoryNotebookDocument(viewType: string, notebook: INotebookTextModel): void {
-		let provider = this._notebookProviders.get(viewType);
-
-		if (provider) {
-			provider.controller.removeNotebookDocument(notebook);
-		}
+		this._onWillDisposeDocument(notebook);
 	}
 
 	updateActiveNotebookEditor(editor: INotebookEditor | null) {
@@ -641,7 +648,8 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 	}
 
 	updateVisibleNotebookEditor(editors: string[]) {
-		this._onDidChangeVisibleEditors.fire(editors);
+		const alreadyCreated = editors.filter(editorId => this._notebookEditors.has(editorId));
+		this._onDidChangeVisibleEditors.fire(alreadyCreated);
 	}
 
 	setToCopy(items: NotebookCellTextModel[]) {
@@ -680,13 +688,33 @@ export class NotebookService extends Disposable implements INotebookService, ICu
 		}
 	}
 
-	private _onWillDispose(model: INotebookTextModel): void {
+	private _onWillDisposeDocument(model: INotebookTextModel): void {
 		let modelId = MODEL_ID(model.uri);
-		let modelData = this._models[modelId];
 
-		delete this._models[modelId];
-		modelData?.dispose();
+		let modelData = this._models.get(modelId);
+		this._models.delete(modelId);
+
+		if (modelData) {
+			// delete editors and documents
+			const willRemovedEditors: INotebookEditor[] = [];
+			this._notebookEditors.forEach(editor => {
+				if (editor.textModel === modelData!.model) {
+					willRemovedEditors.push(editor);
+				}
+			});
+
+			willRemovedEditors.forEach(e => this._notebookEditors.delete(e.getId()));
+
+			let provider = this._notebookProviders.get(modelData!.model.viewType);
 
-		// this._onModelRemoved.fire(model);
+			if (provider) {
+				provider.controller.removeNotebookDocument(modelData!.model);
+			}
+
+
+			this._onNotebookEditorsRemove.fire(willRemovedEditors.map(e => e));
+			this._onNotebookDocumentRemove.fire([modelData.model.uri]);
+			modelData?.dispose();
+		}
 	}
 }

src/vs/workbench/contrib/notebook/common/notebookService.ts

@@ -34,7 +34,8 @@ export interface INotebookService {
 	onDidChangeActiveEditor: Event<string | null>;
 	onDidChangeVisibleEditors: Event<string[]>;
 	onNotebookEditorAdd: Event<IEditor>;
-	onNotebookEditorRemove: Event<IEditor>;
+	onNotebookEditorsRemove: Event<IEditor[]>;
+	onNotebookDocumentRemove: Event<URI[]>;
 	onDidChangeKernels: Event<void>;
 	registerNotebookController(viewType: string, extensionData: NotebookExtensionDescription, controller: IMainNotebookController): void;
 	unregisterNotebookProvider(viewType: string): void;
@@ -69,6 +70,7 @@ export interface INotebookService {
 	addNotebookEditor(editor: IEditor): void;
 	removeNotebookEditor(editor: IEditor): void;
 	listNotebookEditors(): readonly IEditor[];
+	listVisibleNotebookEditors(): readonly IEditor[];
 	listNotebookDocuments(): readonly NotebookTextModel[];
 
 }