Undo/Redo for cell manipulation

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

@@ -885,9 +885,38 @@ registerAction2(class extends Action2 {
 			return;
 		}
 
-		if (viewModel.canUndo()) {
-			viewModel.undo();
+		viewModel.undo();
+	}
+});
+
+registerAction2(class extends Action2 {
+	constructor() {
+		super({
+			id: 'workbench.action.notebook.redo',
+			title: 'Notebook Redo',
+			keybinding: {
+				when: ContextKeyExpr.and(NOTEBOOK_EDITOR_FOCUSED, ContextKeyExpr.not(InputFocusedContextKey)),
+				primary: KeyMod.CtrlCmd | KeyMod.Shift | KeyCode.KEY_Z,
+				weight: KeybindingWeight.WorkbenchContrib
+			}
+		});
+	}
+
+	async run(accessor: ServicesAccessor): Promise<void> {
+		const editorService = accessor.get(IEditorService);
+
+		const editor = getActiveNotebookEditor(editorService);
+		if (!editor) {
+			return;
 		}
+
+		const viewModel = editor.viewModel;
+
+		if (!viewModel) {
+			return;
+		}
+
+		viewModel.redo();
 	}
 });
 

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

@@ -27,7 +27,7 @@ import { INotebookService } from 'vs/workbench/contrib/notebook/browser/notebook
 import { OutputRenderer } from 'vs/workbench/contrib/notebook/browser/view/output/outputRenderer';
 import { BackLayerWebView } from 'vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView';
 import { CodeCellRenderer, MarkdownCellRenderer, NotebookCellListDelegate } from 'vs/workbench/contrib/notebook/browser/view/renderers/cellRenderer';
