Highlight all matches in Notebook on Find (#15562)

* iterate over every cell and highlight ranges.

* fix yellow for all matches and orange for current

* fix

* avoid duplicate deltaDecorations call

* initialize on declare

src/sql/workbench/contrib/notebook/browser/cellViews/interfaces.ts

@@ -28,7 +28,7 @@ export abstract class CellView extends AngularDisposable implements OnDestroy, I
 
 	public abstract cellGuid(): string;
 
-	public deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
+	public deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
 
 	}
 }

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

@@ -102,6 +102,7 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 	public previewFeaturesEnabled: boolean = false;
 	public doubleClickEditEnabled: boolean;
 	private _highlightRange: NotebookRange;
+	private _isFindActive: boolean = false;
 
 	constructor(
 		@Inject(forwardRef(() => ChangeDetectorRef)) private _changeRef: ChangeDetectorRef,
@@ -247,7 +248,9 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 				outputElement.style.lineHeight = this.markdownPreviewLineHeight.toString();
 				this.cellModel.renderedOutputTextContent = this.getRenderedTextOutput();
 				outputElement.focus();
-				this.addDecoration();
+				if (this._isFindActive) {
+					this.addDecoration();
+				}
 			}
 		}
 	}
@@ -353,58 +356,81 @@ export class TextCellComponent extends CellView implements OnInit, OnChanges {
 		return this.cellModel && this.cellModel.id === this.activeCellId;
 	}
 
