Merge pull request #102760 from microsoft/joh/celldocs

Clean up and simplify cell text document world

extensions/vscode-notebook-tests/src/notebook.test.ts

@@ -131,6 +131,46 @@ suite('Notebook API tests', () => {
 		await firstDocumentClose;
 	});
 
+	test('notebook open/close, all cell-documents are ready', async function () {
+		const resource = await createRandomFile('', undefined, 'first', '.vsctestnb');
+
+		const p = getEventOncePromise(vscode.notebook.onDidOpenNotebookDocument).then(notebook => {
+			for (let cell of notebook.cells) {
+				const doc = vscode.workspace.textDocuments.find(doc => doc.uri.toString() === cell.uri.toString());
+				assert.ok(doc);
+				assert.strictEqual(doc === cell.document, true);
+				assert.strictEqual(doc?.languageId, cell.language);
+				assert.strictEqual(doc?.isDirty, false);
+				assert.strictEqual(doc?.isClosed, false);
+			}
+		});
+
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+		await p;
+		await vscode.commands.executeCommand('workbench.action.closeAllEditors');
+	});
+
+	test('notebook open/close, notebook ready when cell-document open event is fired', async function () {
+		const resource = await createRandomFile('', undefined, 'first', '.vsctestnb');
+		let didHappen = false;
+		const p = getEventOncePromise(vscode.workspace.onDidOpenTextDocument).then(doc => {
+			if (doc.uri.scheme !== 'vscode-notebook-cell') {
+				return;
+			}
+			const notebook = vscode.notebook.notebookDocuments.find(notebook => {
+				const cell = notebook.cells.find(cell => cell.document === doc);
+				return Boolean(cell);
+			});
+			assert.ok(notebook, `notebook for cell ${doc.uri} NOT found`);
+			didHappen = true;
+		});
+
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+		await p;
+		assert.strictEqual(didHappen, true);
+		await vscode.commands.executeCommand('workbench.action.closeAllEditors');
+	});
+
 	test('shared document in notebook editors', async function () {
 		assertInitalState();
 

src/vs/base/common/network.ts

@@ -58,6 +58,8 @@ export namespace Schemas {
 
 	export const vscodeNotebook = 'vscode-notebook';
 
+	export const vscodeNotebookCell = 'vscode-notebook-cell';
+
 	export const vscodeSettings = 'vscode-settings';
 
 	export const webviewPanel = 'webview-panel';

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

@@ -203,7 +203,7 @@ export class MainThreadNotebooks extends Disposable implements MainThreadNoteboo
 
 	async removeNotebookTextModel(uri: URI): Promise<void> {
 		// TODO@rebornix, remove cell should use emitDelta as well to ensure document/editor events are sent together
-		await this._proxy.$acceptDocumentAndEditorsDelta({ removedDocuments: [uri] });
+		this._proxy.$acceptDocumentAndEditorsDelta({ removedDocuments: [uri] });
 		let textModelDisposableStore = this._editorEventListenersMapping.get(uri.toString());
 		textModelDisposableStore?.dispose();
 		this._editorEventListenersMapping.delete(URI.from(uri).toString());

src/vs/workbench/api/common/extHost.protocol.ts

@@ -1633,7 +1633,7 @@ export interface ExtHostNotebookShape {
 	$acceptModelChanged(uriComponents: UriComponents, event: NotebookCellsChangedEvent): void;
 	$acceptModelSaved(uriComponents: UriComponents): void;
 	$acceptEditorPropertiesChanged(uriComponents: UriComponents, data: INotebookEditorPropertiesChangeData): void;
-	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): Promise<void>;
+	$acceptDocumentAndEditorsDelta(delta: INotebookDocumentsAndEditorsDelta): void;
 	$undoNotebook(viewType: string, uri: UriComponents, editId: number, isDirty: boolean): Promise<void>;
 	$redoNotebook(viewType: string, uri: UriComponents, editId: number, isDirty: boolean): Promise<void>;
 

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

@@ -53,7 +53,7 @@ export class ExtHostDocuments implements ExtHostDocumentsShape {
 	}
 
 	public getAllDocumentData(): ExtHostDocumentData[] {
-		return this._documentsAndEditors.allDocuments();
+		return [...this._documentsAndEditors.allDocuments()];
 	}
 
 	public getDocumentData(resource: vscode.Uri): ExtHostDocumentData | undefined {
@@ -69,8 +69,8 @@ export class ExtHostDocuments implements ExtHostDocumentsShape {
 
 	public getDocument(resource: vscode.Uri): vscode.TextDocument {
 		const data = this.getDocumentData(resource);
-		if (!data || !data.document) {
-			throw new Error('Unable to retrieve document from URI');
+		if (!data?.document) {
+			throw new Error(`Unable to retrieve document from URI '${resource}'`);
 		}
 		return data.document;
 	}

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

@@ -15,6 +15,19 @@ import { ExtHostTextEditor } from 'vs/workbench/api/common/extHostTextEditor';
 import * as typeConverters from 'vs/workbench/api/common/extHostTypeConverters';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ResourceMap } from 'vs/base/common/map';
+import { Schemas } from 'vs/base/common/network';
+import { Iterable } from 'vs/base/common/iterator';
+
+class Reference<T> {
+	private _count = 0;
+	constructor(readonly value: T) { }
+	ref() {
+		this._count++;
+	}
+	unref() {
+		return --this._count === 0;
+	}
+}
 
 export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsShape {
 
@@ -23,7 +36,7 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 	private _activeEditorId: string | null = null;
 
 	private readonly _editors = new Map<string, ExtHostTextEditor>();
-	private readonly _documents = new ResourceMap<ExtHostDocumentData>();
+	private readonly _documents = new ResourceMap<Reference<ExtHostDocumentData>>();
 
 	private readonly _onDidAddDocuments = new Emitter<ExtHostDocumentData[]>();
 	private readonly _onDidRemoveDocuments = new Emitter<ExtHostDocumentData[]>();
@@ -50,29 +63,40 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 			for (const uriComponent of delta.removedDocuments) {
 				const uri = URI.revive(uriComponent);
 				const data = this._documents.get(uri);
-				this._documents.delete(uri);
-				if (data) {
-					removedDocuments.push(data);
+				if (data?.unref()) {
+					this._documents.delete(uri);
+					removedDocuments.push(data.value);
 				}
 			}
 		}
 
 		if (delta.addedDocuments) {
 			for (const data of delta.addedDocuments) {
 				const resource = URI.revive(data.uri);
-				assert.ok(!this._documents.has(resource), `document '${resource} already exists!'`);
-
-				const documentData = new ExtHostDocumentData(
-					this._extHostRpc.getProxy(MainContext.MainThreadDocuments),
-					resource,
-					data.lines,
-					data.EOL,
-					data.modeId,
-					data.versionId,
-					data.isDirty
-				);
-				this._documents.set(resource, documentData);
-				addedDocuments.push(documentData);
+				let ref = this._documents.get(resource);
+
+				// double check -> only notebook cell documents should be
+				// referenced/opened more than once...
+				if (ref) {
+					if (resource.scheme !== Schemas.vscodeNotebookCell) {
+						throw new Error(`document '${resource} already exists!'`);
+					}
+				}
+				if (!ref) {
+					ref = new Reference(new ExtHostDocumentData(
+						this._extHostRpc.getProxy(MainContext.MainThreadDocuments),
+						resource,
+						data.lines,
+						data.EOL,
+						data.modeId,
+						data.versionId,
+						data.isDirty
+					));
+					this._documents.set(resource, ref);
+					addedDocuments.push(ref.value);
+				}
+
+				ref.ref();
 			}
 		}
 
@@ -92,7 +116,7 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 				assert.ok(this._documents.has(resource), `document '${resource}' does not exist`);
 				assert.ok(!this._editors.has(data.id), `editor '${data.id}' already exists!`);
 
-				const documentData = this._documents.get(resource)!;
+				const documentData = this._documents.get(resource)!.value;
 				const editor = new ExtHostTextEditor(
 					data.id,
 					this._extHostRpc.getProxy(MainContext.MainThreadTextEditors),
@@ -132,11 +156,11 @@ export class ExtHostDocumentsAndEditors implements ExtHostDocumentsAndEditorsSha
 	}
 
 	getDocument(uri: URI): ExtHostDocumentData | undefined {
-		return this._documents.get(uri);
+		return this._documents.get(uri)?.value;
 	}
 
-	allDocuments(): ExtHostDocumentData[] {
-		return [...this._documents.values()];
+	allDocuments(): Iterable<ExtHostDocumentData> {
+		return Iterable.map(this._documents.values(), ref => ref.value);
 	}
 
 	getEditor(id: string): ExtHostTextEditor | undefined {