Fix split cell undo (#17845)

* fix split cell undo

src/sql/workbench/contrib/notebook/test/electron-browser/notebookModel.test.ts

@@ -15,7 +15,7 @@ import { TestDialogService } from 'vs/platform/dialogs/test/common/testDialogSer
 import { URI } from 'vs/base/common/uri';
 
 import { ExecuteManagerStub, SerializationManagerStub } from 'sql/workbench/contrib/notebook/test/stubs';
-import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
+import { NotebookModel, SplitCell } from 'sql/workbench/services/notebook/browser/models/notebookModel';
 import { ModelFactory } from 'sql/workbench/services/notebook/browser/models/modelFactory';
 import { IClientSession, INotebookModelOptions, NotebookContentChange, IClientSessionOptions, ICellModel } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
 import { ClientSession } from 'sql/workbench/services/notebook/browser/models/clientSession';
@@ -605,6 +605,49 @@ suite('notebook model', function (): void {
 		assert.strictEqual(model.cells.indexOf(secondCell), 0, 'Failed to redo');
 	});
 
+	test('Should merge when undoing split cells', async function (): Promise<void> {
+		let expectedNotebookContentSplitCells: nb.INotebookContents = {
+			cells: [{
+				cell_type: CellTypes.Code,
+				source: ['foobar '],
+				execution_count: 1
+			}, {
+				cell_type: CellTypes.Code,
+				source: ['    hello'],
+				execution_count: 1
+			}],
+			metadata: {
+				kernelspec: {
+					name: 'mssql',
+					language: 'sql',
+					display_name: 'SQL'
+				},
+				language_info: {
+					name: 'sql'
+				}
+			},
+			nbformat: 4,
+			nbformat_minor: 5
+		};
+		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader);
+		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContentSplitCells));
+		defaultModelOptions.contentLoader = mockContentManager.object;
+
+		// When I initialize the model
+		let model = new NotebookModel(defaultModelOptions, undefined, logService, undefined, new NullAdsTelemetryService(), queryConnectionService.object, configurationService, undoRedoService);
+		await model.loadContents();
+
+		let splitCells: SplitCell[] = [
+			{ cell: model.cells[0], prefix: undefined },
+			{ cell: model.cells[1], prefix: '\n' }
+		];
+
+		// Merge cells
+		model.mergeCells(splitCells);
+		assert.strictEqual(model.cells.length, 1, 'Cells not deleted after merging');
+		assert.notStrictEqual(model.cells[0].source, ['foobar ', '\n', '    hello'], 'Cell source is not copied correctly');
+	});
+
 	test('Should notify cell on metadata change', async function (): Promise<void> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentLoader);
 		mockContentManager.setup(c => c.loadContent()).returns(() => Promise.resolve(expectedNotebookContent));

src/sql/workbench/services/notebook/browser/models/cellEdit.ts

@@ -5,7 +5,7 @@
 
 import { IResourceUndoRedoElement, UndoRedoElementType } from 'vs/platform/undoRedo/common/undoRedo';
 import { ICellModel, MoveDirection } from 'sql/workbench/services/notebook/browser/models/modelInterfaces';
-import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/notebookModel';
+import { NotebookModel, SplitCell } from 'sql/workbench/services/notebook/browser/models/notebookModel';
 import * as TelemetryKeys from 'sql/platform/telemetry/common/telemetryKeys';
 import { localize } from 'vs/nls';
 
