Add search for output rich content in notebooks

src/vs/workbench/contrib/notebook/browser/contrib/find/notebookFindReplaceWidget.ts

@@ -38,6 +38,7 @@ import { ISashEvent, Orientation, Sash } from 'vs/base/browser/ui/sash/sash';
 import { INotebookEditor } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { defaultInputBoxStyles, defaultProgressBarStyles, defaultToggleStyles } from 'vs/platform/theme/browser/defaultStyles';
 import { IToggleStyles } from 'vs/base/browser/ui/toggle/toggle';
+import { Disposable } from 'vs/base/common/lifecycle';
 
 const NLS_FIND_INPUT_LABEL = nls.localize('label.find', "Find");
 const NLS_FIND_INPUT_PLACEHOLDER = nls.localize('placeholder.find', "Find");
@@ -150,31 +151,26 @@ class NotebookFindFilterActionViewItem extends DropdownMenuActionViewItem {
 	}
 }
 
-class NotebookFindInput extends FindInput {
+export class NotebookFindInputFilter extends Disposable {
 	private _filterButtonContainer: HTMLElement;
 	private _actionbar: ActionBar | null = null;
-	private _filterChecked: boolean = false;
 	private _filtersAction: IAction;
 	private _toggleStyles: IToggleStyles;
 
 	constructor(
 		readonly filters: NotebookFindFilters,
-		contextKeyService: IContextKeyService,
 		readonly contextMenuService: IContextMenuService,
 		readonly instantiationService: IInstantiationService,
-		parent: HTMLElement | null,
-		contextViewProvider: IContextViewProvider,
-		options: IFindInputOptions
+		options: IFindInputOptions,
+		tooltip: string = NOTEBOOK_FIND_FILTERS,
 	) {
-		super(parent, contextViewProvider, options);
 
+		super();
 		this._toggleStyles = options.toggleStyles;
 
-		this._register(registerAndCreateHistoryNavigationContext(contextKeyService, this.inputBox));
-		this._filtersAction = new Action('notebookFindFilterAction', NOTEBOOK_FIND_FILTERS, 'notebook-filters ' + ThemeIcon.asClassName(filterIcon));
+		this._filtersAction = new Action('notebookFindFilterAction', tooltip, 'notebook-filters ' + ThemeIcon.asClassName(filterIcon));
 		this._filtersAction.checked = false;
 		this._filterButtonContainer = dom.$('.find-filter-button');
-		this.controls.appendChild(this._filterButtonContainer);
 		this.createFilters(this._filterButtonContainer);
 
 		this._register(this.filters.onDidChange(() => {
@@ -184,14 +180,24 @@ class NotebookFindInput extends FindInput {
 				this._filtersAction.checked = false;
 			}
 		}));
+	}
 
-		this.inputBox.paddingRight = (this.caseSensitive?.width() ?? 0) + (this.wholeWords?.width() ?? 0) + (this.regex?.width() ?? 0) + this.getFilterWidth();
+	get container() {
+		return this._filterButtonContainer;
 	}
 
-	private getFilterWidth() {
+	get width() {
 		return 2 /*margin left*/ + 2 /*border*/ + 2 /*padding*/ + 16 /* icon width */;
 	}
 
+	applyStyles(filterChecked: boolean): void {
+		const toggleStyles = this._toggleStyles;
+
+		this._filterButtonContainer.style.borderColor = (filterChecked && toggleStyles.inputActiveOptionBorder) || '';
+		this._filterButtonContainer.style.color = (filterChecked && toggleStyles.inputActiveOptionForeground) || 'inherit';
+		this._filterButtonContainer.style.backgroundColor = (filterChecked && toggleStyles.inputActiveOptionBackground) || '';
+	}
+
 	private createFilters(container: HTMLElement): void {
 		this._actionbar = this._register(new ActionBar(container, {
 			actionViewItemProvider: action => {
@@ -203,6 +209,29 @@ class NotebookFindInput extends FindInput {
 		}));
 		this._actionbar.push(this._filtersAction, { icon: true, label: false });
 	}
+}
+
+export class NotebookFindInput extends FindInput {
+	private _findFilter: NotebookFindInputFilter;
+	private _filterChecked: boolean = false;
+
+	constructor(
+		readonly filters: NotebookFindFilters,
+		contextKeyService: IContextKeyService,
+		readonly contextMenuService: IContextMenuService,
+		readonly instantiationService: IInstantiationService,
+		parent: HTMLElement | null,
+		contextViewProvider: IContextViewProvider,
+		options: IFindInputOptions,
+	) {
+		super(parent, contextViewProvider, options);
+
+		this._register(registerAndCreateHistoryNavigationContext(contextKeyService, this.inputBox));
+		this._findFilter = this._register(new NotebookFindInputFilter(filters, contextMenuService, instantiationService, options));
+
+		this.inputBox.paddingRight = (this.caseSensitive?.width() ?? 0) + (this.wholeWords?.width() ?? 0) + (this.regex?.width() ?? 0) + this._findFilter.width;
+		this.controls.appendChild(this._findFilter.container);
+	}
 
 	override setEnabled(enabled: boolean) {
 		super.setEnabled(enabled);
@@ -226,15 +255,7 @@ class NotebookFindInput extends FindInput {
 				this.regex.domNode.classList.toggle('disabled', false);
 			}
 		}
-		this.applyStyles();
-	}
-
-	private applyStyles(): void {
-		const toggleStyles = this._toggleStyles;
-
-		this._filterButtonContainer.style.borderColor = (this._filterChecked && toggleStyles.inputActiveOptionBorder) || '';
-		this._filterButtonContainer.style.color = (this._filterChecked && toggleStyles.inputActiveOptionForeground) || 'inherit';
-		this._filterButtonContainer.style.backgroundColor = (this._filterChecked && toggleStyles.inputActiveOptionBackground) || '';
+		this._findFilter.applyStyles(this._filterChecked);
 	}
 
 	getCellToolbarActions(menu: IMenu): { primary: IAction[]; secondary: IAction[] } {

src/vs/workbench/contrib/notebook/browser/notebookEditorWidget.ts

@@ -2415,7 +2415,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		return Promise.all(requests);
 	}
 
-	async find(query: string, options: INotebookSearchOptions, token: CancellationToken, skipWarmup: boolean = false): Promise<CellFindMatchWithIndex[]> {
+	async find(query: string, options: INotebookSearchOptions, token: CancellationToken, skipWarmup: boolean = false, shouldGetSearchPreviewInfo = false): Promise<CellFindMatchWithIndex[]> {
 		if (!this._notebookViewModel) {
 			return [];
 		}
@@ -2445,7 +2445,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		if (this._webview) {
 			// request all outputs to be rendered
 			await this._warmupAll(!!options.includeOutput);
-			const webviewMatches = await this._webview.find(query, { caseSensitive: options.caseSensitive, wholeWord: options.wholeWord, includeMarkup: !!options.includeMarkupPreview, includeOutput: !!options.includeOutput });
+			const webviewMatches = await this._webview.find(query, { caseSensitive: options.caseSensitive, wholeWord: options.wholeWord, includeMarkup: !!options.includeMarkupPreview, includeOutput: !!options.includeOutput, shouldGetSearchPreviewInfo });
 			// attach webview matches to model find matches
 			webviewMatches.forEach(match => {
 				if (!options.includeMarkupPreview && match.type === 'preview') {

src/vs/workbench/contrib/notebook/browser/view/renderers/backLayerWebView.ts

@@ -1454,7 +1454,7 @@ export class BackLayerWebView<T extends ICommonCellInfo> extends Themable {
 		});
 	}
 
-	async find(query: string, options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean }): Promise<IFindMatch[]> {
+	async find(query: string, options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean; shouldGetSearchPreviewInfo: boolean }): Promise<IFindMatch[]> {
 		if (query === '') {
 			return [];
 		}

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewMessages.ts

@@ -395,7 +395,7 @@ export interface ITokenizedStylesChangedMessage {
 export interface IFindMessage {
 	readonly type: 'find';
 	readonly query: string;
-	readonly options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean };
+	readonly options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean; shouldGetSearchPreviewInfo: boolean };
 }
 
 

src/vs/workbench/contrib/notebook/browser/view/renderers/webviewPreloads.ts

@@ -829,9 +829,18 @@ async function webviewPreloads(ctx: PreloadContext) {
 		container: Node;
 		originalRange: Range;
 		isShadow: boolean;
+		searchPreviewInfo?: ISearchPreviewInfo;
 		highlightResult?: IHighlightResult;
 	}
 
+	interface ISearchPreviewInfo {
+		line: string;
+		range: {
+			start: number;
+			end: number;
+		};
+	}
+
 	interface IHighlighter {
 		highlightCurrentMatch(index: number): void;
 		unHighlightCurrentMatch(index: number): void;
@@ -956,7 +965,82 @@ async function webviewPreloads(ctx: PreloadContext) {
 		}
 	}
 
-	const find = (query: string, options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean }) => {
+	function extractSelectionLine(selection: Selection): ISearchPreviewInfo {
+		const range = selection.getRangeAt(0);
+
+		// we need to keep a reference to the old selection range to re-apply later
+		const oldRange = range.cloneRange();
+		const captureLength = selection.toString().length;
+
+		// use selection API to modify selection to get entire line (the first line if multi-select)
+
+		// collapse selection to start so that the cursor position is at beginning of match
+		selection.collapseToStart();
+
+		// extend selection in both directions to select the line
+		selection.modify('move', 'backward', 'lineboundary');
+		selection.modify('extend', 'forward', 'lineboundary');
+
+		const line = selection.toString();
+
+		// using the original range and the new range, we can find the offset of the match from the line start.
+		const rangeStart = getStartOffset(selection.getRangeAt(0), oldRange);
+
+		// line range for match
+		const lineRange = {
+			start: rangeStart,
+			end: rangeStart + captureLength,
+		};
+
+		// re-add the old range so that the selection is restored
+		selection.removeAllRanges();
+		selection.addRange(oldRange);
+
+		return { line, range: lineRange };
+	}
+
+	function getStartOffset(lineRange: Range, originalRange: Range) {
+		// sometimes, the old and new range are in different DOM elements (ie: when the match is inside of <b></b>)
+		// so we need to find the first common ancestor DOM element and find the positions of the old and new range relative to that.
+		const firstCommonAncestor = findFirstCommonAncestor(lineRange.startContainer, originalRange.startContainer);
+
+		const selectionOffset = getSelectionOffsetRelativeTo(firstCommonAncestor, lineRange.startContainer) + lineRange.startOffset;
+		const textOffset = getSelectionOffsetRelativeTo(firstCommonAncestor, originalRange.startContainer) + originalRange.startOffset;
+		return textOffset - selectionOffset;
+	}
+
+	// modified from https://stackoverflow.com/a/68583466/16253823
+	function findFirstCommonAncestor(nodeA: Node, nodeB: Node) {
+		const range = new Range();
+		range.setStart(nodeA, 0);
+		range.setEnd(nodeB, 0);
+		return range.commonAncestorContainer;
+	}
+
+	// modified from https://stackoverflow.com/a/48812529/16253823
+	function getSelectionOffsetRelativeTo(parentElement: Node, currentNode: Node | null): number {
+		if (!currentNode) {
+			return 0;
+		}
+		let offset = 0;
+
+		if (currentNode === parentElement || !parentElement.contains(currentNode)) {
+			return offset;
+		}
+
+
+		// count the number of chars before the current dom elem and the start of the dom
+		let prevSibling = currentNode.previousSibling;
+		while (prevSibling) {
+			const nodeContent = prevSibling.nodeValue || '';
+			offset += nodeContent.length;
+			prevSibling = prevSibling.previousSibling;
+		}
+
+		return offset + getSelectionOffsetRelativeTo(parentElement, currentNode.parentNode);
+	}
+
+	const find = (query: string, options: { wholeWord?: boolean; caseSensitive?: boolean; includeMarkup: boolean; includeOutput: boolean; shouldGetSearchPreviewInfo: boolean }) => {
 		let find = true;
 		const matches: IFindMatch[] = [];
 
@@ -1001,7 +1085,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 								cellId: preview.id,
 								container: preview,
 								isShadow: true,
-								originalRange: shadowSelection.getRangeAt(0)
+								originalRange: shadowSelection.getRangeAt(0),
+								searchPreviewInfo: options.shouldGetSearchPreviewInfo ? extractSelectionLine(shadowSelection) : undefined,
 							});
 						}
 					}
@@ -1021,7 +1106,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 								cellId: cellId,
 								container: outputNode,
 								isShadow: true,
-								originalRange: shadowSelection.getRangeAt(0)
+								originalRange: shadowSelection.getRangeAt(0),
+								searchPreviewInfo: options.shouldGetSearchPreviewInfo ? extractSelectionLine(shadowSelection) : undefined,
 							});
 						}
 					}