-	public deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
-		if (oldDecorationRange) {
-			this._highlightRange = oldDecorationRange === this._highlightRange ? undefined : this._highlightRange;
-			this.removeDecoration(oldDecorationRange);
+	public deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
+		if (newDecorationsRange) {
+			this._isFindActive = true;
+			if (Array.isArray(newDecorationsRange)) {
+				this.highlightAllMatches();
+			} else {
+				this._highlightRange = newDecorationsRange;
+				this.addDecoration(newDecorationsRange);
+			}
 		}
-
-		if (newDecorationRange) {
-			this._highlightRange = newDecorationRange;
-			this.addDecoration(newDecorationRange);
+		if (oldDecorationsRange) {
+			if (Array.isArray(oldDecorationsRange)) {
+				this.removeDecoration();
+				this._isFindActive = false;
+			} else {
+				this._highlightRange = oldDecorationsRange === this._highlightRange ? undefined : this._highlightRange;
+				this.removeDecoration(oldDecorationsRange);
+			}
 		}
 	}
 
 	private addDecoration(range?: NotebookRange): void {
 		range = range ?? this._highlightRange;
-		if (range && this.output && this.output.nativeElement) {
-			let markAllOccurances = new Mark(this.output.nativeElement); // to highlight all occurances in the element.
-			let elements = this.getHtmlElements();
-			if (elements?.length >= range.startLineNumber) {
-				let elementContainingText = elements[range.startLineNumber - 1];
-				let markCurrent = new Mark(elementContainingText); // to highlight the current item of them all.
-				let editor = this._notebookService.findNotebookEditor(this.model.notebookUri);
-				if (editor) {
-					let findModel = (editor.notebookParams.input as NotebookInput).notebookFindModel;
-					if (findModel?.findMatches?.length > 0) {
-						let searchString = findModel.findExpression;
-						markAllOccurances.mark(searchString, {
-							className: findHighlightClass
-						});
-					}
+		if (this.output && this.output.nativeElement) {
+			this.highlightAllMatches();
+			if (range) {
+				let elements = this.getHtmlElements();
+				if (elements?.length >= range.startLineNumber) {
+					let elementContainingText = elements[range.startLineNumber - 1];
+					let markCurrent = new Mark(elementContainingText); // to highlight the current item of them all.
+
+					markCurrent.markRanges([{
+						start: range.startColumn - 1, //subtracting 1 since markdown html is 0 indexed.
+						length: range.endColumn - range.startColumn
+					}], {
+						className: findRangeSpecificClass,
+						each: function (node, range) {
+							// node is the marked DOM element
+							node.scrollIntoView({ behavior: 'smooth', block: 'center' });
+						}
+					});
 				}
-				markCurrent.markRanges([{
-					start: range.startColumn - 1, //subtracting 1 since markdown html is 0 indexed.
-					length: range.endColumn - range.startColumn
-				}], {
-					className: findRangeSpecificClass,
-					each: function (node, range) {
-						// node is the marked DOM element
-						node.scrollIntoView({ behavior: 'smooth', block: 'center' });
-					}
-				});
 			}
 		}
 	}
 
-	private removeDecoration(range: NotebookRange): void {
-		if (range && this.output && this.output.nativeElement) {
-			let markAllOccurances = new Mark(this.output.nativeElement);
-			let elements = this.getHtmlElements();
-			let elementContainingText = elements[range.startLineNumber - 1];
-			let markCurrent = new Mark(elementContainingText);
-			markAllOccurances.unmark({ acrossElements: true, className: findHighlightClass });
-			markCurrent.unmark({ acrossElements: true, className: findRangeSpecificClass });
+	private highlightAllMatches(): void {
+		if (this.output && this.output.nativeElement) {
+			let markAllOccurances = new Mark(this.output.nativeElement); // to highlight all occurances in the element.
+			let editor = this._notebookService.findNotebookEditor(this.model.notebookUri);
+			if (editor) {
+				let findModel = (editor.notebookParams.input as NotebookInput).notebookFindModel;
+				if (findModel?.findMatches?.length > 0) {
+					let searchString = findModel.findExpression;
+					markAllOccurances.mark(searchString, {
+						className: findHighlightClass
+					});
+				}
+			}
+		}
+	}
+
+	private removeDecoration(range?: NotebookRange): void {
+		if (this.output && this.output.nativeElement) {
+			if (range) {
+				let elements = this.getHtmlElements();
+				let elementContainingText = elements[range.startLineNumber - 1];
+				let markCurrent = new Mark(elementContainingText);
+				markCurrent.unmark({ acrossElements: true, className: findRangeSpecificClass });
+			} else {
+				let markAllOccurances = new Mark(this.output.nativeElement);
+				markAllOccurances.unmark({ acrossElements: true, className: findHighlightClass });
+				markAllOccurances.unmark({ acrossElements: true, className: findRangeSpecificClass });
+				this._highlightRange = undefined;
+			}
 		}
 	}
 

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

@@ -15,20 +15,17 @@ import { NotebookRange } from 'sql/workbench/services/notebook/browser/notebookS
 
 export class NotebookFindDecorations implements IDisposable {
 
-	private _decorations: string[];
-	private _overviewRulerApproximateDecorations: string[];
-	private _findScopeDecorationId: string | null;
-	private _rangeHighlightDecorationId: string | null;
-	private _highlightedDecorationId: string | null;
+	private _decorations: string[] = [];
+	private _overviewRulerApproximateDecorations: string[] = [];
+	private _findScopeDecorationIds: string[] = [];
+	private _codeCellFindScopeDecorationIds: string[] = [];
+	private _rangeHighlightDecorationId: string | null = null;
+	private _highlightedDecorationId: string | null = null;
 	private _startPosition: NotebookRange;
 	private _currentMatch: NotebookRange;
+	private _codeCellDecorations: Map<string, string[]> = new Map<string, string[]>();
 
 	constructor(private readonly _editor: NotebookEditor) {
-		this._decorations = [];
-		this._overviewRulerApproximateDecorations = [];
-		this._findScopeDecorationId = null;
-		this._rangeHighlightDecorationId = null;
-		this._highlightedDecorationId = null;
 		this._startPosition = this._editor.getPosition();
 	}
 
@@ -37,7 +34,9 @@ export class NotebookFindDecorations implements IDisposable {
 
 		this._decorations = [];
 		this._overviewRulerApproximateDecorations = [];
-		this._findScopeDecorationId = null;
+		this._findScopeDecorationIds = [];
+		this._codeCellDecorations = new Map<string, string[]>();
+		this._codeCellFindScopeDecorationIds = [];
 		this._rangeHighlightDecorationId = null;
 		this._highlightedDecorationId = null;
 	}
@@ -45,7 +44,9 @@ export class NotebookFindDecorations implements IDisposable {
 	public reset(): void {
 		this._decorations = [];
 		this._overviewRulerApproximateDecorations = [];
-		this._findScopeDecorationId = null;
+		this._findScopeDecorationIds = [];
+		this._codeCellDecorations = new Map<string, string[]>();
+		this._codeCellFindScopeDecorationIds = [];
 		this._rangeHighlightDecorationId = null;
 		this._highlightedDecorationId = null;
 	}
@@ -61,6 +62,18 @@ export class NotebookFindDecorations implements IDisposable {
 		return null;
 	}
 
+	public getFindScopes(): NotebookRange[] | null {
+		if (this._findScopeDecorationIds.length) {
+			const scopes = this._findScopeDecorationIds.map(findScopeDecorationId =>
+				this._editor.notebookFindModel.getDecorationRange(findScopeDecorationId)
+			).filter(element => !!element);
+			if (scopes.length) {
+				return scopes as NotebookRange[];
+			}
+		}
+		return null;
+	}
+
 	public getStartPosition(): NotebookRange {
 		return this._startPosition;
 	}
@@ -73,7 +86,31 @@ export class NotebookFindDecorations implements IDisposable {
 	}
 
 	public clearDecorations(): void {
-		this.removePrevDecorations();
+		// clear markdown decorations
+		let ranges = this.getFindScopes();
+		if (ranges) {
+			this._editor.updateDecorations(undefined, ranges);
+		}
+		// clear code cell decorations
+		for (let cellGuid of this._codeCellDecorations.keys()) {
+			this._editor.getCellEditor(cellGuid).getControl().changeDecorations((changeAccessor: IModelDecorationsChangeAccessor) => {
+				this._codeCellDecorations.get(cellGuid).forEach(decorationId => changeAccessor.removeDecoration(decorationId));
+				changeAccessor.deltaDecorations(this._codeCellFindScopeDecorationIds, []);
+			});
+		}
+		// remove the current highlight
+		this.removeLastDecoration();
+	}
+
+	public addDecorations(): void {
+		let findScopes = this.getFindScopes();
+		if (findScopes) {
+			// add markdown decorations
+			this._editor.updateDecorations(findScopes, undefined);
+			// add code cell decorations
+			this.setCodeCellDecorations(this._editor.notebookFindModel.findMatches, findScopes);
+		}
+
 	}
 
 	public setCurrentFindMatch(nextMatch: NotebookRange | null): number {
@@ -84,7 +121,6 @@ export class NotebookFindDecorations implements IDisposable {
 				let range = this._editor.notebookFindModel.getDecorationRange(this._decorations[i]);
 				if (nextMatch.equalsRange(range)) {
 					newCurrentDecorationId = this._decorations[i];
-					this._findScopeDecorationId = newCurrentDecorationId;
 					matchPosition = (i + 1);
 					break;
 				}
@@ -92,11 +128,11 @@ export class NotebookFindDecorations implements IDisposable {
 		}
 
 		if (this._highlightedDecorationId !== null || newCurrentDecorationId !== null) {
-			this.removePrevDecorations();
+			this.removeLastDecoration();
 			if (this.checkValidEditor(nextMatch)) {
 				this._editor.getCellEditor(nextMatch.cell.cellGuid).getControl().changeDecorations((changeAccessor: IModelDecorationsChangeAccessor) => {
 					if (this._highlightedDecorationId !== null) {
-						changeAccessor.changeDecorationOptions(this._highlightedDecorationId, NotebookFindDecorations._FIND_MATCH_DECORATION);
+						changeAccessor.changeDecorationOptions(this._highlightedDecorationId, NotebookFindDecorations._RANGE_HIGHLIGHT_DECORATION);
 						this._highlightedDecorationId = null;
 					}
 					if (newCurrentDecorationId !== null) {
@@ -111,7 +147,7 @@ export class NotebookFindDecorations implements IDisposable {
 							let lineBeforeEndMaxColumn = this._editor.notebookFindModel.getLineMaxColumn(lineBeforeEnd);
 							rng = new NotebookRange(rng.cell, rng.startLineNumber, rng.startColumn, lineBeforeEnd, lineBeforeEndMaxColumn);
 						}
-						this._rangeHighlightDecorationId = changeAccessor.addDecoration(rng, NotebookFindDecorations._RANGE_HIGHLIGHT_DECORATION);
+						this._rangeHighlightDecorationId = changeAccessor.addDecoration(rng, NotebookFindDecorations._FIND_MATCH_DECORATION);
 						this._revealRangeInCenterIfOutsideViewport(nextMatch);
 						this._currentMatch = nextMatch;
 					}
@@ -126,7 +162,7 @@ export class NotebookFindDecorations implements IDisposable {
 		return matchPosition;
 	}
 
-	private removePrevDecorations(): void {
+	private removeLastDecoration(): void {
 		if (this._currentMatch && this._currentMatch.cell) {
 			let prevEditor = this._currentMatch.cell.cellType === 'markdown' && !this._currentMatch.isMarkdownSourceCell ? undefined : this._editor.getCellEditor(this._currentMatch.cell.cellGuid);
 			if (prevEditor) {
@@ -156,76 +192,115 @@ export class NotebookFindDecorations implements IDisposable {
 		return range && range.cell && !!(this._editor.getCellEditor(range.cell.cellGuid)) && (range.cell.cellType === 'code' || range.isMarkdownSourceCell);
 	}
 
-	public set(findMatches: NotebookFindMatch[], findScope: NotebookRange | null): void {
-		this._editor.changeDecorations((accessor) => {
+	public set(findMatches: NotebookFindMatch[], findScopes: NotebookRange[] | null): void {
+		if (findScopes) {
+			this._editor.updateDecorations(findScopes, undefined);
 
-			let findMatchesOptions: ModelDecorationOptions = NotebookFindDecorations._FIND_MATCH_DECORATION;
-			let newOverviewRulerApproximateDecorations: IModelDeltaDecoration[] = [];
+			this._editor.changeDecorations((accessor) => {
+				let findMatchesOptions = NotebookFindDecorations._FIND_MATCH_NO_OVERVIEW_DECORATION;
+				// Find matches
+				let newFindMatchesDecorations: IModelDeltaDecoration[] = new Array<IModelDeltaDecoration>(findMatches.length);
+				for (let i = 0, len = findMatches.length; i < len; i++) {
+					newFindMatchesDecorations[i] = {
+						range: findMatches[i].range,
+						options: findMatchesOptions
+					};
+				}
+				this._decorations = accessor.deltaDecorations(this._decorations, newFindMatchesDecorations);
 
-			if (findMatches.length > 1000) {
-				// we go into a mode where the overview ruler gets "approximate" decorations
-				// the reason is that the overview ruler paints all the decorations in the file and we don't want to cause freezes
-				findMatchesOptions = NotebookFindDecorations._FIND_MATCH_NO_OVERVIEW_DECORATION;
+				// Find scope
+				if (this._findScopeDecorationIds.length) {
+					this._findScopeDecorationIds.forEach(findScopeDecorationId => accessor.removeDecoration(findScopeDecorationId));
+					this._findScopeDecorationIds = [];
+				}
+				if (findScopes.length) {
+					this._findScopeDecorationIds = findScopes.map(findScope => accessor.addDecoration(findScope, NotebookFindDecorations._FIND_SCOPE_DECORATION));
+				}
+			});
 
-				// approximate a distance in lines where matches should be merged
-				const lineCount = this._editor.notebookFindModel.getLineCount();
-				const height = this._editor.getConfiguration().layoutInfo.height;
-				const approxPixelsPerLine = height / lineCount;
-				const mergeLinesDelta = Math.max(2, Math.ceil(3 / approxPixelsPerLine));
+			this.setCodeCellDecorations(findMatches, findScopes);
+		}
+	}
 
-				// merge decorations as much as possible
-				let prevStartLineNumber = findMatches[0].range.startLineNumber;
-				let prevEndLineNumber = findMatches[0].range.endLineNumber;
-				for (let i = 1, len = findMatches.length; i < len; i++) {
-					const range: NotebookRange = findMatches[i].range;
-					if (prevEndLineNumber + mergeLinesDelta >= range.startLineNumber) {
-						if (range.endLineNumber > prevEndLineNumber) {
+	private setCodeCellDecorations(findMatches: NotebookFindMatch[], findScopes: NotebookRange[] | null): void {
+		//get all code cells which have matches
+		const codeCellsFindMatches = findScopes.filter((c, i, ranges) => {
+			return ranges.indexOf(ranges.find(t => t.cell.cellGuid === c.cell.cellGuid && t.cell.cellType === 'code')) === i;
+		});
+		codeCellsFindMatches.forEach(findMatch => {
+			this._editor.getCellEditor(findMatch.cell.cellGuid)?.getControl().changeDecorations((accessor) => {
+
+				let findMatchesOptions: ModelDecorationOptions = NotebookFindDecorations._RANGE_HIGHLIGHT_DECORATION;
+				let newOverviewRulerApproximateDecorations: IModelDeltaDecoration[] = [];
+
+				let cellFindScopes = findScopes.filter(f => f.cell.cellGuid === findMatch.cell.cellGuid);
+				let findMatchesInCell = findMatches?.filter(m => m.range.cell.cellGuid === findMatch.cell.cellGuid) || [];
+				let _cellFindScopeDecorationIds: string[] = [];
+				if (findMatchesInCell.length > 1000) {
+					// we go into a mode where the overview ruler gets "approximate" decorations
+					// the reason is that the overview ruler paints all the decorations in the file and we don't want to cause freezes
+					findMatchesOptions = NotebookFindDecorations._FIND_MATCH_NO_OVERVIEW_DECORATION;
+
+					// approximate a distance in lines where matches should be merged
+					const lineCount = this._editor.notebookFindModel.getLineCount();
+					const height = this._editor.getConfiguration().layoutInfo.height;
+					const approxPixelsPerLine = height / lineCount;
+					const mergeLinesDelta = Math.max(2, Math.ceil(3 / approxPixelsPerLine));
+
+					// merge decorations as much as possible
+					let prevStartLineNumber = findMatchesInCell[0].range.startLineNumber;
+					let prevEndLineNumber = findMatchesInCell[0].range.endLineNumber;
+					for (let i = 1, len = findMatchesInCell.length; i < len; i++) {
+						const range = findMatchesInCell[i].range;
+						if (prevEndLineNumber + mergeLinesDelta >= range.startLineNumber) {
+							if (range.endLineNumber > prevEndLineNumber) {
+								prevEndLineNumber = range.endLineNumber;
+							}
+						} else {
+							newOverviewRulerApproximateDecorations.push({
+								range: new Range(prevStartLineNumber, 1, prevEndLineNumber, 1),
+								options: NotebookFindDecorations._FIND_MATCH_ONLY_OVERVIEW_DECORATION
+							});
+							prevStartLineNumber = range.startLineNumber;
 							prevEndLineNumber = range.endLineNumber;
 						}
-					} else {
-						newOverviewRulerApproximateDecorations.push({
-							range: new NotebookRange(range.cell, prevStartLineNumber, 1, prevEndLineNumber, 1),
-							options: NotebookFindDecorations._FIND_MATCH_ONLY_OVERVIEW_DECORATION
-						});
-						prevStartLineNumber = range.startLineNumber;
-						prevEndLineNumber = range.endLineNumber;
 					}
+
+					newOverviewRulerApproximateDecorations.push({
+						range: new Range(prevStartLineNumber, 1, prevEndLineNumber, 1),
+						options: NotebookFindDecorations._FIND_MATCH_ONLY_OVERVIEW_DECORATION
+					});
 				}
 
-				newOverviewRulerApproximateDecorations.push({
-					range: new NotebookRange(findMatches[0].range.cell, prevStartLineNumber, 1, prevEndLineNumber, 1),
-					options: NotebookFindDecorations._FIND_MATCH_ONLY_OVERVIEW_DECORATION
-				});
-			}
+				// Find matches
+				let newFindMatchesDecorations: IModelDeltaDecoration[] = new Array<IModelDeltaDecoration>(findMatchesInCell.length);
+				for (let i = 0, len = findMatchesInCell.length; i < len; i++) {
+					newFindMatchesDecorations[i] = {
+						range: findMatchesInCell[i].range,
+						options: findMatchesOptions
+					};
+				}
+				let decorations = accessor.deltaDecorations(this._decorations, newFindMatchesDecorations);
+				this._codeCellDecorations.set(findMatch.cell.cellGuid, decorations);
+				// Overview ruler approximate decorations
+				this._overviewRulerApproximateDecorations = accessor.deltaDecorations(this._overviewRulerApproximateDecorations, newOverviewRulerApproximateDecorations);
 
-			// Find matches
-			let newFindMatchesDecorations: IModelDeltaDecoration[] = new Array<IModelDeltaDecoration>(findMatches.length);
-			for (let i = 0, len = findMatches.length; i < len; i++) {
-				newFindMatchesDecorations[i] = {
-					range: findMatches[i].range,
-					options: findMatchesOptions
-				};
-			}
-			this._decorations = accessor.deltaDecorations(this._decorations, newFindMatchesDecorations);
+				// Range highlight
+				if (this._rangeHighlightDecorationId) {
+					accessor.removeDecoration(this._rangeHighlightDecorationId);
+					this._rangeHighlightDecorationId = null;
+				}
 
-			// Overview ruler approximate decorations
-			this._overviewRulerApproximateDecorations = accessor.deltaDecorations(this._overviewRulerApproximateDecorations, newOverviewRulerApproximateDecorations);
-
-			// Range highlight
-			if (this._rangeHighlightDecorationId) {
-				accessor.removeDecoration(this._rangeHighlightDecorationId);
-				this._rangeHighlightDecorationId = null;
-			}
-
-			// Find scope
-			if (this._findScopeDecorationId) {
-				accessor.removeDecoration(this._findScopeDecorationId);
-				this._findScopeDecorationId = null;
-			}
-			if (findScope) {
-				this._currentMatch = findScope;
-				this._findScopeDecorationId = accessor.addDecoration(findScope, NotebookFindDecorations._FIND_SCOPE_DECORATION);
-			}
+				// Find scope
+				if (_cellFindScopeDecorationIds.length) {
+					_cellFindScopeDecorationIds.forEach(findScopeDecorationId => accessor.removeDecoration(findScopeDecorationId));
+					_cellFindScopeDecorationIds = [];
+				}
+				if (cellFindScopes.length) {
+					_cellFindScopeDecorationIds = cellFindScopes.map(findScope => accessor.addDecoration(findScope, NotebookFindDecorations._FIND_SCOPE_DECORATION));
+					this._codeCellFindScopeDecorationIds.push(..._cellFindScopeDecorationIds);
+				}
+			});
 		});
 	}
 
@@ -233,8 +308,8 @@ export class NotebookFindDecorations implements IDisposable {
 		let result: string[] = [];
 		result = result.concat(this._decorations);
 		result = result.concat(this._overviewRulerApproximateDecorations);
-		if (this._findScopeDecorationId) {
-			result.push(this._findScopeDecorationId);
+		if (this._findScopeDecorationIds.length) {
+			result.push(...this._findScopeDecorationIds);
 		}
 		if (this._rangeHighlightDecorationId) {
 			result.push(this._rangeHighlightDecorationId);

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

@@ -518,7 +518,7 @@ export class NotebookFindModel extends Disposable implements INotebookFindModel
 
 	public get findMatches(): NotebookFindMatch[] {
 		let findMatches: NotebookFindMatch[] = [];
-		this._findArray.forEach(element => {
+		this._findArray?.forEach(element => {
 			findMatches = findMatches.concat(new NotebookFindMatch(element, null));
 		});
 		return findMatches;

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

@@ -154,14 +154,40 @@ export class NotebookComponent extends AngularDisposable implements OnInit, OnDe
 		return editors;
 	}
 
-	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);
-			cell[cell.length - 1].deltaDecorations(newDecorationRange, undefined);
+	public deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
+		if (oldDecorationsRange) {
+			if (Array.isArray(oldDecorationsRange)) {
+				let decoratedCells: string[] = [];
+				oldDecorationsRange.forEach(oldDecorationRange => {
+					if (oldDecorationRange.cell.cellType === 'markdown' && decoratedCells.indexOf(oldDecorationRange.cell.cellGuid) === -1) {
+						let cell = this.cellEditors.filter(c => c.cellGuid() === oldDecorationRange.cell.cellGuid);
+						cell[cell.length - 1].deltaDecorations(undefined, [oldDecorationRange]);
+						decoratedCells.push(...oldDecorationRange.cell.cellGuid);
+					}
+				});
+			} else {
+				if (oldDecorationsRange.cell.cellType === 'markdown') {
+					let cell = this.cellEditors.filter(c => c.cellGuid() === oldDecorationsRange.cell.cellGuid);
+					cell[cell.length - 1].deltaDecorations(undefined, oldDecorationsRange);
+				}
+			}
 		}
-		if (oldDecorationRange && oldDecorationRange.cell && oldDecorationRange.cell.cellType === 'markdown') {
-			let cell = this.cellEditors.filter(c => c.cellGuid() === oldDecorationRange.cell.cellGuid);
-			cell[cell.length - 1].deltaDecorations(undefined, oldDecorationRange);
+		if (newDecorationsRange) {
+			if (Array.isArray(newDecorationsRange)) {
+				let decoratedCells: string[] = [];
+				newDecorationsRange.forEach(newDecorationRange => {
+					if (newDecorationRange.cell.cellType === 'markdown' && decoratedCells.indexOf(newDecorationRange.cell.cellGuid) === -1) {
+						let cell = this.cellEditors.filter(c => c.cellGuid() === newDecorationRange.cell.cellGuid);
+						cell[cell.length - 1].deltaDecorations([newDecorationRange], undefined);
+						decoratedCells.push(...newDecorationRange.cell.cellGuid);
+					}
+				});
+			} else {
+				if (newDecorationsRange.cell.cellType === 'markdown') {
+					let cell = this.cellEditors.filter(c => c.cellGuid() === newDecorationsRange.cell.cellGuid);
+					cell[cell.length - 1].deltaDecorations(newDecorationsRange, undefined);
+				}
+			}
 		}
 	}
 

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

@@ -103,10 +103,10 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 
 	// updateDecorations is only used for modifying decorations on markdown cells
 	// changeDecorations is the function that handles the decorations w.r.t codeEditor cells.
-	public updateDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
+	public updateDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
 		let editorImpl = this._notebookService.findNotebookEditor(this.notebookInput.notebookUri);
 		if (editorImpl) {
-			editorImpl.deltaDecorations(newDecorationRange, oldDecorationRange);
+			editorImpl.deltaDecorations(newDecorationsRange, oldDecorationsRange);
 		}
 	}
 
@@ -282,6 +282,7 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 				this._finder.focusFindInput();
 				this._updateFinderMatchState();
 				// if find is closed and opened again, highlight the last position.
+				this._findDecorations.addDecorations();
 				this._findDecorations.setStartPosition(this.getPosition());
 			} else {
 				this._finder.getDomNode().style.visibility = 'hidden';
@@ -325,7 +326,7 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 					}
 					this._updateFinderMatchState();
 					this._finder.focusFindInput();
-					this._findDecorations.set(this.notebookFindModel.findMatches, this._currentMatch);
+					this._findDecorations.set(this.notebookFindModel.findMatches, this.notebookFindModel.findArray);
 					this._findState.changeMatchInfo(
 						this.notebookFindModel.getFindIndex(),
 						this._findDecorations.getCount(),
@@ -339,7 +340,7 @@ export class NotebookEditor extends EditorPane implements IFindNotebookControlle
 		}
 		if (e.searchScope) {
 			await this.notebookInput.notebookFindModel.find(this._findState.searchString, this._findState.matchCase, this._findState.wholeWord, NOTEBOOK_MAX_MATCHES);
-			this._findDecorations.set(this.notebookFindModel.findMatches, this._currentMatch);
+			this._findDecorations.set(this.notebookFindModel.findMatches, this.notebookFindModel.findArray);
 			this._findState.changeMatchInfo(
 				this.notebookFindModel.getIndexByRange(this._currentMatch),
 				this._findDecorations.getCount(),

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

@@ -472,7 +472,7 @@ export class FutureStub implements nb.IFuture {
 export class NotebookComponentStub implements INotebookEditor {
 	cellEditors: ICellEditorProvider[];
 	viewMode: string;
-	deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
+	deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
 		throw new Error('Method not implemented.');
 	}
 	get notebookParams(): INotebookParams {
@@ -715,7 +715,7 @@ export class NotebookEditorStub implements INotebookEditor {
 	navigateToSection(sectionId: string): void {
 		throw new Error('Method not implemented.');
 	}
-	deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
+	deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
 		throw new Error('Method not implemented.');
 	}
 	addCell(cellType: CellType, index?: number, event?: UIEvent) {
@@ -733,7 +733,7 @@ export class CellEditorProviderStub implements ICellEditorProvider {
 	getEditor(): QueryTextEditor {
 		throw new Error('Method not implemented.');
 	}
-	deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void {
+	deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void {
 		throw new Error('Method not implemented.');
 	}
 }

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

@@ -186,7 +186,7 @@ export interface ICellEditorProvider {
 	hasEditor(): boolean;
 	cellGuid(): string;
 	getEditor(): BaseTextEditor;
-	deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void;
+	deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void;
 }
 
 export class NotebookRange extends Range {
@@ -220,7 +220,7 @@ export interface INotebookEditor {
 	clearAllOutputs(): Promise<boolean>;
 	getSections(): INotebookSection[];
 	navigateToSection(sectionId: string): void;
-	deltaDecorations(newDecorationRange: NotebookRange, oldDecorationRange: NotebookRange): void;
+	deltaDecorations(newDecorationsRange: NotebookRange | NotebookRange[], oldDecorationsRange: NotebookRange | NotebookRange[]): void;
 	addCell(cellType: CellType, index?: number, event?: UIEvent);
 }