@@ -36,11 +36,11 @@ export class SplitCellEdit implements IResourceUndoRedoElement {
 	resource = this.model.notebookUri;
 	private readonly cellOperation = { cell_operation: 'split_cell' };
 
-	constructor(private model: NotebookModel, private firstCell: ICellModel, private secondCell: ICellModel, private newLinesRemoved: string[]) {
+	constructor(private model: NotebookModel, private cells: SplitCell[]) {
 	}
 
 	undo(): void {
-		this.model.mergeCells(this.firstCell, this.secondCell, this.newLinesRemoved);
+		this.model.mergeCells(this.cells);
 		this.model.sendNotebookTelemetryActionEvent(TelemetryKeys.NbTelemetryAction.UndoCell, this.cellOperation);
 	}
 

src/sql/workbench/services/notebook/browser/models/notebookModel.ts

@@ -58,6 +58,11 @@ interface INotebookMetadataInternal extends nb.INotebookMetadata {
 	azdata_notebook_guid?: string;
 }
 
+export type SplitCell = {
+	cell: ICellModel;
+	prefix: string | undefined;
+};
+
 type NotebookMetadataKeys = Required<nb.INotebookMetadata>;
 const expectedMetadataKeys: NotebookMetadataKeys = {
 	kernelspec: undefined,
@@ -573,7 +578,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 				let newCell = undefined, tailCell = undefined, partialSource = undefined;
 				let newCellIndex = index;
 				let tailCellIndex = index;
-				let newLinesRemoved: string[] = [];
+				let splitCells: SplitCell[] = [];
 
 				// Save UI state
 				let showMarkdown = this.cells[index].showMarkdown;
@@ -608,6 +613,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						headsource = headsource.concat(partialSource.toString());
 					}
 					this.cells[index].source = headsource;
+					splitCells.push({ cell: this.cells[index], prefix: undefined });
 				}
 
 				if (newCellContent.length) {
@@ -630,6 +636,7 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						newCell.source = newSource;
 						newCellIndex++;
 						this.insertCell(newCell, newCellIndex, false);
+						splitCells.push({ cell: this.cells[newCellIndex], prefix: undefined });
 					}
 					else { //update the existing cell
 						this.cells[index].source = newSource;
@@ -644,21 +651,22 @@ export class NotebookModel extends Disposable implements INotebookModel {
 						partialSource = source.slice(tailRange.startLineNumber - 1, tailRange.startLineNumber)[0].slice(tailRange.startColumn - 1);
 						tailSource.splice(0, 1, partialSource);
 					}
+					let newlinesBeforeTailCellContent: string;
 					//Remove the trailing empty line after the cursor
 					if (tailSource[0] === '\r\n' || tailSource[0] === '\n') {
-						newLinesRemoved = tailSource.splice(0, 1);
+						newlinesBeforeTailCellContent = tailSource.splice(0, 1)[0];
 					}
 					tailCell.source = tailSource;
 					tailCellIndex = newCellIndex + 1;
 					this.insertCell(tailCell, tailCellIndex, false);
+					splitCells.push({ cell: this.cells[tailCellIndex], prefix: newlinesBeforeTailCellContent });
 				}
 
 				let activeCell = newCell ? newCell : (headContent.length ? tailCell : this.cells[index]);
 				let activeCellIndex = newCell ? newCellIndex : (headContent.length ? tailCellIndex : index);
 
 				if (addToUndoStack) {
-					let headCell = newCell ? newCell : this.cells[index];
+					this.undoService.pushElement(new SplitCellEdit(this, splitCells));
-					this.undoService.pushElement(new SplitCellEdit(this, headCell, tailCell, newLinesRemoved));
 				}
 				//make new cell Active
 				this.updateActiveCell(activeCell);
@@ -679,19 +687,22 @@ export class NotebookModel extends Disposable implements INotebookModel {
 		return undefined;
 	}
 
-	public mergeCells(cell: ICellModel, secondCell: ICellModel, newLinesRemoved: string[] | undefined): void {
+	public mergeCells(cells: SplitCell[]): void {
-		let index = this._cells.findIndex(cell => cell.equals(cell));
+		let firstCell = cells[0].cell;
-		if (index > -1) {
+		// Append the other cell sources to the first cell
-			cell.source = newLinesRemoved.length > 0 ? [...cell.source, ...newLinesRemoved, ...secondCell.source] : [...cell.source, ...secondCell.source];
+		for (let i = 1; i < cells.length; i++) {
-			cell.isEditMode = true;
+			firstCell.source = cells[i].prefix ? [...firstCell.source, ...cells[i].prefix, ...cells[i].cell.source] : [...firstCell.source, ...cells[i].cell.source];
-			// Set newly created cell as active cell
+		}
-			this.updateActiveCell(cell);
+		firstCell.isEditMode = true;
-			this._contentChangedEmitter.fire({
+		// Set newly created cell as active cell
-				changeType: NotebookChangeType.CellsModified,
+		this.updateActiveCell(firstCell);
-				cells: [cell],
+		this._contentChangedEmitter.fire({
-				cellIndex: index
+			changeType: NotebookChangeType.CellsModified,
-			});
+			cells: [firstCell],
-			this.deleteCell(secondCell, false);
+			cellIndex: 0
+		});
+		for (let i = 1; i < cells.length; i++) {
+			this.deleteCell(cells[i].cell, false);
 		}
 	}