@@ -1039,7 +1125,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 								cellId: lastEl.cellId,
 								container: lastEl.container,
 								isShadow: false,
-								originalRange: selection.getRangeAt(0)
+								originalRange: selection.getRangeAt(0),
+								searchPreviewInfo: options.shouldGetSearchPreviewInfo ? extractSelectionLine(selection) : undefined,
 							});
 
 						} else {
@@ -1059,7 +1146,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 											cellId: cellId,
 											container: node,
 											isShadow: false,
-											originalRange: selection.getRangeAt(0)
+											originalRange: selection.getRangeAt(0),
+											searchPreviewInfo: options.shouldGetSearchPreviewInfo ? extractSelectionLine(selection) : undefined,
 										});
 									}
 									break;
@@ -1095,7 +1183,8 @@ async function webviewPreloads(ctx: PreloadContext) {
 				type: match.type,
 				id: match.id,
 				cellId: match.cellId,
-				index
+				index,
+				searchPreviewInfo: match.searchPreviewInfo,
 			}))
 		});
 	};

src/vs/workbench/contrib/search/browser/media/searchview.css

@@ -376,3 +376,22 @@
 .monaco-workbench .search-view .monaco-list.element-focused .monaco-list-row.focused.selected:not(.highlighted) .action-label:focus {
 	outline-color: var(--vscode-list-activeSelectionForeground)
 }