-import { NotebookCellsSplice, IOutput, CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { IOutput, CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IWebviewService } from 'vs/workbench/contrib/webview/browser/webview';
 import { getExtraColor } from 'vs/workbench/contrib/welcome/walkThrough/common/walkThroughUtils';
 import { IEditorGroup, IEditorGroupsService } from 'vs/workbench/services/editor/common/editorGroupsService';
@@ -329,8 +329,18 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		const viewState = this.loadTextEditorViewState(input);
 		this.notebookViewModel.restoreEditorViewState(viewState);
 
-		this.localStore.add(this.notebookViewModel.onDidChangeCells((e) => {
-			this.updateViewCells(e);
+		this.localStore.add(this.notebookViewModel.onDidChangeViewCells((e) => {
+			if (e.synchronous) {
+				e.splices.reverse().forEach((diff) => {
+					this.list?.splice(diff[0], diff[1], diff[2]);
+				});
+			} else {
+				DOM.scheduleAtNextAnimationFrame(() => {
+					e.splices.reverse().forEach((diff) => {
+						this.list?.splice(diff[0], diff[1], diff[2]);
+					});
+				});
+			}
 		}));
 
 		this.webview?.updateRendererPreloads(this.notebookViewModel.renderers);
@@ -533,16 +543,6 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		});
 	}
 
-	updateViewCells(splices: NotebookCellsSplice[]) {
-		DOM.scheduleAtNextAnimationFrame(() => {
-			splices.reverse().forEach((diff) => {
-				this.list?.splice(diff[0], diff[1], diff[2].map(cell => {
-					return this.instantiationService.createInstance(CellViewModel, this.notebookViewModel!.viewType, this.notebookViewModel!.handle, cell);
-				}));
-			});
-		});
-	}
-
 	async insertNotebookCell(cell: ICellViewModel, type: CellKind, direction: 'above' | 'below', initialText: string = ''): Promise<void> {
 		const newLanguages = this.notebookViewModel!.languages;
 		const language = newLanguages && newLanguages.length ? newLanguages[0] : 'markdown';
@@ -550,9 +550,7 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 		const insertIndex = direction === 'above' ? index : index + 1;
 		const newModeCell = await this.notebookService.createNotebookCell(this.notebookViewModel!.viewType, this.notebookViewModel!.uri, insertIndex, language, type);
 		newModeCell!.source = initialText.split(/\r?\n/g);
-		const newCell = this.notebookViewModel!.insertCell(insertIndex, newModeCell!);
-
-		this.list?.splice(insertIndex, 0, [newCell]);
+		const newCell = this.notebookViewModel!.insertCell(insertIndex, newModeCell!, true);
 		this.list?.setFocus([insertIndex]);
 
 		if (type === CellKind.Markdown) {
@@ -567,8 +565,7 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 	async deleteNotebookCell(cell: ICellViewModel): Promise<void> {
 		const index = this.notebookViewModel!.getViewCellIndex(cell);
 		await this.notebookService.deleteNotebookCell(this.notebookViewModel!.viewType, this.notebookViewModel!.uri, index);
-		this.notebookViewModel!.deleteCell(index);
-		this.list?.splice(index, 1);
+		this.notebookViewModel!.deleteCell(index, true);
 	}
 
 	moveCellDown(cell: ICellViewModel): void {
@@ -584,13 +581,10 @@ export class NotebookEditor extends BaseEditor implements INotebookEditor {
 	}
 
 	private moveCellToIndex(cell: ICellViewModel, index: number, newIdx: number): void {
-		if (!this.notebookViewModel!.moveCellToIdx(index, newIdx)) {
+		if (!this.notebookViewModel!.moveCellToIdx(index, newIdx, true)) {
 			return;
 		}
 
-		this.list?.splice(index, 1);
-		this.list!.splice(newIdx, 0, [cell as CellViewModel]);
-
 		DOM.scheduleAtNextAnimationFrame(() => {
 			this.list?.revealInCenterIfOutsideViewport(index + 1);
 		});

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

+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ICell } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { IResourceUndoRedoElement, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
+import { URI } from 'vs/base/common/uri';
+import { CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookCellViewModel';
+import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
+import { NotebookViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModel';
+
+
+/**
+ * It should not modify Undo/Redo stack
+ */
+export interface ICellEditingDelegate {
+	insertCell?(index: number, viewCell: CellViewModel): void;
+	deleteCell?(index: number, cell: ICell): void;
+	moveCell?(fromIndex: number, toIndex: number): void;
+}
+
+export class InsertCellEdit implements IResourceUndoRedoElement {
+	type: UndoRedoElementType.Resource = UndoRedoElementType.Resource;
+	label: string = 'Insert Cell';
+	constructor(
+		public resource: URI,
+		private insertIndex: number,
+		private cell: CellViewModel,
+		private editingDelegate: ICellEditingDelegate
+	) {
+	}
+
+	undo(): void | Promise<void> {
+		if (!this.editingDelegate.deleteCell) {
+			throw new Error('Notebook Delete Cell not implemented for Undo/Redo');
+		}
+
+		this.editingDelegate.deleteCell(this.insertIndex, this.cell.cell);
+	}
+	redo(): void | Promise<void> {
+		if (!this.editingDelegate.insertCell) {
+			throw new Error('Notebook Insert Cell not implemented for Undo/Redo');
+		}
+
+		this.editingDelegate.insertCell(this.insertIndex, this.cell);
+	}
+}
+
+export class DeleteCellEdit implements IResourceUndoRedoElement {
+	type: UndoRedoElementType.Resource = UndoRedoElementType.Resource;
+	label: string = 'Delete Cell';
+
+	private _rawCell: ICell;
+	constructor(
+		public resource: URI,
+		private insertIndex: number,
+		cell: CellViewModel,
+		private editingDelegate: ICellEditingDelegate,
+		private instantiationService: IInstantiationService,
+		private notebookViewModel: NotebookViewModel
+	) {
+		this._rawCell = cell.cell;
+
+		// save inmem text to `ICell`
+		this._rawCell.source = [cell.getText()];
+	}
+
+	undo(): void | Promise<void> {
+		if (!this.editingDelegate.insertCell) {
+			throw new Error('Notebook Insert Cell not implemented for Undo/Redo');
+		}
+
+		const cell = this.instantiationService.createInstance(CellViewModel, this.notebookViewModel.viewType, this.notebookViewModel.handle, this._rawCell);
+		this.editingDelegate.insertCell(this.insertIndex, cell);
+	}
+
+	redo(): void | Promise<void> {
+		if (!this.editingDelegate.deleteCell) {
+			throw new Error('Notebook Delete Cell not implemented for Undo/Redo');
+		}
+
+		this.editingDelegate.deleteCell(this.insertIndex, this._rawCell);
+	}
+}
+
+export class MoveCellEdit implements IResourceUndoRedoElement {
+	type: UndoRedoElementType.Resource = UndoRedoElementType.Resource;
+	label: string = 'Delete Cell';
+
+	constructor(
+		public resource: URI,
+		private fromIndex: number,
+		private toIndex: number,
+		private editingDelegate: ICellEditingDelegate
+	) {
+	}
+
+	undo(): void | Promise<void> {
+		if (!this.editingDelegate.moveCell) {
+			throw new Error('Notebook Move Cell not implemented for Undo/Redo');
+		}
+
+		this.editingDelegate.moveCell(this.toIndex, this.fromIndex);
+	}
+
+	redo(): void | Promise<void> {
+		if (!this.editingDelegate.moveCell) {
+			throw new Error('Notebook Move Cell not implemented for Undo/Redo');
+		}
+
+		this.editingDelegate.moveCell(this.fromIndex, this.toIndex);
+	}
+}

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

@@ -9,7 +9,7 @@ import * as editorCommon from 'vs/editor/common/editorCommon';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { NotebookEditorModel } from 'vs/workbench/contrib/notebook/browser/notebookEditorInput';
 import { CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookCellViewModel';
-import { NotebookCellsSplice, ICell } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { ICell } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { IModelDeltaDecoration } from 'vs/editor/common/model';
 import { onUnexpectedError } from 'vs/base/common/errors';
 import { CellFindMatch, CellState, ICellViewModel } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
@@ -18,6 +18,7 @@ import { Range } from 'vs/editor/common/core/range';
 import { WorkspaceTextEdit } from 'vs/editor/common/modes';
 import { URI } from 'vs/base/common/uri';
 import { IUndoRedoService } from 'vs/platform/undoRedo/common/undoRedo';
+import { InsertCellEdit, DeleteCellEdit, MoveCellEdit } from 'vs/workbench/contrib/notebook/browser/viewModel/cellEdit';
 
 export interface INotebookEditorViewState {
 	editingCells: { [key: number]: boolean };
@@ -41,6 +42,17 @@ export interface IModelDecorationsChangeAccessor {
 const invalidFunc = () => { throw new Error(`Invalid change accessor`); };
 
 
+export type NotebookViewCellsSplice = [
+	number /* start */,
+	number /* delete count */,
+	CellViewModel[]
+];
+
+export interface INotebookViewCellsUpdateEvent {
+	synchronous: boolean;
+	splices: NotebookViewCellsSplice[];
+}
+
 export class NotebookViewModel extends Disposable {
 	private _localStore: DisposableStore = this._register(new DisposableStore());
 	private _viewCells: CellViewModel[] = [];
@@ -69,8 +81,8 @@ export class NotebookViewModel extends Disposable {
 		return this._model.notebook.uri;
 	}
 
-	private readonly _onDidChangeCells = new Emitter<NotebookCellsSplice[]>();
-	get onDidChangeCells(): Event<NotebookCellsSplice[]> { return this._onDidChangeCells.event; }
+	private readonly _onDidChangeViewCells = new Emitter<INotebookViewCellsUpdateEvent>();
+	get onDidChangeViewCells(): Event<INotebookViewCellsUpdateEvent> { return this._onDidChangeViewCells.event; }
 
 	private _lastNotebookEditResource: URI[] = [];
 
@@ -90,7 +102,15 @@ export class NotebookViewModel extends Disposable {
 	) {
 		super();
 
-		this._register(this._model.onDidChangeCells(e => this._onDidChangeCells.fire(e)));
+		this._register(this._model.onDidChangeCells(e => {
+			this._onDidChangeViewCells.fire({
+				synchronous: true,
+				splices: e.map(splice => {
+					return [splice[0], splice[1], splice[2].map(cell => this.instantiationService.createInstance(CellViewModel, this.viewType, this.handle, cell))];
+				})
+			});
+		}));
+
 		this._viewCells = this._model!.notebook!.cells.map(cell => {
 			const viewCell = this.instantiationService.createInstance(CellViewModel, this.viewType, this._model!.notebook!.handle, cell);
 			this._localStore.add(viewCell);
@@ -114,22 +134,53 @@ export class NotebookViewModel extends Disposable {
 		return this._viewCells.indexOf(cell as CellViewModel);
 	}
 
-	insertCell(index: number, cell: ICell): CellViewModel {
+	insertCell(index: number, cell: ICell, synchronous: boolean): CellViewModel {
 		const newCell = this.instantiationService.createInstance(CellViewModel, this.viewType, this.handle, cell);
 		this._viewCells!.splice(index, 0, newCell);
 		this._model.insertCell(newCell.cell, index);
 		this._localStore.add(newCell);
+		this.undoService.pushElement(new InsertCellEdit(this.uri, index, newCell, {
+			insertCell: (insertIndex: number, viewCell: CellViewModel) => {
+				this._viewCells!.splice(insertIndex, 0, viewCell);
+				this._model.insertCell(viewCell.cell, insertIndex);
+				this._localStore.add(viewCell);
+				this._onDidChangeViewCells.fire({ synchronous: true, splices: [[insertIndex, 0, [viewCell]]] });
+			},
+			deleteCell: (deleteIndex: number, cell: ICell) => {
+				this._viewCells.splice(deleteIndex, 1);
+				this._model.deleteCell(cell);
+				this._onDidChangeViewCells.fire({ synchronous: true, splices: [[deleteIndex, 1, []]] });
+			}
+		}));
+
+		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 0, [newCell]]] });
 		return newCell;
 	}
 
-	deleteCell(index: number) {
+	deleteCell(index: number, synchronous: boolean) {
 		let viewCell = this._viewCells[index];
 		this._viewCells.splice(index, 1);
 		this._model.deleteCell(viewCell.cell);
+
+		this.undoService.pushElement(new DeleteCellEdit(this.uri, index, viewCell, {
+			insertCell: (insertIndex: number, viewCell: CellViewModel) => {
+				this._viewCells!.splice(insertIndex, 0, viewCell);
+				this._model.insertCell(viewCell.cell, insertIndex);
+				this._localStore.add(viewCell);
+				this._onDidChangeViewCells.fire({ synchronous: true, splices: [[insertIndex, 0, [viewCell]]] });
+			},
+			deleteCell: (deleteIndex: number, cell: ICell) => {
+				this._viewCells.splice(deleteIndex, 1);
+				this._model.deleteCell(cell);
+				this._onDidChangeViewCells.fire({ synchronous: true, splices: [[deleteIndex, 1, []]] });
+			}
+		}, this.instantiationService, this));
+
+		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 1, []]] });
 		viewCell.dispose();
 	}
 
-	moveCellToIdx(index: number, newIdx: number): boolean {
+	moveCellToIdx(index: number, newIdx: number, synchronous: boolean, pushedToUndoStack: boolean = true): boolean {
 		const viewCell = this.viewCells[index] as CellViewModel;
 		if (!viewCell) {
 			return false;
@@ -141,6 +192,17 @@ export class NotebookViewModel extends Disposable {
 		this.viewCells!.splice(newIdx, 0, viewCell);
 		this._model.insertCell(viewCell.cell, newIdx);
 
+		if (pushedToUndoStack) {
+			this.undoService.pushElement(new MoveCellEdit(this.uri, index, newIdx, {
+				moveCell: (fromIndex: number, toIndex: number) => {
+					this.moveCellToIdx(fromIndex, toIndex, true, false);
+				}
+			}));
+		}
+
+		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[index, 1, []]] });
+		this._onDidChangeViewCells.fire({ synchronous: synchronous, splices: [[newIdx, 0, [viewCell]]] });
+
 		return true;
 	}
 
@@ -296,34 +358,15 @@ export class NotebookViewModel extends Disposable {
 	}
 
 	canUndo(): boolean {
-		const lastResource = this.lastNotebookEditResource;
-
-		if (!lastResource) {
-			return false;
-		}
-
-		const lastElement = this.undoService.getLastElement(lastResource);
-
-		if (lastElement?.label === 'Notebook Replace' || lastElement?.label === 'Notebook Replace All') {
-			return true;
-		}
-
-		return false;
+		return this.undoService.canUndo(this.uri);
 	}
 
 	undo() {
-		const lastResource = this.lastNotebookEditResource;
-
-		if (!lastResource) {
-			return;
-		}
-
-		const lastElement = this.undoService.getLastElement(lastResource);
+		this.undoService.undo(this.uri);
+	}
 
-		if (lastElement?.label === 'Notebook Replace' || lastElement?.label === 'Notebook Replace All') {
-			this.undoService.undo(lastResource);
-			this._lastNotebookEditResource.pop();
-		}
+	redo() {
+		this.undoService.redo(this.uri);
 	}
 
 	equal(model: NotebookEditorModel) {

src/vs/workbench/contrib/notebook/test/notebookViewModel.test.ts

@@ -35,12 +35,12 @@ suite('NotebookViewModel', () => {
 				[['var b = 2;'], 'javascript', CellKind.Code, []]
 			],
 			(editor, viewModel) => {
-				const cell = viewModel.insertCell(1, new TestCell(viewModel.viewType, 0, ['var c = 3;'], 'javascript', CellKind.Code, []));
+				const cell = viewModel.insertCell(1, new TestCell(viewModel.viewType, 0, ['var c = 3;'], 'javascript', CellKind.Code, []), true);
 				assert.equal(viewModel.viewCells.length, 3);
 				assert.equal(viewModel.notebookDocument.cells.length, 3);
 				assert.equal(viewModel.getViewCellIndex(cell), 1);
 
-				viewModel.deleteCell(1);
+				viewModel.deleteCell(1, true);
 				assert.equal(viewModel.viewCells.length, 2);
 				assert.equal(viewModel.notebookDocument.cells.length, 2);
 				assert.equal(viewModel.getViewCellIndex(cell), -1);
@@ -62,13 +62,13 @@ suite('NotebookViewModel', () => {
 				const lastViewCell = viewModel.viewCells[viewModel.viewCells.length - 1];
 
 				const insertIndex = viewModel.getViewCellIndex(firstViewCell) + 1;
-				const cell = viewModel.insertCell(insertIndex, new TestCell(viewModel.viewType, 3, ['var c = 3;'], 'javascript', CellKind.Code, []));
+				const cell = viewModel.insertCell(insertIndex, new TestCell(viewModel.viewType, 3, ['var c = 3;'], 'javascript', CellKind.Code, []), true);
 
 				const addedCellIndex = viewModel.getViewCellIndex(cell);
-				viewModel.deleteCell(addedCellIndex);
+				viewModel.deleteCell(addedCellIndex, true);
 
 				const secondInsertIndex = viewModel.getViewCellIndex(lastViewCell) + 1;
-				const cell2 = viewModel.insertCell(secondInsertIndex, new TestCell(viewModel.viewType, 4, ['var d = 4;'], 'javascript', CellKind.Code, []));
+				const cell2 = viewModel.insertCell(secondInsertIndex, new TestCell(viewModel.viewType, 4, ['var d = 4;'], 'javascript', CellKind.Code, []), true);
 
 				assert.equal(viewModel.viewCells.length, 3);
 				assert.equal(viewModel.notebookDocument.cells.length, 3);