Support moving notebook cells during chat edits (#242534)

* Support moving a nb cell whilst in chat edits

* Fix tests

* Fixes to insert/delete of cells

* More fixes

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookEntry.ts

@@ -48,7 +48,7 @@ import { CellEditState, getNotebookEditorFromEditorPane } from '../../../noteboo
 import { INotebookEditorService } from '../../../notebook/browser/services/notebookEditorService.js';
 import { NotebookCellTextModel } from '../../../notebook/common/model/notebookCellTextModel.js';
 import { NotebookTextModel } from '../../../notebook/common/model/notebookTextModel.js';
-import { CellEditType, ICellDto2, ICellEditOperation, ICellReplaceEdit, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookTextModelChangedEvent, TransientOptions } from '../../../notebook/common/notebookCommon.js';
+import { CellEditType, ICell, ICellDto2, ICellEditOperation, ICellReplaceEdit, IResolvedNotebookEditorModel, NotebookCellsChangeType, NotebookCellsModelMoveEvent, NotebookCellTextModelSplice, NotebookTextModelChangedEvent, TransientOptions } from '../../../notebook/common/notebookCommon.js';
 import { computeDiff } from '../../../notebook/common/notebookDiff.js';
 import { INotebookEditorModelResolverService } from '../../../notebook/common/notebookEditorModelResolverService.js';
 import { INotebookLoggingService } from '../../../notebook/common/notebookLoggingService.js';
@@ -243,14 +243,7 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 	}
 
 	mirrorNotebookEdits(e: NotebookTextModelChangedEvent) {
-		/**
-		 * TODO@DonJayamanne
-		 * If user deletes cells, invoke this.disposeDeletedCellEntries();
-		 * If user makes any changes, invoke this.computeStateAfterAcceptingRejectingChanges(true);
-		 * If user deletes/inserts cells manually, do we need to apply those to the original snapshot? Unlikely, double check.
-		 */
 		if (this._isEditFromUs || Array.from(this.cellEntryMap.values()).some(entry => entry.isEditFromUs)) {
-			// TODO@DonJayamanne Apply this same edit to the original notebook.
 			return;
 		}
 
@@ -288,44 +281,16 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 					break;
 				}
 				case NotebookCellsChangeType.ModelChange: {
+					let cellDiffs = sortCellChanges(this._cellsDiffInfo.get()).slice();
 					event.changes.forEach(change => {
-						const cells = change[2].map(cell => {
-							return {
-								cellKind: cell.cellKind,
-								language: cell.language,
-								metadata: cell.metadata,
-								outputs: cell.outputs,
-								source: cell.getValue(),
-								mime: undefined,
-								internalMetadata: cell.internalMetadata
-							} satisfies ICellDto2;
-						});
-						const cellDiffs = sortCellChanges(this._cellsDiffInfo.get()).slice();
-						const wasInsertedAsFirstCell = change[0] === 0;
-						const wasInsertedAsLastCell = change[0] === this.modifiedModel.cells.length - 1;
-						const diffEntryIndex = wasInsertedAsFirstCell ? 0 : (wasInsertedAsLastCell ? this.originalModel.cells.length : (cellDiffs.findIndex(d => d.modifiedCellIndex === change[0])));
-						const indexToInsertInOriginalModel = (wasInsertedAsFirstCell || diffEntryIndex === -1) ? 0 : (wasInsertedAsLastCell ? this.originalModel.cells.length : (((cellDiffs.slice(0, diffEntryIndex).reverse().find(c => typeof c.originalCellIndex === 'number')?.originalCellIndex ?? -1) + 1)));
-						const edit: ICellEditOperation = {
-							editType: CellEditType.Replace,
-							cells,
-							index: indexToInsertInOriginalModel,
-							count: change[1]
-						};
-						this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, true);
-						// If cells were deleted we handled that with this.disposeDeletedCellEntries();
-						if (diffEntryIndex >= 0) {
-							cellDiffs.splice(diffEntryIndex, change[1]);
-						}
-						// For inserted cells, we need to ensure that we create a corresponding CellEntry.
-						// So that any edits to the inserted cell is handled and mirrored over to the corresponding cell in original model.
-						cells.forEach((_, i) => {
-							const originalCellIndex = i + indexToInsertInOriginalModel;
-							const modifiedCellIndex = change[0] + i;
-							const unchangedCell = this.createModifiedCellDiffInfo(modifiedCellIndex, originalCellIndex);
-							cellDiffs.splice(diffEntryIndex === -1 ? 0 : diffEntryIndex, 0, unchangedCell);
-							this.updateCellDiffInfo(cellDiffs, undefined);
-						});
+						cellDiffs = adjustCellDiffAndOriginalModelBasedOnCellAddDelete(change,
+							cellDiffs,
+							this.modifiedModel.cells.length,
+							this.originalModel.cells.length,
+							this.originalModel.applyEdits.bind(this.originalModel),
+							this.createModifiedCellDiffInfo.bind(this));
 					});
