Merge pull request microsoft#226545 from microsoft/rebornix/coherent-orangutan

Experimental undo/redo handling for notebook multicursor

Author: rebornix

src/vs/workbench/contrib/notebook/browser/contrib/multicursor/notebookMulticursor.ts

@@ -8,6 +8,7 @@ import { Emitter, Event } from 'vs/base/common/event';
 import { KeyCode, KeyMod } from 'vs/base/common/keyCodes';
 import { Disposable, DisposableStore } from 'vs/base/common/lifecycle';
 import { ResourceMap } from 'vs/base/common/map';
+import { URI } from 'vs/base/common/uri';
 import { EditorConfiguration } from 'vs/editor/browser/config/editorConfiguration';
 import { ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { CodeEditorWidget } from 'vs/editor/browser/widget/codeEditor/codeEditorWidget';
@@ -32,12 +33,15 @@ import { IConfigurationService } from 'vs/platform/configuration/common/configur
 import { ContextKeyExpr, IContextKeyService, RawContextKey } from 'vs/platform/contextkey/common/contextkey';
 import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
 import { KeybindingWeight } from 'vs/platform/keybinding/common/keybindingsRegistry';
+import { IPastFutureElements, IUndoRedoElement, IUndoRedoService, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
 import { INotebookActionContext, NotebookAction } from 'vs/workbench/contrib/notebook/browser/controller/coreActions';
 import { getNotebookEditorFromEditorPane, ICellViewModel, INotebookEditor, INotebookEditorContribution } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { registerNotebookContribution } from 'vs/workbench/contrib/notebook/browser/notebookEditorExtensions';
 import { CellEditorOptions } from 'vs/workbench/contrib/notebook/browser/view/cellParts/cellEditorOptions';
 import { NOTEBOOK_CELL_EDITOR_FOCUSED, NOTEBOOK_IS_ACTIVE_EDITOR } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
+import { RedoCommand, UndoCommand } from 'vs/editor/browser/editorExtensions';
+import { registerWorkbenchContribution2, WorkbenchPhase } from 'vs/workbench/common/contributions';
 
 const NOTEBOOK_ADD_FIND_MATCH_TO_SELECTION_ID = 'notebook.addFindMatchToSelection';
 
@@ -53,6 +57,7 @@ interface TrackedMatch {
 	wordSelections: Selection[];
 	config: IEditorConfiguration;
 	decorationIds: string[];
+	undoRedoHistory: IPastFutureElements;
 }
 
 export const NOTEBOOK_MULTI_SELECTION_CONTEXT = {
@@ -88,6 +93,7 @@ export class NotebookMultiCursorController extends Disposable implements INotebo
 		@ILanguageConfigurationService private readonly languageConfigurationService: ILanguageConfigurationService,
 		@IAccessibilityService private readonly accessibilityService: IAccessibilityService,
 		@IConfigurationService private readonly configurationService: IConfigurationService,
+		@IUndoRedoService private readonly undoRedoService: IUndoRedoService,
 	) {
 		super();
 
@@ -258,10 +264,67 @@ export class NotebookMultiCursorController extends Disposable implements INotebo
 		}));
 	}
 
+	private updateFinalUndoRedo() {
+		const anchorCellModel = this.anchorCell?.[1].getModel();
+		if (!anchorCellModel) {
+			// should not happen
+			return;
+		}
+
+		const newElementsMap: ResourceMap<IUndoRedoElement[]> = new ResourceMap<IUndoRedoElement[]>();
+		const resources: URI[] = [];
+
+		this.trackedMatches.forEach(trackedMatch => {
+			const undoRedoState = trackedMatch.undoRedoHistory;
+			if (!undoRedoState) {
+				return;
+			}
+
+			resources.push(trackedMatch.cellViewModel.uri);
+
+			const currentPastElements = this.undoRedoService.getElements(trackedMatch.cellViewModel.uri).past.slice();
+			const oldPastElements = trackedMatch.undoRedoHistory.past.slice();
+			const newElements = currentPastElements.slice(oldPastElements.length);
+			if (newElements.length === 0) {
+				return;
+			}
+
+			newElementsMap.set(trackedMatch.cellViewModel.uri, newElements);
+
+			this.undoRedoService.removeElements(trackedMatch.cellViewModel.uri);
+			oldPastElements.forEach(element => {
+				this.undoRedoService.pushElement(element);
+			});
+		});
+
+		this.undoRedoService.pushElement({
+			type: UndoRedoElementType.Workspace,
+			resources: resources,
+			label: 'Multi Cursor Edit',
+			code: 'multiCursorEdit',
+			confirmBeforeUndo: false,
+			undo: async () => {
+				newElementsMap.forEach(async value => {
+					value.reverse().forEach(async element => {
+						await element.undo();
+					});
+				});
+			},
+			redo: async () => {
+				newElementsMap.forEach(async value => {
+					value.forEach(async element => {
+						await element.redo();
+					});
+				});
+			}
+		});
+	}
+
 	public resetToIdleState() {
 		this.state = NotebookMultiCursorState.Idle;
 		this._nbMultiSelectState.set(NotebookMultiCursorState.Idle);
 		this._nbIsMultiSelectSession.set(false);
+		this.updateFinalUndoRedo();
 
 		this.trackedMatches.forEach(match => {
 			this.clearDecorations(match);
@@ -304,14 +367,17 @@ export class NotebookMultiCursorController extends Disposable implements INotebo
 				throw new Error('Active cell is not an instance of CodeEditorWidget');
 			}
 
+			textModel.pushStackElement();
+
 			this.trackedMatches = [];
 			const editorConfig = this.constructCellEditorOptions(this.anchorCell[0]);
 			const newMatch: TrackedMatch = {
 				cellViewModel: cell,
 				initialSelection: inputSelection,
 				wordSelections: [newSelection],
 				config: editorConfig, // cache this in the match so we can create new cursors controllers with the correct language config
-				decorationIds: []
+				decorationIds: [],
+				undoRedoHistory: this.undoRedoService.getElements(cell.uri)
 			};
 			this.trackedMatches.push(newMatch);
 
@@ -362,12 +428,16 @@ export class NotebookMultiCursorController extends Disposable implements INotebo
 					throw new Error('Active cell is not an instance of CodeEditorWidget');
 				}
 
+				const textModel = await resultCellViewModel.resolveTextModel();
+				textModel.pushStackElement();
+
 				newMatch = {
 					cellViewModel: resultCellViewModel,
 					initialSelection: initialSelection,
 					wordSelections: [newSelection],
 					config: this.constructCellEditorOptions(this.anchorCell[0]),
-					decorationIds: []
+					decorationIds: [],
+					undoRedoHistory: this.undoRedoService.getElements(resultCellViewModel.uri)
 				} satisfies TrackedMatch;
 				this.trackedMatches.push(newMatch);
 
@@ -407,6 +477,30 @@ export class NotebookMultiCursorController extends Disposable implements INotebo
 		});
 	}
 
