Improve notebook edit navigation (#241535)

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

@@ -393,6 +393,8 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 				return d;
 			});
 			this._cellDiffInfo.set(diff, undefined);
+			const changeCount = countChanges(this._cellDiffInfo.get());
+			this._changesCount.set(changeCount, undefined);
 		}
 	}
 

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

@@ -59,8 +59,13 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 	private readonly _currentCell = observableValue<NotebookCellTextModel | undefined>(this, undefined);
 	readonly currentCell: IObservable<NotebookCellTextModel | undefined> = this._currentCell;
 
+	private readonly _currentChange = observableValue<{ change: ICellDiffInfo; index: number } | undefined>(this, undefined);
+	readonly currentChange: IObservable<{ change: ICellDiffInfo; index: number } | undefined> = this._currentChange;
+
 	private readonly cellEditorIntegrations = new Map<NotebookCellTextModel, { integration: ChatEditingCodeEditorIntegration; diff: ISettableObservable<IDocumentDiff2> }>();
 
+	private readonly insertDeleteDecorators: IObservable<{ insertedCellDecorator: NotebookInsertedCellDecorator; deletedCellDecorator: NotebookDeletedCellDecorator } | undefined>;
+
 	constructor(
 		private readonly _entry: IModifiedFileEntry,
 		private readonly notebookEditor: INotebookEditor,
@@ -99,7 +104,9 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 		}));
 
 		this._register(autorun(r => {
-			const changes = cellChanges.read(r).filter(c => c.type !== 'unchanged' && c.type !== 'delete');
+			const sortedCellChanges = sortCellChanges(cellChanges.read(r));
+
+			const changes = sortedCellChanges.filter(c => c.type !== 'unchanged' && c.type !== 'delete');
 			onDidChangeVisibleRanges.read(r);
 			if (!changes.length) {
 				this.cellEditorIntegrations.forEach(({ diff }) => {
@@ -121,13 +128,6 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 				} else {
 					const diff2 = observableValue(`diff${cell.handle}`, diff.diff);
 					const integration = this.instantiationService.createInstance(ChatEditingCodeEditorIntegration, _entry, editor, diff2);
-					this._register(autorun((r) => {
-						const current = integration.currentIndex.read(r);
-						const indexOfChange = current >= 0 ? this.getIndexOfChange(cell, current) : -1;
-						if (indexOfChange >= 0) {
-							this._currentIndex.set(indexOfChange, undefined);
-						}
-					}));
 					this.cellEditorIntegrations.set(cell, { integration, diff: diff2 });
 					this._register(integration);
 					this._register(editor.onDidDispose(() => {
@@ -150,10 +150,29 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 					this.cellEditorIntegrations.delete(cell);
 				}
 			});
+
+			// set initial index
+			this._currentIndex.set(0, undefined);
+			this._revealChange(sortedCellChanges[0]);
+
+			this._register(autorun(r => {
+				const currentChange = this.currentChange.read(r);
+				if (currentChange) {
+					const change = currentChange.change;
+					const indexInChange = currentChange.index;
+					const diffChangeIndex = sortCellChanges(this.cellChanges.get().filter(c => c.type !== 'unchanged')).findIndex(c => c === change);
+
+					if (diffChangeIndex !== -1) {
+						this._currentIndex.set(diffChangeIndex + indexInChange, undefined);
+					}
+				} else {
+					this._currentIndex.set(-1, undefined);
+				}
+			}));
 		}));
 
 
-		const insertDeleteDecorators = derivedWithStore((r, store) => {
+		this.insertDeleteDecorators = derivedWithStore((r, store) => {
 			if (!notebookEdotirViewModelAttached.read(r)) {
 				return;
 			}
@@ -169,7 +188,7 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 
 		this._register(autorun(r => {
 			const changes = cellChanges.read(r);
-			const decorators = insertDeleteDecorators.read(r);
+			const decorators = this.insertDeleteDecorators.read(r);
 			if (decorators) {
 				decorators.insertedCellDecorator.apply(changes);
 				decorators.deletedCellDecorator.apply(changes, originalModel);
@@ -177,18 +196,6 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 		}));
 	}
 
-	private getIndexOfChange(cell: NotebookCellTextModel, change: number): number {
-		if (this.getCurrentCell()?.cell !== cell) {
-			return -1;
-		}
-		const cellIndex = this.notebookModel.cells.findIndex(c => c.handle === cell.handle);
-		const indexOfCellChange = this.cellChanges.get().findIndex(c => c.modifiedCellIndex === cellIndex);
-		if (indexOfCellChange === -1) {
-			return -1;
-		}
-		// Count all changes upto this cell.
-		return countChanges(this.cellChanges.get().slice(0, indexOfCellChange)) + change;
-	}
 	getCurrentCell() {
 		const activeCell = this.notebookModel.cells.find(c => c.handle === this.notebookEditor.getActiveCell()?.handle) || this._currentCell.get();
 		if (!activeCell) {
@@ -246,55 +253,156 @@ export class ChatEditingNotebookEditorIntegration extends Disposable implements
 	}
 
 	reveal(firstOrLast: boolean): void {
-		const changes = this.cellChanges.get().filter(c => c.type === 'modified' || c.type === 'insert');
+		const changes = sortCellChanges(this.cellChanges.get().filter(c => c.type !== 'unchanged'));
 		if (!changes.length) {
 			return undefined;
 		}
-		const index = firstOrLast ?
+		const change = firstOrLast ? changes[0] : changes[changes.length - 1];
-			changes.reduce((prev, curr) => prev.modifiedCellIndex < curr.modifiedCellIndex ? prev : curr).modifiedCellIndex :
+		this._revealChange(change, firstOrLast);
-			changes.reduce((prev, curr) => prev.modifiedCellIndex > curr.modifiedCellIndex ? prev : curr).modifiedCellIndex;
+	}
-		const info = this.getCell(index);
+
-		if (info) {
+	private _revealChange(change: ICellDiffInfo, firstOrLast?: boolean) {
-			this.selectCell(info.cell);
+		switch (change.type) {
-			info.integration.reveal(firstOrLast);
+			case 'insert':
+			case 'modified':
+				{
+					const cell = this.getCell(change.modifiedCellIndex);
+					if (!cell) {
+						return false;
+					}
+
+					cell.integration.reveal(firstOrLast ?? true);
+					this._currentChange.set({ change: change, index: cell.integration.currentIndex.get() }, undefined);
+
+					return true;
+				}
+			case 'delete':
+				// reveal the deleted cell decorator
+				this._currentCell.set(undefined, undefined);
+				this.insertDeleteDecorators.get()?.deletedCellDecorator.reveal(change.originalCellIndex);
+				this._currentChange.set({ change: change, index: 0 }, undefined);
+				return true;
+			default:
+				break;
 		}
+
+		return false;
 	}
 
 	next(wrap: boolean): boolean {
-		const info = this.getCurrentCell();
+		const changes = sortCellChanges(this.cellChanges.get().filter(c => c.type !== 'unchanged'));
-		if (!info) {
+		const currentChange = this.currentChange.get();
-			return false;
+		if (!currentChange) {
-		}
+			const firstChange = changes[0];
-		if (info.integration.next(wrap)) {
+
-			return true;
+			if (firstChange) {
-		} else {
+				this._currentCell.set(undefined, undefined);
-			const info = this.getNextCell(true);
+				return this._revealChange(firstChange);
-			if (info) {
-				this.selectCell(info.cell);
-				this.selectCell(info.cell);
-				info.integration.reveal(true);
-				return true;
 			}
+
 			return false;
 		}
+
+		// go to next
+		// first check if we are at the end of the current change
+		switch (currentChange.change.type) {
+			case 'modified':
+				{
+					const currentChangeInfo = this.getCell(currentChange.change.modifiedCellIndex);
+					if (!currentChangeInfo) {
+						return false;
+					}
+
+					if (currentChangeInfo.integration.next(false)) {
+						this._currentChange.set({ change: currentChange.change, index: currentChangeInfo.integration.currentIndex.get() }, undefined);
+						return true;
+					} else {
+						const nextChange = changes[changes.indexOf(currentChange.change) + 1];
+						if (nextChange) {
+							return this._revealChange(nextChange, true);
+						}
+					}
+				}
+				break;
+			case 'insert':
+			case 'delete':
+				{
+					// go to next change directly
+					const nextChange = changes[changes.indexOf(currentChange.change) + 1];
+					if (nextChange) {
+						return this._revealChange(nextChange, true);
+					}
+				}
+				break;
+			default:
+				break;
+		}
+
+		if (wrap) {
+			return this.next(false);
+		}
+
+		return false;
 	}
+
 	previous(wrap: boolean): boolean {
-		const info = this.getCurrentCell();
+		const changes = sortCellChanges(this.cellChanges.get().filter(c => c.type !== 'unchanged'));
-		if (!info) {
+		const currentChange = this.currentChange.get();
-			return false;
+		if (!currentChange) {
-		}
+			const lastChange = changes[changes.length - 1];
-		if (info.integration.previous(wrap)) {
+			if (lastChange) {
-			return true;
+				this._currentCell.set(undefined, undefined);
-		} else {
+				return this._revealChange(lastChange, false);
-			const info = this.getNextCell(false);
-			if (info) {
-				this.selectCell(info.cell);
-				info.integration.reveal(false);
-				return true;
 			}
+
 			return false;
 		}
+
+		// go to previous
+		// first check if we are at the start of the current change
+		switch (currentChange.change.type) {
+			case 'modified':
+				{
+					const currentChangeInfo = this.getCell(currentChange.change.modifiedCellIndex);
+					if (!currentChangeInfo) {
+						return false;
+					}
+
+					if (currentChangeInfo.integration.previous(false)) {
+						this._currentChange.set({ change: currentChange.change, index: currentChangeInfo.integration.currentIndex.get() }, undefined);
+						return true;
+					} else {
+						const prevChange = changes[changes.indexOf(currentChange.change) - 1];
+						if (prevChange) {
+							return this._revealChange(prevChange, false);
+						}
+					}
+				}
+				break;
+			case 'insert':
+			case 'delete':
+				{
+					// go to previous change directly
+					const prevChange = changes[changes.indexOf(currentChange.change) - 1];
+					if (prevChange) {
+						return this._revealChange(prevChange, false);
+					}
+				}
+				break;
+			default:
+				break;
+		}
+
+		if (wrap) {
+			const lastChange = changes[changes.length - 1];
+			if (lastChange) {
+				return this._revealChange(lastChange, false);
+			}
+		}
+
+		return false;
 	}
+
 	enableAccessibleDiffView(): void {
 		this.getCurrentCell()?.integration.enableAccessibleDiffView();
 	}
@@ -335,3 +443,31 @@ export function countChanges(changes: ICellDiffInfo[]): number {
 	}, 0);
 
 }
+
+export function sortCellChanges(changes: ICellDiffInfo[]): ICellDiffInfo[] {
+	return [...changes].sort((a, b) => {
+		// For unchanged and modified, use modifiedCellIndex
+		if ((a.type === 'unchanged' || a.type === 'modified') &&
+			(b.type === 'unchanged' || b.type === 'modified')) {
+			return a.modifiedCellIndex - b.modifiedCellIndex;
+		}
+
+		// For delete entries, use originalCellIndex
+		if (a.type === 'delete' && b.type === 'delete') {
+			return a.originalCellIndex - b.originalCellIndex;
+		}
+
+		// For insert entries, use modifiedCellIndex
+		if (a.type === 'insert' && b.type === 'insert') {
+			return a.modifiedCellIndex - b.modifiedCellIndex;
+		}
+
+		// Mixed types: compare based on available indices
+		const aIndex = a.type === 'delete' ? a.originalCellIndex :
+			(a.type === 'insert' ? a.modifiedCellIndex : a.modifiedCellIndex);
+		const bIndex = b.type === 'delete' ? b.originalCellIndex :
+			(b.type === 'insert' ? b.modifiedCellIndex : b.modifiedCellIndex);
+
+		return aIndex - bIndex;
+	});
+}

src/vs/workbench/contrib/notebook/browser/diff/inlineDiff/notebookDeletedCellDecorator.ts

@@ -40,7 +40,7 @@ export class NotebookDeletedCellDecorator extends Disposable implements INoteboo
 		}
 		const cells = this._notebookEditor.getCellsInRange({ start: info.previousIndex, end: info.previousIndex + 1 });
 		if (!cells.length) {
-			return;
+			return this._notebookEditor.getLayoutInfo().height + info.offset;
 		}
 		const cell = cells[0];
 		const cellHeight = this._notebookEditor.getHeightOfElement(cell);
@@ -48,6 +48,22 @@ export class NotebookDeletedCellDecorator extends Disposable implements INoteboo
 		return top + cellHeight + info.offset;
 	}
 
+	reveal(deletedIndex: number) {
+		const top = this.getTop(deletedIndex);
+		if (typeof top === 'number') {
+			this._notebookEditor.focusContainer();
+			this._notebookEditor.revealOffsetInCenterIfOutsideViewport(top);
+
+			const info = this.deletedCellInfos.get(deletedIndex);
+
+			if (info) {
+				const prevIndex = info.previousIndex;
+				this._notebookEditor.setFocus({ start: prevIndex, end: prevIndex });
+				this._notebookEditor.setSelections([{ start: prevIndex, end: prevIndex }]);
+			}
+		}
+	}
+
 	public apply(diffInfo: CellDiffInfo[], original: NotebookTextModel): void {
 		this.clear();