Add ability to search in notebook inputs (under experimental flag)

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

@@ -3,11 +3,12 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 import { Disposable } from 'vs/base/common/lifecycle';
-import { FindMatch, IModelDeltaDecoration } from 'vs/editor/common/model';
+import { IModelDeltaDecoration } from 'vs/editor/common/model';
 import { ModelDecorationOptions } from 'vs/editor/common/model/textModel';
 import { FindDecorations } from 'vs/editor/contrib/find/browser/findDecorations';
+import { Range } from 'vs/editor/common/core/range';
 import { overviewRulerSelectionHighlightForeground, overviewRulerFindMatchForeground } from 'vs/platform/theme/common/colorRegistry';
-import { CellFindMatchWithIndex, CellWebviewFindMatch, ICellModelDecorations, ICellModelDeltaDecorations, ICellViewModel, INotebookDeltaDecoration, INotebookEditor, NotebookOverviewRulerLane, } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
+import { CellFindMatchWithIndex, ICellModelDecorations, ICellModelDeltaDecorations, ICellViewModel, INotebookDeltaDecoration, INotebookEditor, NotebookOverviewRulerLane, } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 
 export class FindMatchDecorationModel extends Disposable {
 	private _allMatchesDecorations: ICellModelDecorations[] = [];
@@ -25,65 +26,72 @@ export class FindMatchDecorationModel extends Disposable {
 		return this._currentMatchDecorations;
 	}
 
-	public async highlightCurrentFindMatchDecoration(cell: ICellViewModel, match: FindMatch | CellWebviewFindMatch): Promise<number | null> {
+	public clearDecorations() {
+		this.clearCurrentFindMatchDecoration();
+		this.setAllFindMatchesDecorations([]);
+	}
 
-		if (match instanceof FindMatch) {
-			this.clearCurrentFindMatchDecoration();
 
-			// match is an editor FindMatch, we update find match decoration in the editor
-			// we will highlight the match in the webview
-			this._notebookEditor.changeModelDecorations(accessor => {
-				const findMatchesOptions: ModelDecorationOptions = FindDecorations._CURRENT_FIND_MATCH_DECORATION;
-
-				const decorations: IModelDeltaDecoration[] = [
-					{ range: match.range, options: findMatchesOptions }
-				];
-				const deltaDecoration: ICellModelDeltaDecorations = {
-					ownerId: cell.handle,
-					decorations: decorations
-				};
-
-				this._currentMatchDecorations = {
-					kind: 'input',
-					decorations: accessor.deltaDecorations(this._currentMatchDecorations?.kind === 'input' ? this._currentMatchDecorations.decorations : [], [deltaDecoration])
-				};
-			});
+	public async highlightCurrentFindMatchDecorationInCell(cell: ICellViewModel, cellRange: Range): Promise<number | null> {
 
-			this._currentMatchCellDecorations = this._notebookEditor.deltaCellDecorations(this._currentMatchCellDecorations, [{
+		this.clearCurrentFindMatchDecoration();
+
+		// match is an editor FindMatch, we update find match decoration in the editor
+		// we will highlight the match in the webview
+		this._notebookEditor.changeModelDecorations(accessor => {
+			const findMatchesOptions: ModelDecorationOptions = FindDecorations._CURRENT_FIND_MATCH_DECORATION;
+
+			const decorations: IModelDeltaDecoration[] = [
+				{ range: cellRange, options: findMatchesOptions }
+			];
+			const deltaDecoration: ICellModelDeltaDecorations = {
 				ownerId: cell.handle,
-				handle: cell.handle,
-				options: {
-					overviewRuler: {
-						color: overviewRulerSelectionHighlightForeground,
-						modelRanges: [match.range],
-						includeOutput: false,
-						position: NotebookOverviewRulerLane.Center
-					}
+				decorations: decorations
+			};
+
+			this._currentMatchDecorations = {
+				kind: 'input',
+				decorations: accessor.deltaDecorations(this._currentMatchDecorations?.kind === 'input' ? this._currentMatchDecorations.decorations : [], [deltaDecoration])
+			};
+		});
+
+		this._currentMatchCellDecorations = this._notebookEditor.deltaCellDecorations(this._currentMatchCellDecorations, [{
+			ownerId: cell.handle,
+			handle: cell.handle,
+			options: {
+				overviewRuler: {
+					color: overviewRulerSelectionHighlightForeground,
+					modelRanges: [cellRange],
+					includeOutput: false,
+					position: NotebookOverviewRulerLane.Center
 				}
-			} as INotebookDeltaDecoration]);
+			}
+		} as INotebookDeltaDecoration]);
 
-			return null;
-		} else {
-			this.clearCurrentFindMatchDecoration();
+		return null;
+	}
 
-			const offset = await this._notebookEditor.highlightFind(cell, match.index);
-			this._currentMatchDecorations = { kind: 'output', index: match.index };
+	public async highlightCurrentFindMatchDecorationInWebview(cell: ICellViewModel, index: number): Promise<number | null> {
 
-			this._currentMatchCellDecorations = this._notebookEditor.deltaCellDecorations(this._currentMatchCellDecorations, [{
-				ownerId: cell.handle,
-				handle: cell.handle,
-				options: {
-					overviewRuler: {
-						color: overviewRulerSelectionHighlightForeground,
-						modelRanges: [],
-						includeOutput: true,
-						position: NotebookOverviewRulerLane.Center
-					}
+		this.clearCurrentFindMatchDecoration();
+
+		const offset = await this._notebookEditor.highlightFind(index);
+		this._currentMatchDecorations = { kind: 'output', index: index };
+
+		this._currentMatchCellDecorations = this._notebookEditor.deltaCellDecorations(this._currentMatchCellDecorations, [{
+			ownerId: cell.handle,
+			handle: cell.handle,
+			options: {
+				overviewRuler: {
+					color: overviewRulerSelectionHighlightForeground,
+					modelRanges: [],
+					includeOutput: true,
+					position: NotebookOverviewRulerLane.Center
 				}
-			} as INotebookDeltaDecoration]);
+			}
+		} as INotebookDeltaDecoration]);
 
-			return offset;
-		}
+		return offset;
 	}
 
 	public clearCurrentFindMatchDecoration() {

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

@@ -457,11 +457,12 @@ export class FindModel extends Disposable {
 	private async highlightCurrentFindMatchDecoration(cellIndex: number, matchIndex: number): Promise<number | null> {
 		const cell = this._findMatches[cellIndex].cell;
 		const match = this._findMatches[cellIndex].getMatch(matchIndex);
-		return this._findMatchDecorationModel.highlightCurrentFindMatchDecoration(cell,
-			(matchIndex < this._findMatches[cellIndex].contentMatches.length) ?
-				(match as FindMatch) :
-				(match as CellWebviewFindMatch)
-		);
+
+		if (matchIndex < this._findMatches[cellIndex].contentMatches.length) {
+			return this._findMatchDecorationModel.highlightCurrentFindMatchDecorationInCell(cell, (match as FindMatch).range);
+		} else {
+			return this._findMatchDecorationModel.highlightCurrentFindMatchDecorationInWebview(cell, (match as CellWebviewFindMatch).index);
+		}
 	}
 
 	clear() {

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

@@ -666,7 +666,7 @@ export interface INotebookEditor {
 	getNextVisibleCellIndex(index: number): number | undefined;
 	getPreviousVisibleCellIndex(index: number): number | undefined;
 	find(query: string, options: INotebookSearchOptions, token: CancellationToken): Promise<CellFindMatchWithIndex[]>;
-	highlightFind(cell: ICellViewModel, matchIndex: number): Promise<number>;
+	highlightFind(matchIndex: number): Promise<number>;
 	unHighlightFind(matchIndex: number): Promise<void>;
 	findStop(): void;
 	showProgress(): void;
@@ -730,8 +730,17 @@ export interface IActiveNotebookEditorDelegate extends INotebookEditorDelegate {
 	getNextVisibleCellIndex(index: number): number;
 }
 
+export interface ISearchPreviewInfo {
+	line: string;
+	range: {
+		start: number;
+		end: number;
+	};
+}
+
 export interface CellWebviewFindMatch {
 	readonly index: number;
+	readonly searchPreviewInfo?: ISearchPreviewInfo;
 }
 
 export type CellContentFindMatch = FindMatch;

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

@@ -124,6 +124,8 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 	readonly onDidChangeCellState = this._onDidChangeCellState.event;
 	private readonly _onDidChangeViewCells = this._register(new Emitter<INotebookViewCellsUpdateEvent>());
 	readonly onDidChangeViewCells: Event<INotebookViewCellsUpdateEvent> = this._onDidChangeViewCells.event;
+	private readonly _onWillChangeModel = this._register(new Emitter<NotebookTextModel | undefined>());
+	readonly onWillChangeModel: Event<NotebookTextModel | undefined> = this._onWillChangeModel.event;
 	private readonly _onDidChangeModel = this._register(new Emitter<NotebookTextModel | undefined>());
 	readonly onDidChangeModel: Event<NotebookTextModel | undefined> = this._onDidChangeModel.event;
 	private readonly _onDidChangeOptions = this._register(new Emitter<void>());
@@ -157,7 +159,6 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 	private readonly _onDidResizeOutputEmitter = this._register(new Emitter<ICellViewModel>());
 	readonly onDidResizeOutput = this._onDidResizeOutputEmitter.event;
 
-
 	//#endregion
 	private _overlayContainer!: HTMLElement;
 	private _notebookTopToolbarContainer!: HTMLElement;
@@ -208,6 +209,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 	}
 
 	set viewModel(newModel: NotebookViewModel | undefined) {
+		this._onWillChangeModel.fire(this._notebookViewModel?.notebookDocument);
 		this._notebookViewModel = newModel;
 		this._onDidChangeModel.fire(newModel?.notebookDocument);
 	}
@@ -2459,7 +2461,7 @@ export class NotebookEditorWidget extends Disposable implements INotebookEditorD
 		return ret;
 	}
 
-	async highlightFind(cell: CodeCellViewModel, matchIndex: number): Promise<number> {
+	async highlightFind(matchIndex: number): Promise<number> {
 		if (!this._webview) {
 			return 0;
 		}

src/vs/workbench/contrib/notebook/browser/services/notebookEditorService.ts

@@ -9,6 +9,8 @@ import { NotebookEditorInput } from 'vs/workbench/contrib/notebook/common/notebo
 import { INotebookEditor, INotebookEditorCreationOptions } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { Event } from 'vs/base/common/event';
 import { Dimension } from 'vs/base/browser/dom';
+import { NotebookEditorWidget } from 'vs/workbench/contrib/notebook/browser/notebookEditorWidget';
+import { URI } from 'vs/base/common/uri';
 
 export const INotebookEditorService = createDecorator<INotebookEditorService>('INotebookEditorWidgetService');
 
@@ -21,6 +23,8 @@ export interface INotebookEditorService {
 
 	retrieveWidget(accessor: ServicesAccessor, group: IEditorGroup, input: NotebookEditorInput, creationOptions?: INotebookEditorCreationOptions, dimension?: Dimension): IBorrowValue<INotebookEditor>;
 
+	retrieveExistingWidgetFromURI(resource: URI): IBorrowValue<NotebookEditorWidget> | undefined;
+	retrieveAllExistingWidgets(): IBorrowValue<NotebookEditorWidget>[];
 	onDidAddNotebookEditor: Event<INotebookEditor>;
 	onDidRemoveNotebookEditor: Event<INotebookEditor>;
 	addNotebookEditor(editor: INotebookEditor): void;

src/vs/workbench/contrib/notebook/browser/services/notebookEditorServiceImpl.ts

@@ -14,6 +14,7 @@ import { INotebookEditor, INotebookEditorCreationOptions } from 'vs/workbench/co
 import { Emitter } from 'vs/base/common/event';
 import { GroupIdentifier } from 'vs/workbench/common/editor';
 import { Dimension } from 'vs/base/browser/dom';
+import { URI } from 'vs/base/common/uri';
 
 export class NotebookEditorWidgetService implements INotebookEditorService {
 
@@ -129,6 +130,26 @@ export class NotebookEditorWidgetService implements INotebookEditorService {
 		targetMap.set(input.resource, widget);
 	}
 
+	retrieveExistingWidgetFromURI(resource: URI): IBorrowValue<NotebookEditorWidget> | undefined {
+		for (const widgetInfo of this._borrowableEditors.values()) {
+			const widget = widgetInfo.get(resource);
+			if (widget) {
+				return this._createBorrowValue(widget.token!, widget);
+			}
+		}
+		return undefined;
+	}
+
+	retrieveAllExistingWidgets(): IBorrowValue<NotebookEditorWidget>[] {
+		const ret: IBorrowValue<NotebookEditorWidget>[] = [];
+		for (const widgetInfo of this._borrowableEditors.values()) {
+			for (const widget of widgetInfo.values()) {
+				ret.push(this._createBorrowValue(widget.token!, widget));
+			}
+		}
+		return ret;
+	}
+
 	retrieveWidget(accessor: ServicesAccessor, group: IEditorGroup, input: NotebookEditorInput, creationOptions?: INotebookEditorCreationOptions, initialDimension?: Dimension): IBorrowValue<NotebookEditorWidget> {
 
 		let value = this._borrowableEditors.get(group.id)?.get(input.resource);

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

@@ -411,11 +411,20 @@ export interface IFindStopMessage {
 	readonly type: 'findStop';
 }
 
+export interface ISearchPreviewInfo {
+	line: string;
+	range: {
+		start: number;
+		end: number;
+	};
+}
+
 export interface IFindMatch {
 	readonly type: 'preview' | 'output';
 	readonly cellId: string;
 	readonly id: string;
 	readonly index: number;
+	readonly searchPreviewInfo?: ISearchPreviewInfo;
 }
 
 export interface IDidFindMessage extends BaseToWebviewMessage {

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

@@ -277,9 +277,23 @@ suite('Notebook Find', () => {
 				mdModel.contentMatches.push(new FindMatch(new Range(1, 1, 1, 2), []));
 				assert.strictEqual(mdModel.length, 1);
 				mdModel.webviewMatches.push({
-					index: 0
+					index: 0,
+					searchPreviewInfo: {
+						line: '',
+						range: {
+							start: 0,
+							end: 0,
+						}
+					}
 				}, {
-					index: 1
+					index: 1,
+					searchPreviewInfo: {
+						line: '',
+						range: {
+							start: 0,
+							end: 0,
+						}
+					}
 				});
 
 				assert.strictEqual(mdModel.length, 3);

src/vs/workbench/contrib/search/common/replace.ts → src/vs/workbench/contrib/search/browser/replace.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Match, FileMatch, FileMatchOrMatch } from 'vs/workbench/contrib/search/common/searchModel';
+import { Match, FileMatch, FileMatchOrMatch } from 'vs/workbench/contrib/search/browser/searchModel';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { IProgress, IProgressStep } from 'vs/platform/progress/common/progress';
 

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

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 import { InstantiationType, registerSingleton } from 'vs/platform/instantiation/common/extensions';
-import { IReplaceService } from 'vs/workbench/contrib/search/common/replace';
+import { IReplaceService } from 'vs/workbench/contrib/search/browser/replace';
 import { ReplaceService, ReplacePreviewContentProvider } from 'vs/workbench/contrib/search/browser/replaceService';
 import { Registry } from 'vs/platform/registry/common/platform';
 import { IWorkbenchContributionsRegistry, Extensions as WorkbenchExtensions } from 'vs/workbench/common/contributions';