+					this.updateCellDiffInfo(cellDiffs, undefined);
 					this.disposeDeletedCellEntries();
 					break;
 				}
@@ -379,15 +344,11 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 					break;
 				}
 				case NotebookCellsChangeType.Move: {
-					const edit: ICellEditOperation = {
-						editType: CellEditType.Move,
-						index: event.index,
-						length: event.length,
-						newIdx: event.newIdx
-					};
-					this.originalModel.applyEdits([edit], true, undefined, () => undefined, undefined, true);
-					// TODO@DonJayamanne
-					// We need to update the entries in _cellDiffInfo.
+					const result = adjustCellDiffAndOriginalModelBasedOnCellMovements(event, this._cellsDiffInfo.get().slice());
+					if (result) {
+						this.originalModel.applyEdits(result[1], true, undefined, () => undefined, undefined, true);
+						this._cellsDiffInfo.set(result[0], undefined);
+					}
 					break;
 				}
 				default: {
@@ -1280,3 +1241,203 @@ class ChatEditingNotebookCellEntry extends ObservableDisposable {
 		}
 	}
 }
+
+export function adjustCellDiffAndOriginalModelBasedOnCellAddDelete(change: NotebookCellTextModelSplice<ICell>,
+	cellDiffInfo: ICellDiffInfo[],
+	modifiedModelCellCount: number,
+	originalModelCellCount: number,
+	applyEdits: typeof NotebookTextModel.prototype.applyEdits,
+	createModifiedCellDiffInfo: (modifiedCellIndex: number, originalCellIndex: number) => ICellDiffInfo,
+): ICellDiffInfo[] {
+	cellDiffInfo = sortCellChanges(cellDiffInfo).slice();
+	const numberOfCellsInserted = change[2].length;
+	const numberOfCellsDeleted = change[1];
+	const cells = change[2].map(cell => {
+		return {
+			cellKind: cell.cellKind,
+			language: cell.language,
+			metadata: cell.metadata,
+			outputs: cell.outputs,
+			source: cell.getValue(),
+			mime: undefined,
+			internalMetadata: cell.internalMetadata
+		} satisfies ICellDto2;
+	});
+	const wasInsertedAsFirstCell = change[0] === 0;
+	const wasInsertedAsLastCell = change[0] === modifiedModelCellCount - 1;
+	const diffEntryIndex = wasInsertedAsFirstCell ? 0 : (wasInsertedAsLastCell ? cellDiffInfo.length - 1 : (cellDiffInfo.findIndex(d => d.modifiedCellIndex === change[0])));
+	const indexToInsertInOriginalModel = (wasInsertedAsFirstCell || diffEntryIndex === -1) ? 0 : (wasInsertedAsLastCell ? originalModelCellCount : (((cellDiffInfo.slice(0, diffEntryIndex).reverse().find(c => typeof c.originalCellIndex === 'number')?.originalCellIndex ?? -1) + 1)));
+	if (cells.length) {
+		const edit: ICellEditOperation = {
+			editType: CellEditType.Replace,
+			cells,
+			index: indexToInsertInOriginalModel,
+			count: change[1]
+		};
+		applyEdits([edit], true, undefined, () => undefined, undefined, true);
+	}
+	// If cells were deleted we handled that with this.disposeDeletedCellEntries();
+	if (numberOfCellsDeleted) {
+		// Adjust the indexes.
+		let numberOfOriginalCellsRemovedSoFar = 0;
+		let numberOfModifiedCellsRemovedSoFar = 0;
+		const modifiedIndexesToRemove = new Set<number>();
+		for (let i = 0; i < numberOfCellsDeleted; i++) {
+			modifiedIndexesToRemove.add(change[0] + i);
+		}
+		const itemsToRemove = new Set<ICellDiffInfo>();
+		for (let i = 0; i < cellDiffInfo.length; i++) {
+			const diff = cellDiffInfo[i];
+			if (i < diffEntryIndex) {
+				continue;
+			}
+
+			let changed = false;
+			if (typeof diff.modifiedCellIndex === 'number' && modifiedIndexesToRemove.has(diff.modifiedCellIndex)) {
+				// This will be removed.
+				numberOfModifiedCellsRemovedSoFar++;
+				if (typeof diff.originalCellIndex === 'number') {
+					numberOfOriginalCellsRemovedSoFar++;
+				}
+				itemsToRemove.add(diff);
+				continue;
+			}
+			if (typeof diff.modifiedCellIndex === 'number' && numberOfModifiedCellsRemovedSoFar) {
+				diff.modifiedCellIndex -= numberOfModifiedCellsRemovedSoFar;
+				changed = true;
+			}
+			if (typeof diff.originalCellIndex === 'number' && numberOfOriginalCellsRemovedSoFar) {
+				diff.originalCellIndex -= numberOfOriginalCellsRemovedSoFar;
+				changed = true;
+			}
+			if (changed) {
+				cellDiffInfo[i] = { ...diff };
+			}
+		}
+		cellDiffInfo = cellDiffInfo.filter(d => !itemsToRemove.has(d));
+	}
+
+	if (numberOfCellsInserted) {
+		for (let i = 0; i < cellDiffInfo.length; i++) {
+			const diff = cellDiffInfo[i];
+			if (i < diffEntryIndex) {
+				continue;
+			}
+			let changed = false;
+			if (typeof diff.modifiedCellIndex === 'number') {
+				diff.modifiedCellIndex += numberOfCellsInserted;
+				changed = true;
+			}
+			if (typeof diff.originalCellIndex === 'number') {
+				diff.originalCellIndex += numberOfCellsInserted;
+				changed = true;
+			}
+			if (changed) {
+				cellDiffInfo[i] = { ...diff };
+			}
+		}
+	}
+
+	// For inserted cells, we need to ensure that we create a corresponding CellEntry.
+	// So that any edits to the inserted cell is handled and mirrored over to the corresponding cell in original model.
+	cells.forEach((_, i) => {
+		const originalCellIndex = i + indexToInsertInOriginalModel;
+		const modifiedCellIndex = change[0] + i;
+		const unchangedCell = createModifiedCellDiffInfo(modifiedCellIndex, originalCellIndex);
+		cellDiffInfo.splice((diffEntryIndex === -1 ? 0 : diffEntryIndex) + i, 0, unchangedCell);
+	});
+	return cellDiffInfo;
+}
+/**
+ * Given the movements of cells in modified notebook, adjust the ICellDiffInfo[] array
+ * and generate edits for the old notebook (if required).
+ * TODO@DonJayamanne Handle bulk moves (movements of more than 1 cell).
+ */
+export function adjustCellDiffAndOriginalModelBasedOnCellMovements(event: NotebookCellsModelMoveEvent<ICell>, cellDiffInfo: ICellDiffInfo[]): [ICellDiffInfo[], ICellEditOperation[]] | undefined {
+	const minimumIndex = Math.min(event.index, event.newIdx);
+	const maximumIndex = Math.max(event.index, event.newIdx);
+	const cellDiffs = cellDiffInfo.slice();
+	const indexOfEntry = cellDiffs.findIndex(d => d.modifiedCellIndex === event.index);
+	const indexOfEntryToPlaceBelow = cellDiffs.findIndex(d => d.modifiedCellIndex === event.newIdx);
+	if (indexOfEntry === -1 || indexOfEntryToPlaceBelow === -1) {
+		return undefined;
+	}
+	// Create a new object so that the observable value is triggered.
+	// Besides we'll be updating the values of this object in place.
+	const entryToBeMoved = { ...cellDiffs[indexOfEntry] };
+	const moveDirection = event.newIdx > event.index ? 'down' : 'up';
+
+
+	const startIndex = cellDiffs.findIndex(d => d.modifiedCellIndex === minimumIndex);
+	const endIndex = cellDiffs.findIndex(d => d.modifiedCellIndex === maximumIndex);
+	const movingExistingCell = typeof entryToBeMoved.originalCellIndex === 'number';
+	let originalCellsWereEffected = false;
+	for (let i = 0; i < cellDiffs.length; i++) {
+		const diff = cellDiffs[i];
+		let changed = false;
+		if (moveDirection === 'down') {
+			if (i > startIndex && i <= endIndex) {
+				if (typeof diff.modifiedCellIndex === 'number') {
+					changed = true;
+					diff.modifiedCellIndex = diff.modifiedCellIndex - 1;
+				}
+				if (typeof diff.originalCellIndex === 'number' && movingExistingCell) {
+					diff.originalCellIndex = diff.originalCellIndex - 1;
+					originalCellsWereEffected = true;
+					changed = true;
+				}
+			}
+		} else {
+			if (i >= startIndex && i < endIndex) {
+				if (typeof diff.modifiedCellIndex === 'number') {
+					changed = true;
+					diff.modifiedCellIndex = diff.modifiedCellIndex + 1;
+				}
+				if (typeof diff.originalCellIndex === 'number' && movingExistingCell) {
+					diff.originalCellIndex = diff.originalCellIndex + 1;
+					originalCellsWereEffected = true;
+					changed = true;
+				}
+			}
+		}
+		// Create a new object so that the observable value is triggered.
+		// Do only if there's a change.
+		if (changed) {
+			cellDiffs[i] = { ...diff };
+		}
+	}
+	entryToBeMoved.modifiedCellIndex = event.newIdx;
+	const originalCellIndex = entryToBeMoved.originalCellIndex;
+	if (moveDirection === 'down') {
+		cellDiffs.splice(endIndex + 1, 0, entryToBeMoved);
+		cellDiffs.splice(startIndex, 1);
+		// If we're moving a new cell up/down, then we need just adjust just the modified indexes of the cells in between.
+		// If we're moving an existing up/down, then we need to adjust the original indexes as well.
+		if (typeof entryToBeMoved.originalCellIndex === 'number') {
+			entryToBeMoved.originalCellIndex = cellDiffs.slice(0, endIndex).reduce((lastOriginalIndex, diff) => typeof diff.originalCellIndex === 'number' ? Math.max(lastOriginalIndex, diff.originalCellIndex) : lastOriginalIndex, -1) + 1;
+		}
+	} else {
+		cellDiffs.splice(endIndex, 1);
+		cellDiffs.splice(startIndex, 0, entryToBeMoved);
+		// If we're moving a new cell up/down, then we need just adjust just the modified indexes of the cells in between.
+		// If we're moving an existing up/down, then we need to adjust the original indexes as well.
+		if (typeof entryToBeMoved.originalCellIndex === 'number') {
+			entryToBeMoved.originalCellIndex = cellDiffs.slice(0, startIndex).reduce((lastOriginalIndex, diff) => typeof diff.originalCellIndex === 'number' ? Math.max(lastOriginalIndex, diff.originalCellIndex) : lastOriginalIndex, -1) + 1;
+		}
+	}
+
+	// If this is a new cell that we're moving, and there are no existing cells in between, then we can just move the new cell.
+	// I.e. no need to update the original notebook model.
+	if (typeof entryToBeMoved.originalCellIndex === 'number' && originalCellsWereEffected && typeof originalCellIndex === 'number' && entryToBeMoved.originalCellIndex !== originalCellIndex) {
+		const edit: ICellEditOperation = {
+			editType: CellEditType.Move,
+			index: originalCellIndex,
+			length: event.length,
+			newIdx: entryToBeMoved.originalCellIndex
+		};
+
+		return [cellDiffs, [edit]];
+	}
+
+	return [cellDiffs, []];
+}

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingNotebookEditorIntegration.ts

@@ -595,9 +595,17 @@ export function sortCellChanges(changes: ICellDiffInfo[]): ICellDiffInfo[] {
 			return a.modifiedCellIndex - b.modifiedCellIndex;
 		}
 
+		if (a.type === 'delete' && b.type === 'insert') {
+			return -1;
+		}
+		if (a.type === 'insert' && b.type === 'delete') {
+			return 1;
+		}
+
 		if ((a.type === 'delete' && b.type !== 'insert') || (a.type !== 'insert' && b.type === 'delete')) {
 			return a.originalCellIndex - b.originalCellIndex;
 		}
+
 		// Mixed types: compare based on available indices
 		const aIndex = a.type === 'delete' ? a.originalCellIndex :
 			(a.type === 'insert' ? a.modifiedCellIndex : a.modifiedCellIndex);