Allow search provider results to be outside of folder

src/vs/workbench/api/test/node/extHostSearch.test.ts

@@ -702,6 +702,25 @@ suite('ExtHostSearch', () => {
 			const { results } = await runFileSearch(query);
 			compareURIs(results, reportedResults);
 		});
+
+		test('Works when provider returns files that are not in the original workspace', async () => {
+			const reportedResults = [
+				joinPath(rootFolderB, 'file1.ts'),
+				joinPath(rootFolderA, 'file2.ts'),
+				joinPath(rootFolderA, 'subfolder/file3.ts')
+			];
+
+			await registerTestFileSearchProvider({
+				provideFileSearchResults(query: vscode.FileSearchQuery, options: vscode.FileSearchOptions, token: vscode.CancellationToken): Promise<URI[]> {
+					return Promise.resolve(reportedResults);
+				}
+			});
+
+			const { results, stats } = await runFileSearch(getSimpleQuery());
+			assert(!stats.limitHit);
+			assert.strictEqual(results.length, 3);
+			compareURIs(results, reportedResults);
+		});
 	});
 
 	suite('Text:', () => {
@@ -1296,5 +1315,24 @@ suite('ExtHostSearch', () => {
 			const { results } = await runTextSearch(query);
 			assertResults(results, providedResults);
 		});
+
+
+		test('Works when provider returns files that are not in the original workspace', async () => {
+			const providedResults: vscode.TextSearchResult[] = [
+				makeTextResult(rootFolderB, 'file1.ts'),
+				makeTextResult(rootFolderA, 'file2.ts')
+			];
+
+			await registerTestTextSearchProvider({
+				provideTextSearchResults(query: vscode.TextSearchQuery, options: vscode.TextSearchOptions, progress: vscode.Progress<vscode.TextSearchResult>, token: vscode.CancellationToken): Promise<vscode.TextSearchComplete> {
+					providedResults.forEach(r => progress.report(r));
+					return Promise.resolve(null!);
+				}
+			});
+
+			const { results, stats } = await runTextSearch(getSimpleQuery('foo'));
+			assert(!stats.limitHit);
+			assertResults(results, providedResults);
+		});
 	});
 });

src/vs/workbench/services/search/common/fileSearchManager.ts

@@ -12,9 +12,9 @@ import { StopWatch } from '../../../../base/common/stopwatch.js';
 import { URI } from '../../../../base/common/uri.js';
 import { IFileMatch, IFileSearchProviderStats, IFolderQuery, ISearchCompleteStats, IFileQuery, QueryGlobTester, resolvePatternsForProvider, hasSiblingFn, excludeToGlobPattern, DEFAULT_MAX_SEARCH_RESULTS } from './search.js';
 import { FileSearchProviderFolderOptions, FileSearchProviderNew, FileSearchProviderOptions } from './searchExtTypes.js';
-import { TernarySearchTree } from '../../../../base/common/ternarySearchTree.js';
 import { Disposable } from '../../../../base/common/lifecycle.js';
 import { OldFileSearchProviderConverter } from './searchExtConversionTypes.js';
