From 0b43b16873f883c37c352dfa5198511d5ae8c0c6 Mon Sep 17 00:00:00 2001
From: Rob Lourens <roblourens@gmail.com>
Date: Mon, 31 Dec 2018 15:52:06 -0800
Subject: [PATCH] Fix #64926

---
 .../services/search/node/textSearchManager.ts | 45 +++++++++++++++----
 .../api/extHostSearch.test.ts                 |  4 +-
 2 files changed, 38 insertions(+), 11 deletions(-)

diff --git a/src/vs/workbench/services/search/node/textSearchManager.ts b/src/vs/workbench/services/search/node/textSearchManager.ts
index 2d1d7ba49e4..6f953a3d549 100644
--- a/src/vs/workbench/services/search/node/textSearchManager.ts
+++ b/src/vs/workbench/services/search/node/textSearchManager.ts
@@ -35,20 +35,26 @@ export class TextSearchManager {
 			this.collector = new TextSearchResultsCollector(onProgress);
 
 			let isCanceled = false;
-			const onResult = (match: vscode.TextSearchResult, folderIdx: number) => {
+			const onResult = (result: vscode.TextSearchResult, folderIdx: number) => {
 				if (isCanceled) {
 					return;
 				}
 
-				if (typeof this.query.maxResults === 'number' && this.resultCount >= this.query.maxResults) {
-					this.isLimitHit = true;
-					isCanceled = true;
-					tokenSource.cancel();
-				}
-
 				if (!this.isLimitHit) {
-					this.resultCount++;
-					this.collector.add(match, folderIdx);
+					const resultSize = this.resultSize(result);
+					if (extensionResultIsMatch(result) && typeof this.query.maxResults === 'number' && this.resultCount + resultSize > this.query.maxResults) {
+						this.isLimitHit = true;
+						isCanceled = true;
+						tokenSource.cancel();
+
+						result = this.trimResultToSize(result, this.query.maxResults - this.resultCount);
+					}
+
+					const newResultSize = this.resultSize(result);
+					this.resultCount += newResultSize;
+					if (newResultSize > 0) {
+						this.collector.add(result, folderIdx);
+					}
 				}
 			};
 
@@ -74,6 +80,27 @@ export class TextSearchManager {
 		});
 	}
 
+	private resultSize(result: vscode.TextSearchResult): number {
+		const match = <vscode.TextSearchMatch>result;
+		return Array.isArray(match.ranges) ?
+			match.ranges.length :
+			1;
+	}
+
+	private trimResultToSize(result: vscode.TextSearchMatch, size: number): vscode.TextSearchMatch {
+		const rangesArr = Array.isArray(result.ranges) ? result.ranges : [result.ranges];
+		const matchesArr = Array.isArray(result.preview.matches) ? result.preview.matches : [result.preview.matches];
+
+		return {
+			ranges: rangesArr.slice(0, size),
+			preview: {
+				matches: matchesArr.slice(0, size),
+				text: result.preview.text
+			},
+			uri: result.uri
+		};
+	}
+
 	private searchInFolder(folderQuery: IFolderQuery<URI>, onResult: (result: vscode.TextSearchResult) => void, token: CancellationToken): Promise<vscode.TextSearchComplete | null | undefined> {
 		const queryTester = new QueryGlobTester(this.query, folderQuery);
 		const testingPs: Promise<void>[] = [];
diff --git a/src/vs/workbench/test/electron-browser/api/extHostSearch.test.ts b/src/vs/workbench/test/electron-browser/api/extHostSearch.test.ts
index 1eb1b00694f..0a0084ddcc7 100644
--- a/src/vs/workbench/test/electron-browser/api/extHostSearch.test.ts
+++ b/src/vs/workbench/test/electron-browser/api/extHostSearch.test.ts
@@ -627,7 +627,7 @@ suite('ExtHostSearch', () => {
 
 		function makePreview(text: string): vscode.TextSearchMatch['preview'] {
 			return {
-				matches: new Range(0, 0, 0, text.length),
+				matches: [new Range(0, 0, 0, text.length)],
 				text
 			};
 		}
@@ -635,7 +635,7 @@ suite('ExtHostSearch', () => {
 		function makeTextResult(baseFolder: URI, relativePath: string): vscode.TextSearchMatch {
 			return {
 				preview: makePreview('foo'),
-				ranges: new Range(0, 0, 0, 3),
+				ranges: [new Range(0, 0, 0, 3)],
 				uri: joinPath(baseFolder, relativePath)
 			};
 		}