Fix/search text cell on edit (#9685)

* find in text cell changes

* remove prev decorations

* update find index on cell edit

* added test

* addressed comments

* emit error

src/sql/workbench/contrib/notebook/browser/cellViews/textCell.component.ts

@@ -25,8 +25,7 @@ import { NotebookModel } from 'sql/workbench/services/notebook/browser/models/no
 import { ISanitizer, defaultSanitizer } from 'sql/workbench/services/notebook/browser/outputs/sanitizer';
 import { CellToggleMoreActions } from 'sql/workbench/contrib/notebook/browser/cellToggleMoreActions';
 import { CodeComponent } from 'sql/workbench/contrib/notebook/browser/cellViews/code.component';
-import { BaseTextEditor } from 'vs/workbench/browser/parts/editor/textEditor';
-import { NotebookRange } from 'sql/workbench/services/notebook/browser/notebookService';
+import { NotebookRange, ICellEditorProvider } from 'sql/workbench/services/notebook/browser/notebookService';
 import { IColorTheme } from 'vs/platform/theme/common/themeService';
 
 export const TEXT_SELECTOR: string = 'text-cell-component';
@@ -106,6 +105,14 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 		}));
 	}
 
+	public get cellEditors(): ICellEditorProvider[] {
+		let editors: ICellEditorProvider[] = [];
+		if (this.markdowncodeCell) {
+			editors.push(...this.markdowncodeCell.toArray());
+		}
+		return editors;
+	}
+
 	//Gets sanitizer from ISanitizer interface
 	private get sanitizer(): ISanitizer {
 		if (this._sanitizer) {
@@ -121,15 +128,6 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 	get activeCellId(): string {
 		return this._activeCellId;
 	}
-	/**
-	 * Returns the code editor of makrdown cell in edit mode.
-	 */
-	getEditor(): BaseTextEditor | undefined {
-		if (this.markdowncodeCell.length > 0) {
-			return this.markdowncodeCell.first.getEditor();
-		}
-		return undefined;
-	}
 
 	private setLoading(isLoading: boolean): void {
 		this.cellModel.loaded = !isLoading;
@@ -231,6 +229,7 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 
 	public toggleEditMode(editMode?: boolean): void {
 		this.isEditMode = editMode !== undefined ? editMode : !this.isEditMode;
+		this.cellModel.isEditMode = this.isEditMode;
 		this.updateMoreActions();
 		this.updatePreview();
 		this._changeRef.detectChanges();

src/sql/workbench/contrib/notebook/browser/models/notebookFindModel.ts

@@ -48,4 +48,6 @@ export interface INotebookFindModel {
 	findExpression: string;
 	/** Emit event when the find count changes */
 	onFindCountChange: Event<number>;
+	/** Get the find index when range is given*/
+	getIndexByRange(range: NotebookRange): number;
 }

src/sql/workbench/contrib/notebook/browser/notebook.component.ts

@@ -71,7 +71,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 	@ViewChild('bookNav', { read: ElementRef }) private bookNav: ElementRef;
 
 	@ViewChildren(CodeCellComponent) private codeCells: QueryList<CodeCellComponent>;
-	@ViewChildren(TextCellComponent) private textCells: QueryList<ICellEditorProvider>;
+	@ViewChildren(TextCellComponent) private textCells: QueryList<TextCellComponent>;
 
 	private _model: NotebookModel;
 	protected _actionBar: Taskbar;
@@ -150,6 +150,7 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 			this.codeCells.toArray().forEach(cell => editors.push(...cell.cellEditors));
 		}
 		if (this.textCells) {
+			this.textCells.toArray().forEach(cell => editors.push(...cell.cellEditors));
 			editors.push(...this.textCells.toArray());
 		}
 		return editors;
@@ -157,12 +158,12 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 
 	public deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
 		if (newDecorationRange && newDecorationRange.cell && newDecorationRange.cell.cellType === 'markdown') {
-			let cell = this.cellEditors.filter(c => c.cellGuid() === newDecorationRange.cell.cellGuid)[0];
-			cell.deltaDecorations(newDecorationRange, undefined);
+			let cell = this.cellEditors.filter(c => c.cellGuid() === newDecorationRange.cell.cellGuid);
+			cell[cell.length - 1].deltaDecorations(newDecorationRange, undefined);
 		}
 		if (oldDecorationRange && oldDecorationRange.cell && oldDecorationRange.cell.cellType === 'markdown') {
-			let cell = this.cellEditors.filter(c => c.cellGuid() === oldDecorationRange.cell.cellGuid)[0];
-			cell.deltaDecorations(undefined, oldDecorationRange);
+			let cell = this.cellEditors.filter(c => c.cellGuid() === oldDecorationRange.cell.cellGuid);
+			cell[cell.length - 1].deltaDecorations(undefined, oldDecorationRange);
 		}
 	}
 

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

@@ -94,7 +94,9 @@ export class NotebookEditor extends BaseEditor implements IFindNotebookControlle
 		let editorImpl = this._notebookService.findNotebookEditor(this.notebookInput.notebookUri);
 		if (editorImpl) {
 			let cellEditorProvider = editorImpl.cellEditors.filter(c => c.cellGuid() === cellGuid)[0];
-			return cellEditorProvider ? cellEditorProvider.getEditor() : undefined;
+			if (cellEditorProvider) {
+				return cellEditorProvider.getEditor();
+			}
 		}
 		return undefined;
 	}
@@ -269,6 +271,7 @@ export class NotebookEditor extends BaseEditor implements IFindNotebookControlle
 		}
 		if (this._findCountChangeListener === undefined && this._notebookModel) {
 			this._findCountChangeListener = this.notebookInput.notebookFindModel.onFindCountChange(() => this._updateFinderMatchState());
+			this.registerModelChanges();
 		}
 		if (e.isRevealed) {
 			if (this._findState.isRevealed) {
@@ -281,14 +284,20 @@ export class NotebookEditor extends BaseEditor implements IFindNotebookControlle
 				this._finder.getDomNode().style.visibility = 'hidden';
 				this._findDecorations.clearDecorations();
 			}
+		} else {
+			if (!this._findState.isRevealed) {
+				this._finder.getDomNode().style.visibility = 'hidden';
+				this._findDecorations.clearDecorations();
+			}
 		}
 
 		if (e.searchString || e.matchCase || e.wholeWord) {
 			this._findDecorations.clearDecorations();
+			// if the search scope changes remove the prev
 			if (this._notebookModel) {
 				if (this._findState.searchString) {
 					let findScope = this._findDecorations.getFindScope();
-					if (this._findState.searchString === this.notebookFindModel.findExpression && findScope !== null && !e.matchCase && !e.wholeWord) {
+					if (this._findState.searchString === this.notebookFindModel.findExpression && findScope !== null && !e.matchCase && !e.wholeWord && !e.searchScope) {
 						if (findScope) {
 							this._updateFinderMatchState();
 							this._findState.changeMatchInfo(
@@ -328,6 +337,46 @@ export class NotebookEditor extends BaseEditor implements IFindNotebookControlle
 				}
 			}
 		}
+		if (e.searchScope) {
+			await this.notebookInput.notebookFindModel.find(this._findState.searchString, this._findState.matchCase, this._findState.wholeWord, NOTEBOOK_MAX_MATCHES).then(findRange => {
+				this._findDecorations.set(this.notebookFindModel.findMatches, this._currentMatch);
+				this._findState.changeMatchInfo(
+					this.notebookFindModel.getIndexByRange(this._currentMatch),
+					this._findDecorations.getCount(),
+					this._currentMatch
+				);
+				if (this._finder.getDomNode().style.visibility === 'visible') {
+					this._setCurrentFindMatch(this._currentMatch);
+				}
+			});
+		}
+	}
+
+	private registerModelChanges(): void {
+		let changeEvent: FindReplaceStateChangedEvent = {
+			moveCursor: true,
+			updateHistory: true,
+			searchString: false,
+			replaceString: false,
+			isRevealed: false,
+			isReplaceRevealed: false,
+			isRegex: false,
+			wholeWord: false,
+			matchCase: false,
+			preserveCase: false,
+			searchScope: true,
+			matchesPosition: false,
+			matchesCount: false,
+			currentMatch: false
+		};
+		this._notebookModel.cells.forEach(cell => {
+			this._register(cell.onCellModeChanged((state) => {
+				this._onFindStateChange(changeEvent).catch(onUnexpectedError);
+			}));
+		});
+		this._register(this._notebookModel.contentChanged(e => {
+			this._onFindStateChange(changeEvent).catch(onUnexpectedError);
+		}));
 	}
 
 	public setSelection(range: NotebookRange): void {

src/sql/workbench/contrib/notebook/find/notebookFindDecorations.ts

@@ -128,9 +128,9 @@ export class NotebookFindDecorations implements IDisposable {
 
 	private removePrevDecorations(): void {
 		if (this._currentMatch && this._currentMatch.cell) {
-			let pevEditor = this._currentMatch.cell.cellType === 'markdown' ? undefined : this._editor.getCellEditor(this._currentMatch.cell.cellGuid);
-			if (pevEditor) {
-				pevEditor.getControl().changeDecorations((changeAccessor: IModelDecorationsChangeAccessor) => {
+			let prevEditor = this._currentMatch.cell.cellType === 'markdown' && !this._currentMatch.isMarkdownSourceCell ? undefined : this._editor.getCellEditor(this._currentMatch.cell.cellGuid);
+			if (prevEditor) {
+				prevEditor.getControl().changeDecorations((changeAccessor: IModelDecorationsChangeAccessor) => {
 					changeAccessor.removeDecoration(this._rangeHighlightDecorationId);
 					this._rangeHighlightDecorationId = null;
 				});
@@ -153,7 +153,7 @@ export class NotebookFindDecorations implements IDisposable {
 	}
 
 	public checkValidEditor(range: NotebookRange): boolean {
-		return range && range.cell && range.cell.cellType === 'code' && !!(this._editor.getCellEditor(range.cell.cellGuid));
+		return range && range.cell && !!(this._editor.getCellEditor(range.cell.cellGuid)) && (range.cell.cellType === 'code' || range.isMarkdownSourceCell);
 	}
 
 	public set(findMatches: NotebookFindMatch[], findScope: NotebookRange | null): void {

src/sql/workbench/contrib/notebook/find/notebookFindModel.ts

@@ -406,7 +406,8 @@ export class NotebookFindModel extends Disposable implements INotebookFindModel
 			if (oldDecorationIndex < oldDecorationsLen) {
 				// (1) get ourselves an old node
 				do {
-					node = this._decorations[oldDecorationsIds[oldDecorationIndex++]].node;
+					let decorationNode = this._decorations[oldDecorationsIds[oldDecorationIndex++]];
+					node = decorationNode?.node;
 				} while (!node && oldDecorationIndex < oldDecorationsLen);
 
 				// (2) remove the node from the tree (if it exists)
@@ -522,11 +523,29 @@ export class NotebookFindModel extends Disposable implements INotebookFindModel
 	}
 
 	public get findArray(): NotebookRange[] {
-		return this.findArray;
+		return this._findArray;
+	}
+
+	getIndexByRange(range: NotebookRange): number {
+		let index = this.findArray.findIndex(r => r.cell.cellGuid === range.cell.cellGuid && r.startColumn === range.startColumn && r.endColumn === range.endColumn && r.startLineNumber === range.startLineNumber && r.endLineNumber === range.endLineNumber && r.isMarkdownSourceCell === range.isMarkdownSourceCell);
+		this._findIndex = index > -1 ? index : this._findIndex;
+		// _findIndex is the 0 based index, return index + 1 for the actual count on UI
+		return this._findIndex + 1;
 	}
 
 	private searchFn(cell: ICellModel, exp: string, matchCase: boolean = false, wholeWord: boolean = false, maxMatches?: number): NotebookRange[] {
 		let findResults: NotebookRange[] = [];
+		if (cell.cellType === 'markdown' && cell.isEditMode && typeof cell.source !== 'string') {
+			let cellSource = cell.source;
+			for (let j = 0; j < cellSource.length; j++) {
+				let findStartResults = this.search(cellSource[j], exp, matchCase, wholeWord, maxMatches - findResults.length);
+				findStartResults.forEach(start => {
+					// lineNumber: j+1 since notebook editors aren't zero indexed.
+					let range = new NotebookRange(cell, j + 1, start, j + 1, start + exp.length, true);
+					findResults.push(range);
+				});
+			}
+		}
 		let cellVal = cell.cellType === 'markdown' ? cell.renderedOutputTextContent : cell.source;
 		if (cellVal) {
 			if (typeof cellVal === 'string') {
@@ -650,10 +669,10 @@ export class NotebookFindModel extends Disposable implements INotebookFindModel
 
 }
 
-export class NotebookIntervalNode {
+export class NotebookIntervalNode extends IntervalNode {
 
 	constructor(public node: IntervalNode, public cell: ICellModel) {
-
+		super(node.id, node.start, node.end);
 	}
 }
 

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

@@ -833,4 +833,39 @@ suite('Cell Model', function (): void {
 			assert.strictEqual(actualMsg, testMsg);
 		});
 	});
+
+	test('Should emit event on markdown cell edit', async function (): Promise<void> {
+		let notebookModel = new NotebookModelStub({
+			name: '',
+			version: '',
+			mimetype: ''
+		});
+		let contents: nb.ICellContents = {
+			cell_type: CellTypes.Markdown,
+			source: ''
+		};
+		let model = factory.createCell(contents, { notebook: notebookModel, isTrusted: false });
+		assert(!model.isEditMode);
+
+		let createCellModePromise = () => {
+			return new Promise((resolve, reject) => {
+				setTimeout((error) => reject(error), 2000);
+				model.onCellModeChanged(isEditMode => {
+					resolve(isEditMode);
+				});
+			});
+		};
+
+		assert(!model.isEditMode);
+		let cellModePromise = createCellModePromise();
+		model.isEditMode = true;
+		let isEditMode = await cellModePromise;
+		assert(isEditMode);
+
+		cellModePromise = createCellModePromise();
+		model.isEditMode = false;
+		isEditMode = await cellModePromise;
+		assert(!isEditMode);
+	});
+
 });

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

@@ -323,6 +323,41 @@ suite('Notebook Find Model', function (): void {
 		assert.equal(notebookFindModel.findMatches.length, 0, 'Find failed to apply match whole word for //');
 	});
 
+	test('Should find results in the code cell on markdown edit', async function (): Promise<void> {
+		let markdownContent: nb.INotebookContents = {
+			cells: [{
+				cell_type: CellTypes.Markdown,
+				source: ['SOP067 - INTERNAL - Install azdata CLI - release candidate', '==========================================================', 'Steps', '-----', '### Parameters'],
+				metadata: { language: 'python' },
+				execution_count: 1
+			}],
+			metadata: {
+				kernelspec: {
+					name: 'mssql',
+					language: 'sql'
+				}
+			},
+			nbformat: 4,
+			nbformat_minor: 5
+		};
+		await initNotebookModel(markdownContent);
+
+		// Need to set rendered text content for 1st cell
+		setRenderedTextContent(0);
+
+		let notebookFindModel = new NotebookFindModel(model);
+		await notebookFindModel.find('SOP', false, false, max_find_count);
+
+		assert.equal(notebookFindModel.findMatches.length, 1, 'Find failed on markdown');
+
+		// fire the edit mode on cell
+		model.cells[0].isEditMode = true;
+		notebookFindModel = new NotebookFindModel(model);
+		await notebookFindModel.find('SOP', false, false, max_find_count);
+
+		assert.equal(notebookFindModel.findMatches.length, 2, 'Find failed on markdown edit');
+	});
+
 
 	async function initNotebookModel(contents: nb.INotebookContents): Promise<void> {
 		let mockContentManager = TypeMoq.Mock.ofType(NotebookEditorContentManager);

src/sql/workbench/contrib/notebook/test/stubs.ts

@@ -123,7 +123,6 @@ export class NotebookModelStub implements INotebookModel {
 }
 
 export class NotebookFindModelStub implements INotebookFindModel {
-
 	getFindCount(): number {
 		throw new Error('Method not implemented.');
 	}
@@ -158,6 +157,9 @@ export class NotebookFindModelStub implements INotebookFindModel {
 	findMatches: NotebookFindMatch[];
 	findExpression: string;
 	onFindCountChange: vsEvent.Event<number>;
+	getIndexByRange(range: NotebookRange): number {
+		throw new Error('Method not implemented.');
+	}
 }
 
 export class NotebookManagerStub implements INotebookManager {

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

@@ -108,6 +108,10 @@ export class CellModel implements ICellModel {
 		return this._onOutputsChanged.event;
 	}
 
+	public get onCellModeChanged(): Event<boolean> {
+		return this._onCellModeChanged.event;
+	}
+
 	public get isEditMode(): boolean {
 		return this._isEditMode;
 	}

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

@@ -484,6 +484,7 @@ export interface ICellModel {
 	readonly onLoaded: Event<string>;
 	isCollapsed: boolean;
 	readonly onCollapseStateChanged: Event<boolean>;
+	readonly onCellModeChanged: Event<boolean>;
 	modelContentChangedEvent: IModelContentChangedEvent;
 	isEditMode: boolean;
 	readonly ariaLabel: string;

src/sql/workbench/services/notebook/browser/notebookService.ts

@@ -179,10 +179,12 @@ export class NotebookRange extends Range {
 		this.cell = cell;
 	}
 	cell: ICellModel;
+	isMarkdownSourceCell: boolean;
 
-	constructor(cell: ICellModel, startLineNumber: number, startColumn: number, endLineNumber: number, endColumn: number) {
+	constructor(cell: ICellModel, startLineNumber: number, startColumn: number, endLineNumber: number, endColumn: number, markdownEditMode?: boolean) {
 		super(startLineNumber, startColumn, endLineNumber, endColumn);
 		this.updateActiveCell(cell);
+		this.isMarkdownSourceCell = markdownEditMode ? markdownEditMode : false;
 	}
 }