+
+.monaco-workbench .search-container .monaco-custom-toggle.disabled {
+	opacity: 0.3;
+	cursor: default;
+	user-select: none;
+	-webkit-user-select: none;
+	pointer-events: none;
+	background-color: inherit !important;
+}
+
+.monaco-workbench .search-container .find-filter-button {
+	color: inherit;
+	margin-left: 2px;
+	float: left;
+	cursor: pointer;
+	box-sizing: border-box;
+	user-select: none;
+	-webkit-user-select: none;
+}

src/vs/workbench/contrib/search/browser/replaceService.ts

@@ -6,7 +6,7 @@
 import * as nls from 'vs/nls';
 import { URI } from 'vs/base/common/uri';
 import * as network from 'vs/base/common/network';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, IReference } from 'vs/base/common/lifecycle';
 import { IReplaceService } from 'vs/workbench/contrib/search/browser/replace';
 import { IEditorService } from 'vs/workbench/services/editor/common/editorService';
 import { IModelService } from 'vs/editor/common/services/model';
@@ -27,7 +27,7 @@ import { ILabelService } from 'vs/platform/label/common/label';
 import { dirname } from 'vs/base/common/resources';
 import { Promises } from 'vs/base/common/async';
 import { SaveSourceRegistry } from 'vs/workbench/common/editor';
