Fix selection after reverting the document

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

@@ -272,3 +272,51 @@ suite('notebook undo redo', () => {
 		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
 	});
 });
+
+suite('notebook working copy', () => {
+	test('notebook revert on close', async function () {
+		const resource = vscode.Uri.parse(join(vscode.workspace.rootPath || '', './first.vsctestnb'));
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+
+		await waitFor(500);
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellBelow');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, '');
+
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellAbove');
+		await vscode.commands.executeCommand('default:type', { text: 'var abc = 0;' });
+
+		// close active editor from command will revert the file
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+
+		await waitFor(500);
+		assert.equal(vscode.notebook.activeNotebookEditor !== undefined, true);
+		assert.equal(vscode.notebook.activeNotebookEditor?.selection !== undefined, true);
+		assert.deepEqual(vscode.notebook.activeNotebookEditor?.document.cells[0], vscode.notebook.activeNotebookEditor?.selection);
+		assert.equal(vscode.notebook.activeNotebookEditor?.selection?.source, 'test');
+
+		await vscode.commands.executeCommand('workbench.action.files.save');
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+	});
+
+	test('notebook revert', async function () {
+		const resource = vscode.Uri.parse(join(vscode.workspace.rootPath || '', './first.vsctestnb'));
+		await vscode.commands.executeCommand('vscode.openWith', resource, 'notebookCoreTest');
+
+		await waitFor(500);
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellBelow');
+		assert.equal(vscode.notebook.activeNotebookEditor!.selection?.source, '');
+
+		await vscode.commands.executeCommand('notebook.cell.insertCodeCellAbove');
+		await vscode.commands.executeCommand('default:type', { text: 'var abc = 0;' });
+		await vscode.commands.executeCommand('workbench.action.files.revert');
+
+		assert.equal(vscode.notebook.activeNotebookEditor !== undefined, true);
+		assert.equal(vscode.notebook.activeNotebookEditor?.selection !== undefined, true);
+		assert.deepEqual(vscode.notebook.activeNotebookEditor?.document.cells[0], vscode.notebook.activeNotebookEditor?.selection);
+		assert.deepEqual(vscode.notebook.activeNotebookEditor?.document.cells.length, 1);
+		assert.equal(vscode.notebook.activeNotebookEditor?.selection?.source, 'test');
+
+		await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+	});
+});

src/vs/workbench/contrib/notebook/browser/view/notebookCellList.ts

@@ -157,7 +157,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 
 					for (let i = diff.start; i < diff.start + diff.deleteCount; i++) {
 						const cell = this.element(i);
-						if (this._viewModel!.hasCell(cell)) {
+						if (this._viewModel!.hasCell(cell.handle)) {
 							hideOutputs.push(...cell?.model.outputs);
 						} else {
 							deletedOutputs.push(...cell?.model.outputs);
@@ -177,7 +177,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 
 						for (let i = diff.start; i < diff.start + diff.deleteCount; i++) {
 							const cell = this.element(i);
-							if (this._viewModel!.hasCell(cell)) {
+							if (this._viewModel!.hasCell(cell.handle)) {
 								hideOutputs.push(...cell?.model.outputs);
 							} else {
 								deletedOutputs.push(...cell?.model.outputs);
@@ -299,7 +299,7 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 
 			for (let i = diff.start; i < diff.start + diff.deleteCount; i++) {
 				const cell = this.element(i);
-				if (this._viewModel!.hasCell(cell)) {
+				if (this._viewModel!.hasCell(cell.handle)) {
 					hideOutputs.push(...cell?.model.outputs);
 				} else {
 					deletedOutputs.push(...cell?.model.outputs);
@@ -316,6 +316,18 @@ export class NotebookCellList extends WorkbenchList<CellViewModel> implements ID
 	splice2(start: number, deleteCount: number, elements: CellViewModel[] = []): void {
 		// we need to convert start and delete count based on hidden ranges
 		super.splice(start, deleteCount, elements);
+
+		const selectionsLeft = [];
+		this._viewModel!.selectionHandles.forEach(handle => {
+			if (this._viewModel!.hasCell(handle)) {
+				selectionsLeft.push(handle);
+			}
+		});
+
+		if (!selectionsLeft.length && this._viewModel!.viewCells) {
+			// after splice, the selected cells are deleted
+			this._viewModel!.selectionHandles = [this._viewModel!.viewCells[0].handle];
+		}
 	}
 
 	getViewIndex(cell: ICellViewModel) {

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel.ts

@@ -268,7 +268,12 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 			});
 
 			diffs.reverse().forEach(diff => {
-				this._viewCells.splice(diff[0], diff[1], ...diff[2]);
+				const deletedCells = this._viewCells.splice(diff[0], diff[1], ...diff[2]);
+
+				deletedCells.forEach(cell => {
+					this._handleToViewCellMapping.delete(cell.handle);
+				});
+
 				diff[2].forEach(cell => {
 					this._handleToViewCellMapping.set(cell.handle, cell);
 					this._localStore.add(cell);
@@ -456,8 +461,8 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		return index + 1;
 	}
 
-	hasCell(cell: ICellViewModel) {
-		return this._handleToViewCellMapping.has(cell.handle);
+	hasCell(handle: number) {
+		return this._handleToViewCellMapping.has(handle);
 	}
 
 	getVersionId() {

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

@@ -70,7 +70,6 @@ export class NotebookEditorModel extends EditorModel implements IWorkingCopy, IN
 	}
 
 	async revert(options?: IRevertOptions | undefined): Promise<void> {
-		console.log(options);
 		await this.load({ forceReadFromDisk: true });
 		this._dirty = false;
 		this._onDidChangeDirty.fire();