+import { ResourceMap } from '../../../../base/common/map.js';
 
 interface IInternalFileMatch {
 	base: URI;
@@ -126,7 +126,7 @@ class FileSearchEngine {
 		};
 
 
-		const folderMappings: TernarySearchTree<URI, FolderQueryInfo> = TernarySearchTree.forUris<FolderQueryInfo>();
+		const folderMappings = new ResourceMap<FolderQueryInfo>();
 		fqs.forEach(fq => {
 			const queryTester = new QueryGlobTester(this.config, fq);
 			const noSiblingsClauses = !queryTester.hasSiblingExcludeClauses();
@@ -153,9 +153,9 @@ class FileSearchEngine {
 
 
 			if (results) {
-				results.forEach(result => {
-
-					const fqFolderInfo = folderMappings.findSubstr(result)!;
+				results.forEach(resultInfo => {
+					const result = resultInfo.result;
+					const fqFolderInfo = folderMappings.get(resultInfo.folder)!;
 					const relativePath = path.posix.relative(fqFolderInfo.folder.path, result.path);
 
 					if (fqFolderInfo.noSiblingsClauses) {

src/vs/workbench/services/search/common/searchExtConversionTypes.ts

@@ -13,7 +13,7 @@ import { CancellationToken } from '../../../../base/common/cancellation.js';
 import { URI } from '../../../../base/common/uri.js';
 import { IProgress } from '../../../../platform/progress/common/progress.js';
 import { DEFAULT_TEXT_SEARCH_PREVIEW_OPTIONS } from './search.js';
-import { Range, FileSearchProviderNew, FileSearchProviderOptions, ProviderResult, TextSearchCompleteNew, TextSearchContextNew, TextSearchMatchNew, TextSearchProviderNew, TextSearchProviderOptions, TextSearchQueryNew, TextSearchResultNew, AITextSearchProviderNew, TextSearchCompleteMessage } from './searchExtTypes.js';
+import { Range, FileSearchProviderNew, FileSearchProviderOptions, ProviderResult, TextSearchCompleteNew, TextSearchContextNew, TextSearchMatchNew, TextSearchProviderNew, TextSearchProviderOptions, TextSearchQueryNew, TextSearchResultNew, AITextSearchProviderNew, TextSearchCompleteMessage, SearchResultFromFolder } from './searchExtTypes.js';
 
 // old types that are retained for backward compatibility
 // TODO: delete this when search apis are adopted by all first-party extensions
@@ -470,13 +470,22 @@ function newToOldFileProviderOptions(options: FileSearchProviderOptions): FileSe
 export class OldFileSearchProviderConverter implements FileSearchProviderNew {
 	constructor(private provider: FileSearchProvider) { }
 
-	provideFileSearchResults(pattern: string, options: FileSearchProviderOptions, token: CancellationToken): ProviderResult<URI[]> {
-		const getResult = async () => {
+	provideFileSearchResults(pattern: string, options: FileSearchProviderOptions, token: CancellationToken): ProviderResult<SearchResultFromFolder<URI>[]> {
+		const getResult: () => Promise<{
+			uris: URI[] | null | undefined;
+			folder: URI;
+		}[]> = async () => {
 			const newOpts = newToOldFileProviderOptions(options);
 			return Promise.all(newOpts.map(
-				o => this.provider.provideFileSearchResults({ pattern }, o, token)));
+				async o => ({ uris: await this.provider.provideFileSearchResults({ pattern }, o, token), folder: o.folder })));
 		};
-		return getResult().then(e => coalesce(e).flat());
+		return getResult().then(rawResult =>
+			coalesce(rawResult.flatMap(e => (
+				e.uris?.map(uri => (
+					{ result: uri, folder: e.folder }
+				)))
+			))
+		);
 	}
 }
 
@@ -528,19 +537,22 @@ export function oldToNewTextSearchResult(result: TextSearchResult): TextSearchRe
 export class OldTextSearchProviderConverter implements TextSearchProviderNew {
 	constructor(private provider: TextSearchProvider) { }
 
-	provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: IProgress<TextSearchResultNew>, token: CancellationToken): ProviderResult<TextSearchCompleteNew> {
+	provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: IProgress<SearchResultFromFolder<TextSearchResultNew>>, token: CancellationToken): ProviderResult<TextSearchCompleteNew> {
 
-		const progressShim = (oldResult: TextSearchResult) => {
+		const progressShim = (oldResult: TextSearchResult, folder: URI) => {
 			if (!validateProviderResult(oldResult)) {
 				return;
 			}
-			progress.report(oldToNewTextSearchResult(oldResult));
+			progress.report({
+				folder,
+				result: oldToNewTextSearchResult(oldResult)
+			});
 		};
 
 		const getResult = async () => {
 			return coalesce(await Promise.all(
 				newToOldTextProviderOptions(options).map(
-					o => this.provider.provideTextSearchResults(query, o, { report: (e) => progressShim(e) }, token))))
+					o => this.provider.provideTextSearchResults(query, o, { report: (e) => progressShim(e, o.folder) }, token))))
 				.reduce(
 					(prev, cur) => ({ limitHit: prev.limitHit || cur.limitHit }),
 					{ limitHit: false }
@@ -559,18 +571,21 @@ export class OldTextSearchProviderConverter implements TextSearchProviderNew {
 export class OldAITextSearchProviderConverter implements AITextSearchProviderNew {
 	constructor(private provider: AITextSearchProvider) { }
 
-	provideAITextSearchResults(query: string, options: TextSearchProviderOptions, progress: IProgress<TextSearchResultNew>, token: CancellationToken): ProviderResult<TextSearchCompleteNew> {
-		const progressShim = (oldResult: TextSearchResult) => {
+	provideAITextSearchResults(query: string, options: TextSearchProviderOptions, progress: IProgress<SearchResultFromFolder<TextSearchResultNew>>, token: CancellationToken): ProviderResult<TextSearchCompleteNew> {
+		const progressShim = (oldResult: TextSearchResult, folder: URI) => {
 			if (!validateProviderResult(oldResult)) {
 				return;
 			}
-			progress.report(oldToNewTextSearchResult(oldResult));
+			progress.report({
+				folder,
+				result: oldToNewTextSearchResult(oldResult)
+			});
 		};
 
 		const getResult = async () => {
 			return coalesce(await Promise.all(
 				newToOldTextProviderOptions(options).map(
-					o => this.provider.provideAITextSearchResults(query, o, { report: (e) => progressShim(e) }, token))))
+					o => this.provider.provideAITextSearchResults(query, o, { report: (e) => progressShim(e, o.folder) }, token))))
 				.reduce(
 					(prev, cur) => ({ limitHit: prev.limitHit || cur.limitHit }),
 					{ limitHit: false }

src/vs/workbench/services/search/common/searchExtTypes.ts

@@ -320,41 +320,6 @@ export class TextSearchContextNew {
  */
 export type TextSearchResultNew = TextSearchMatchNew | TextSearchContextNew;
 
-
-/**
- * A FileSearchProvider provides search results for files in the given folder that match a query string. It can be invoked by quickaccess or other extensions.
- *
- * A FileSearchProvider is the more powerful of two ways to implement file search in VS Code. Use a FileSearchProvider if you wish to search within a folder for
- * all files that match the user's query.
- *
- * The FileSearchProvider will be invoked on every keypress in quickaccess. When `workspace.findFiles` is called, it will be invoked with an empty query string,
- * and in that case, every file in the folder should be returned.
- */
-export interface FileSearchProviderNew {
-	/**
-	 * Provide the set of files that match a certain file path pattern.
-	 * @param query The parameters for this query.
-	 * @param options A set of options to consider while searching files.
-	 * @param progress A progress callback that must be invoked for all results.
-	 * @param token A cancellation token.
-	 */
-	provideFileSearchResults(pattern: string, options: FileSearchProviderOptions, token: CancellationToken): ProviderResult<URI[]>;
-}
-
-/**
- * A TextSearchProvider provides search results for text results inside files in the workspace.
- */
-export interface TextSearchProviderNew {
-	/**
-	 * Provide results that match the given text pattern.
-	 * @param query The parameters for this query.
-	 * @param options A set of options to consider while searching.
-	 * @param progress A progress callback that must be invoked for all results.
-	 * @param token A cancellation token.
-	 */
-	provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: IProgress<TextSearchResultNew>, token: CancellationToken): ProviderResult<TextSearchCompleteNew>;
-}
-
 /**
  * Information collected when text search is complete.
  */
@@ -417,7 +382,7 @@ export interface FileSearchProviderNew {
 	 * @param progress A progress callback that must be invoked for all results.
 	 * @param token A cancellation token.
 	 */
-	provideFileSearchResults(pattern: string, options: FileSearchProviderOptions, token: CancellationToken): ProviderResult<URI[]>;
+	provideFileSearchResults(pattern: string, options: FileSearchProviderOptions, token: CancellationToken): ProviderResult<SearchResultFromFolder<URI>[]>;
 }
 
 /**
@@ -431,7 +396,7 @@ export interface TextSearchProviderNew {
 	 * @param progress A progress callback that must be invoked for all results.
 	 * @param token A cancellation token.
 	 */
-	provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: IProgress<TextSearchResultNew>, token: CancellationToken): ProviderResult<TextSearchCompleteNew>;
+	provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: IProgress<SearchResultFromFolder<TextSearchResultNew>>, token: CancellationToken): ProviderResult<TextSearchCompleteNew>;
 }
 
 /**
@@ -537,5 +502,13 @@ export interface AITextSearchProviderNew {
 	 * @param progress A progress callback that must be invoked for all results.
 	 * @param token A cancellation token.
 	 */
-	provideAITextSearchResults(query: string, options: TextSearchProviderOptions, progress: IProgress<TextSearchResultNew>, token: CancellationToken): ProviderResult<TextSearchCompleteNew>;
+	provideAITextSearchResults(query: string, options: TextSearchProviderOptions, progress: IProgress<SearchResultFromFolder<TextSearchResultNew>>, token: CancellationToken): ProviderResult<TextSearchCompleteNew>;
 }
+
+/**
+ * A wrapper for a search result that indicates the original workspace folder that this result was found for.
+ */
+export type SearchResultFromFolder<T> = {
+	result: T;
+	folder: URI;
+};

src/vs/workbench/services/search/common/textSearchManager.ts

@@ -6,13 +6,13 @@
 import { isThenable } from '../../../../base/common/async.js';
 import { CancellationToken, CancellationTokenSource } from '../../../../base/common/cancellation.js';
 import { toErrorMessage } from '../../../../base/common/errorMessage.js';
+import { ResourceMap } from '../../../../base/common/map.js';
 import { Schemas } from '../../../../base/common/network.js';
 import * as path from '../../../../base/common/path.js';
 import * as resources from '../../../../base/common/resources.js';
-import { TernarySearchTree } from '../../../../base/common/ternarySearchTree.js';
 import { URI } from '../../../../base/common/uri.js';
 import { DEFAULT_MAX_SEARCH_RESULTS, hasSiblingPromiseFn, IAITextQuery, IExtendedExtensionSearchOptions, IFileMatch, IFolderQuery, excludeToGlobPattern, IPatternInfo, ISearchCompleteStats, ITextQuery, ITextSearchContext, ITextSearchMatch, ITextSearchResult, ITextSearchStats, QueryGlobTester, QueryType, resolvePatternsForProvider, ISearchRange, DEFAULT_TEXT_SEARCH_PREVIEW_OPTIONS } from './search.js';
-import { AITextSearchProviderNew, TextSearchCompleteNew, TextSearchMatchNew, TextSearchProviderFolderOptions, TextSearchProviderNew, TextSearchProviderOptions, TextSearchQueryNew, TextSearchResultNew } from './searchExtTypes.js';
+import { AITextSearchProviderNew, SearchResultFromFolder, TextSearchCompleteNew, TextSearchMatchNew, TextSearchProviderFolderOptions, TextSearchProviderNew, TextSearchProviderOptions, TextSearchQueryNew, TextSearchResultNew } from './searchExtTypes.js';
 
 export interface IFileUtils {
 	readdir: (resource: URI) => Promise<string[]>;
@@ -122,20 +122,20 @@ export class TextSearchManager {
 	}
 
 	private async doSearch(folderQueries: IFolderQuery<URI>[], onResult: (result: TextSearchResultNew, folderIdx: number) => void, token: CancellationToken): Promise<TextSearchCompleteNew | null | undefined> {
-		const folderMappings: TernarySearchTree<URI, FolderQueryInfo> = TernarySearchTree.forUris<FolderQueryInfo>();
+		const folderMappings: ResourceMap<FolderQueryInfo> = new ResourceMap<FolderQueryInfo>();
 		folderQueries.forEach((fq, i) => {
 			const queryTester = new QueryGlobTester(this.query, fq);
 			folderMappings.set(fq.folder, { queryTester, folder: fq.folder, folderIdx: i });
 		});
 
 		const testingPs: Promise<void>[] = [];
 		const progress = {
-			report: (result: TextSearchResultNew) => {
-
+			report: (resultInfo: SearchResultFromFolder<TextSearchResultNew>) => {
+				const result = resultInfo.result;
 				if (result.uri === undefined) {
 					throw Error('Text search result URI is undefined. Please check provider implementation.');
 				}
-				const folderQuery = folderMappings.findSubstr(result.uri)!;
+				const folderQuery = folderMappings.get(resultInfo.folder)!;
 				const hasSibling = folderQuery.folder.scheme === Schemas.file ?
 					hasSiblingPromiseFn(() => {
 						return this.fileUtils.readdir(resources.dirname(result.uri));

src/vs/workbench/services/search/node/ripgrepSearchProvider.ts

@@ -6,7 +6,7 @@
 import { CancellationTokenSource, CancellationToken } from '../../../../base/common/cancellation.js';
 import { OutputChannel } from './ripgrepSearchUtils.js';
 import { RipgrepTextSearchEngine } from './ripgrepTextSearchEngine.js';
-import { TextSearchProviderNew, TextSearchCompleteNew, TextSearchResultNew, TextSearchQueryNew, TextSearchProviderOptions, } from '../common/searchExtTypes.js';
+import { TextSearchProviderNew, TextSearchCompleteNew, TextSearchResultNew, TextSearchQueryNew, TextSearchProviderOptions, SearchResultFromFolder, } from '../common/searchExtTypes.js';
 import { Progress } from '../../../../platform/progress/common/progress.js';
 import { Schemas } from '../../../../base/common/network.js';
 import type { RipgrepTextSearchOptions } from '../common/searchExtTypesInternal.js';
@@ -18,7 +18,7 @@ export class RipgrepSearchProvider implements TextSearchProviderNew {
 		process.once('exit', () => this.dispose());
 	}
 
-	async provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: Progress<TextSearchResultNew>, token: CancellationToken): Promise<TextSearchCompleteNew> {
+	async provideTextSearchResults(query: TextSearchQueryNew, options: TextSearchProviderOptions, progress: Progress<SearchResultFromFolder<TextSearchResultNew>>, token: CancellationToken): Promise<TextSearchCompleteNew> {
 		const numThreads = await this.getNumThreads();
 		const engine = new RipgrepTextSearchEngine(this.outputChannel, numThreads);
 
@@ -36,10 +36,10 @@ export class RipgrepSearchProvider implements TextSearchProviderNew {
 				// Ripgrep search engine can only provide file-scheme results, but we want to use it to search some schemes that are backed by the filesystem, but with some other provider as the frontend,
 				// case in point vscode-userdata. In these cases we translate the query to a file, and translate the results back to the frontend scheme.
 				const translatedOptions = { ...extendedOptions, folder: folderOption.folder.with({ scheme: Schemas.file }) };
-				const progressTranslator = new Progress<TextSearchResultNew>(data => progress.report({ ...data, uri: data.uri.with({ scheme: folderOption.folder.scheme }) }));
-				return this.withToken(token, token => engine.provideTextSearchResultsWithRgOptions(query, translatedOptions, progressTranslator, token));
+				const progressTranslator = new Progress<SearchResultFromFolder<TextSearchResultNew>>(data => progress.report({ folder: data.folder, result: { ...data.result, uri: data.result.uri.with({ scheme: folderOption.folder.scheme }) } }));
+				return this.withToken(token, token => engine.provideTextSearchResultsWithRgOptions(query, translatedOptions, folderOption.folder, progressTranslator, token));
 			} else {
-				return this.withToken(token, token => engine.provideTextSearchResultsWithRgOptions(query, extendedOptions, progress, token));
+				return this.withToken(token, token => engine.provideTextSearchResultsWithRgOptions(query, extendedOptions, folderOption.folder, progress, token));
 			}
 		})).then((e => {
 			const complete: TextSearchCompleteNew = {