-import { CellUri } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { CellUri, IResolvedNotebookEditorModel } from 'vs/workbench/contrib/notebook/common/notebookCommon';
 import { INotebookEditorModelResolverService } from 'vs/workbench/contrib/notebook/common/notebookEditorModelResolverService';
 
 const REPLACE_PREVIEW = 'replacePreview';
@@ -116,10 +116,13 @@ export class ReplaceService implements IReplaceService {
 			if (e.resource.scheme === network.Schemas.vscodeNotebookCell) {
 				const notebookResource = CellUri.parse(e.resource)?.notebook;
 				if (notebookResource) {
-					// todo: find whether there is a common API for saving notebooks and text files
-					const ref = await this.notebookEditorModelResolverService.resolve(notebookResource);
-					await ref.object.save({ source: ReplaceService.REPLACE_SAVE_SOURCE });
-					ref.dispose();
+					let ref: IReference<IResolvedNotebookEditorModel> | undefined;
+					try {
+						ref = await this.notebookEditorModelResolverService.resolve(notebookResource);
+						await ref.object.save({ source: ReplaceService.REPLACE_SAVE_SOURCE });
+					} finally {
+						ref?.dispose();
+					}
 				}
 				return;
 			} else {
@@ -196,8 +199,11 @@ export class ReplaceService implements IReplaceService {
 
 		if (arg instanceof Match) {
 			if (arg instanceof NotebookMatch) {
-				const match = <NotebookMatch>arg;
-				edits.push(this.createEdit(match, match.replaceString, match.cell.uri));
+				if (!arg.isWebviewMatch()) {
+					// only apply edits if it's not a webview match, since webview matches are read-only
+					const match = <NotebookMatch>arg;
+					edits.push(this.createEdit(match, match.replaceString, match.cell.uri));
+				}
 			} else {
 				const match = <Match>arg;
 				edits.push(this.createEdit(match, match.replaceString, resource));
@@ -212,13 +218,12 @@ export class ReplaceService implements IReplaceService {
 			arg.forEach(element => {
 				const fileMatch = <FileMatch>element;
 				if (fileMatch.count() > 0) {
-					edits.push(...fileMatch.matches().map(
-						match => this.createEdit(match, match.replaceString, (match instanceof NotebookMatch) ? match.cell.uri : resource)
+					edits.push(...fileMatch.matches().flatMap(
+						match => this.createEdits(match, resource)
 					));
 				}
 			});
 		}
-
 		return edits;
 	}