+	async undo() {
+		const models: ITextModel[] = [];
+		for (const match of this.trackedMatches) {
+			const model = await match.cellViewModel.resolveTextModel();
+			if (model) {
+				models.push(model);
+			}
+		}
+
+		await Promise.all(models.map(model => model.undo()));
+	}
+
+	async redo() {
+		const models: ITextModel[] = [];
+		for (const match of this.trackedMatches) {
+			const model = await match.cellViewModel.resolveTextModel();
+			if (model) {
+				models.push(model);
+			}
+		}
+
+		await Promise.all(models.map(model => model.redo()));
+	}
+
 	private constructCellEditorOptions(cell: ICellViewModel): EditorConfiguration {
 		const cellEditorOptions = new CellEditorOptions(this.notebookEditor.getBaseCellEditorOptions(cell.language), this.notebookEditor.notebookOptions, this.configurationService);
 		const options = cellEditorOptions.getUpdatedValue(cell.internalMetadata, cell.uri);
@@ -602,7 +696,59 @@ class NotebookDeleteLeftMultiSelectionAction extends NotebookAction {
 	}
 }
 
+class NotebookMultiCursorUndoRedoContribution extends Disposable {
+
+	static readonly ID = 'workbench.contrib.notebook.multiCursorUndoRedo';
+
+	constructor(@IEditorService private readonly _editorService: IEditorService, @IConfigurationService private readonly configurationService: IConfigurationService) {
+		super();
+
+		if (!this.configurationService.getValue<boolean>('notebook.multiSelect.enabled')) {
+			return;
+		}
+
+		const PRIORITY = 10005;
+		this._register(UndoCommand.addImplementation(PRIORITY, 'notebook-multicursor-undo-redo', () => {
+			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane);
+			if (!editor) {
+				return false;
+			}
+
+			if (!editor.hasModel()) {
+				return false;
+			}
+
+			const controller = editor.getContribution<NotebookMultiCursorController>(NotebookMultiCursorController.id);
+
+			return controller.undo();
+		}, ContextKeyExpr.and(
+			ContextKeyExpr.equals('config.notebook.multiSelect.enabled', true),
+			NOTEBOOK_IS_ACTIVE_EDITOR,
+			NOTEBOOK_MULTI_SELECTION_CONTEXT.IsNotebookMultiSelect,
+		)));
+
+		this._register(RedoCommand.addImplementation(PRIORITY, 'notebook-multicursor-undo-redo', () => {
+			const editor = getNotebookEditorFromEditorPane(this._editorService.activeEditorPane);
+			if (!editor) {
+				return false;
+			}
+
+			if (!editor.hasModel()) {
+				return false;
+			}
+
+			const controller = editor.getContribution<NotebookMultiCursorController>(NotebookMultiCursorController.id);
+			return controller.redo();
+		}, ContextKeyExpr.and(
+			ContextKeyExpr.equals('config.notebook.multiSelect.enabled', true),
+			NOTEBOOK_IS_ACTIVE_EDITOR,
+			NOTEBOOK_MULTI_SELECTION_CONTEXT.IsNotebookMultiSelect,
+		)));
+	}
+}
+
 registerNotebookContribution(NotebookMultiCursorController.id, NotebookMultiCursorController);
 registerAction2(NotebookAddMatchToMultiSelectionAction);
 registerAction2(NotebookExitMultiSelectionAction);
 registerAction2(NotebookDeleteLeftMultiSelectionAction);
+registerWorkbenchContribution2(NotebookMultiCursorUndoRedoContribution.ID, NotebookMultiCursorUndoRedoContribution, WorkbenchPhase.BlockRestore);