Fix #65451 - keep search tree element collapse state when refreshing

src/vs/workbench/parts/search/browser/searchView.ts

@@ -65,49 +65,6 @@ import { IUntitledEditorService } from 'vs/workbench/services/untitled/common/un
 
 const $ = dom.$;
 
-function createResultIterator(searchResult: SearchResult, collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
-	const folderMatches = searchResult.folderMatches()
-		.filter(fm => !fm.isEmpty())
-		.sort(searchMatchComparer);
-
-	if (folderMatches.length === 1) {
-		return createFolderIterator(folderMatches[0], collapseResults);
-	}
-
-	const foldersIt = Iterator.fromArray(folderMatches);
-	return Iterator.map(foldersIt, folderMatch => {
-		const children = createFolderIterator(folderMatch, collapseResults);
-		return <ITreeElement<RenderableMatch>>{ element: folderMatch, children };
-	});
-}
-
-function createFolderIterator(folderMatch: FolderMatch, collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
-	const filesIt = Iterator.fromArray(
-		folderMatch.matches()
-			.sort(searchMatchComparer));
-
-	return Iterator.map(filesIt, fileMatch => {
-		const children = createFileIterator(fileMatch);
-
-		const collapsed = collapseResults === 'alwaysCollapse' || (fileMatch.matches().length > 10 && collapseResults !== 'alwaysExpand');
-
-		return <ITreeElement<RenderableMatch>>{ element: fileMatch, children, collapsed };
-	});
-}
-
-function createFileIterator(fileMatch: FileMatch): Iterator<ITreeElement<RenderableMatch>> {
-	const matchesIt = Iterator.from(
-		fileMatch.matches()
-			.sort(searchMatchComparer));
-	return Iterator.map(matchesIt, r => (<ITreeElement<RenderableMatch>>{ element: r }));
-}
-
-export function createIterator(match: FolderMatch | FileMatch | SearchResult, collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
-	return match instanceof SearchResult ? createResultIterator(match, collapseResults) :
-		match instanceof FolderMatch ? createFolderIterator(match, collapseResults) :
-			createFileIterator(match);
-}
-
 export class SearchView extends Viewlet implements IViewlet, IPanel {
 
 	private static readonly MAX_TEXT_RESULTS = 10000;
@@ -483,21 +440,68 @@ export class SearchView extends Viewlet implements IViewlet, IPanel {
 	refreshTree(event?: IChangeEvent): void {
 		const collapseResults = this.configurationService.getValue<ISearchConfigurationProperties>('search').collapseResults;
 		if (!event || event.added || event.removed) {
-			this.tree.setChildren(null, createResultIterator(this.viewModel.searchResult, collapseResults));
+			this.tree.setChildren(null, this.createResultIterator(collapseResults));
 		} else {
 			event.elements.forEach(element => {
 				if (element instanceof FolderMatch) {
 					// The folder may or may not be in the tree. Refresh the whole thing.
-					this.tree.setChildren(null, createResultIterator(this.viewModel.searchResult, collapseResults));
+					this.tree.setChildren(null, this.createResultIterator(collapseResults));
 					return;
 				}
 
 				const root = element instanceof SearchResult ? null : element;
-				this.tree.setChildren(root, createIterator(element, collapseResults));
+				this.tree.setChildren(root, this.createIterator(element, collapseResults));
 			});
 		}
 	}
 
+	private createResultIterator(collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
+		const folderMatches = this.searchResult.folderMatches()
+			.filter(fm => !fm.isEmpty())
+			.sort(searchMatchComparer);
+
+		if (folderMatches.length === 1) {
+			return this.createFolderIterator(folderMatches[0], collapseResults);
+		}
+
+		const foldersIt = Iterator.fromArray(folderMatches);
+		return Iterator.map(foldersIt, folderMatch => {
+			const children = this.createFolderIterator(folderMatch, collapseResults);
+			return <ITreeElement<RenderableMatch>>{ element: folderMatch, children };
+		});
+	}
+
+	private createFolderIterator(folderMatch: FolderMatch, collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
+		const filesIt = Iterator.fromArray(
+			folderMatch.matches()
+				.sort(searchMatchComparer));
+
+		return Iterator.map(filesIt, fileMatch => {
+			const children = this.createFileIterator(fileMatch);
+
+			let nodeExists = true;
+			try { this.tree.getNode(fileMatch); } catch (e) { nodeExists = false; }
+
+			const collapsed = nodeExists ? undefined :
+				(collapseResults === 'alwaysCollapse' || (fileMatch.matches().length > 10 && collapseResults !== 'alwaysExpand'));
+
+			return <ITreeElement<RenderableMatch>>{ element: fileMatch, children, collapsed };
+		});
+	}
+
+	private createFileIterator(fileMatch: FileMatch): Iterator<ITreeElement<RenderableMatch>> {
+		const matchesIt = Iterator.from(
+			fileMatch.matches()
+				.sort(searchMatchComparer));
+		return Iterator.map(matchesIt, r => (<ITreeElement<RenderableMatch>>{ element: r }));
+	}
+
+	private createIterator(match: FolderMatch | FileMatch | SearchResult, collapseResults: ISearchConfigurationProperties['collapseResults']): Iterator<ITreeElement<RenderableMatch>> {
+		return match instanceof SearchResult ? this.createResultIterator(collapseResults) :
+			match instanceof FolderMatch ? this.createFolderIterator(match, collapseResults) :
+				this.createFileIterator(match);
+	}
+
 	private replaceAll(): void {
 		if (this.viewModel.searchResult.count() === 0) {
 			return;

src/vs/workbench/parts/search/test/browser/searchViewlet.test.ts

@@ -12,11 +12,8 @@ import { TestInstantiationService } from 'vs/platform/instantiation/test/common/
 import { IFileMatch, ITextSearchMatch, OneLineRange, QueryType } from 'vs/platform/search/common/search';
 import { IWorkspaceContextService } from 'vs/platform/workspace/common/workspace';
 import { TestWorkspace } from 'vs/platform/workspace/test/common/testWorkspace';
-import { FileMatch, Match, SearchResult, RenderableMatch, searchMatchComparer } from 'vs/workbench/parts/search/common/searchModel';
+import { FileMatch, Match, searchMatchComparer, SearchResult } from 'vs/workbench/parts/search/common/searchModel';
 import { TestContextService } from 'vs/workbench/test/workbenchTestServices';
-import { createIterator } from 'vs/workbench/parts/search/browser/searchView';
-import { ITreeElement } from 'vs/base/browser/ui/tree/tree';
-import { Iterator } from 'vs/base/common/iterator';
 
 suite('Search - Viewlet', () => {
 	let instantiation: TestInstantiationService;
@@ -63,12 +60,6 @@ suite('Search - Viewlet', () => {
 
 		assert.equal(fileMatch.id(), 'file:///c%3A/foo');
 		assert.equal(lineMatch.id(), 'file:///c%3A/foo>[2,1 -> 2,2]b');
-
-		const resultIterator = createIterator(result, 'auto');
-		const first = resultIterator.next();
-
-		assert(!!first.value!.children);
-		assert.equal((<Iterator<ITreeElement<RenderableMatch>>>first.value!.children).next().value!.element.id(), 'file:///c%3A/foo>[2,1 -> 2,2]b');
 	});
 
 	test('Comparer', () => {