tree: fix filter & collapse

src/vs/base/browser/ui/tree/tree.ts

@@ -178,8 +178,8 @@ export class Tree<T, TFilterData = void> implements IDisposable {
 		this.model.setCollapsedAll(true);
 	}
 
-	refilter(location?: number[]): void {
-		this.model.refilter(location);
+	refilter(): void {
+		this.model.refilter();
 	}
 
 	private onMouseClick(e: IListMouseEvent<ITreeNode<T, TFilterData>>): void {

src/vs/base/browser/ui/tree/treeModel.ts

@@ -187,28 +187,10 @@ export class TreeModel<T, TFilterData = void> {
 		return this.findNode(location).node.collapsed;
 	}
 
-	refilter(location?: number[]): void {
-		let node: ITreeNode<T, TFilterData>;
-
-		if (!location || location.length === 0) {
-			node = this.root;
-
-			const previousRevealedCount = node.revealedCount;
-			const toInsert = this.updateSubtreeViewState(this.root);
-			this.list.splice(0, previousRevealedCount, toInsert.slice(1));
-		} else {
-			const findResult = this.findNode(location);
-
-			if (!findResult.revealed) {
-				return;
-			}
-
-			node = findResult.node;
-
-			const previousRevealedCount = node.revealedCount;
-			const toInsert = this.updateSubtreeViewState(this.root);
-			this.list.splice(findResult.listIndex, previousRevealedCount, toInsert);
-		}
+	refilter(/* location?: number[] */): void {
+		const previousRevealedCount = this.root.revealedCount;
+		const toInsert = this.updateNodeAfterFilterChange(this.root);
+		this.list.splice(0, previousRevealedCount, toInsert);
 	}
 
 	private _setCollapsed(node: IMutableTreeNode<T, TFilterData>, listIndex: number, revealed: boolean, collapsed?: boolean | undefined): boolean {
@@ -228,7 +210,7 @@ export class TreeModel<T, TFilterData = void> {
 
 		if (revealed) {
 			const previousRevealedCount = node.revealedCount;
-			const toInsert = this.updateSubtreeViewState(node);
+			const toInsert = this.updateNodeAfterCollapseChange(node);
 
 			this.list.splice(listIndex + 1, previousRevealedCount - 1, toInsert.slice(1));
 			this._onDidChangeCollapseState.fire(node);
@@ -244,7 +226,7 @@ export class TreeModel<T, TFilterData = void> {
 
 		this.updateNodeFilterState(node);
 
-		if (revealed && node.visible) {
+		if (revealed) {
 			treeListElements.push(node);
 		}
 
@@ -252,14 +234,17 @@ export class TreeModel<T, TFilterData = void> {
 		node.children = Iterator.collect(Iterator.map(children, el => this.createTreeNode(el, node, revealed && !treeElement.collapsed, treeListElements)));
 		node.collapsible = node.collapsible || node.children.length > 0;
 
-		if (typeof node.visible === 'undefined' && node.children.length === 0) {
-			node.visible = false;
-			treeListElements.pop();
-		} else {
-			node.visible = true;
+		if (typeof node.visible === 'undefined') {
+			node.visible = node.children.length > 0;
 		}
 
-		if (node.visible && !collapsed) {
+		if (!node.visible) {
+			node.revealedCount = 0;
+
+			if (revealed) {
+				treeListElements.pop();
+			}
+		} else if (!collapsed) {
 			node.revealedCount += getRevealedCount(node.children);
 		}
 
@@ -270,65 +255,92 @@ export class TreeModel<T, TFilterData = void> {
 	 * Recursively updates the view state of a subtree, while collecting
 	 * all the visible nodes in an array. Used in expanding/collapsing.
 	 */
-	private updateSubtreeViewState(node: IMutableTreeNode<T, TFilterData>, filterFirst = false): ITreeNode<T, TFilterData>[] {
+	private updateNodeAfterCollapseChange(node: IMutableTreeNode<T, TFilterData>): ITreeNode<T, TFilterData>[] {
 		const previousRevealedCount = node.revealedCount;
 		const result: ITreeNode<T, TFilterData>[] = [];
-		let first = true;
 
-		const recurse = (node: IMutableTreeNode<T, TFilterData>, revealed = true): number => {
-			if (!first || filterFirst) {
-				this.updateNodeFilterState(node);
+		this._updateNodeAfterCollapseChange(node, result);
+		this._updateParentRevealedCount(node.parent, result.length - previousRevealedCount);
+
+		return result;
+	}
+
+	private _updateNodeAfterCollapseChange(node: IMutableTreeNode<T, TFilterData>, result: ITreeNode<T, TFilterData>[]): number {
+		if (node.visible === false) {
+			return 0;
+		}
+
+		result.push(node);
+		node.revealedCount = 1;
+
+		if (!node.collapsed) {
+			for (const child of node.children) {
+				node.revealedCount += this._updateNodeAfterCollapseChange(child, result);
 			}
+		}
+
+		return node.revealedCount;
+	}
+
+	private updateNodeAfterFilterChange(node: IMutableTreeNode<T, TFilterData>): ITreeNode<T, TFilterData>[] {
+		const previousRevealedCount = node.revealedCount;
+		const result: ITreeNode<T, TFilterData>[] = [];
+
+		this._updateNodeAfterFilterChange(node, result);
+		this._updateParentRevealedCount(node.parent, result.length - previousRevealedCount);
+
+		return result;
+	}
+
+	private _updateNodeAfterFilterChange(node: IMutableTreeNode<T, TFilterData>, result: ITreeNode<T, TFilterData>[], revealed = true): boolean {
+		if (node !== this.root) {
+			this.updateNodeFilterState(node);
 
 			if (node.visible === false) {
-				return 0;
+				return false;
 			}
 
-			first = false;
-
 			if (revealed) {
 				result.push(node);
 			}
+		}
 
-			node.revealedCount = 1;
-
-			let childrenRevealedCount = 0;
-			if (!node.collapsed || typeof node.visible === 'undefined') {
-				for (const child of node.children) {
-					childrenRevealedCount += recurse(child, revealed && !node.collapsed);
-				}
-			}
+		const resultStartLength = result.length;
+		node.revealedCount = node === this.root ? 0 : 1;
 
-			if (typeof node.visible === 'undefined' && childrenRevealedCount === 0) {
-				node.visible = false;
-				node.revealedCount = 0;
-				result.pop();
-				return 0;
+		let hasVisibleDescendants = false;
+		if (typeof node.visible === 'undefined' || !node.collapsed) {
+			for (const child of node.children) {
+				hasVisibleDescendants = this._updateNodeAfterFilterChange(child, result, revealed && !node.collapsed) || hasVisibleDescendants;
 			}
+		}
 
-			if (!node.collapsed) {
-				node.revealedCount += childrenRevealedCount;
-			}
+		if (typeof node.visible === 'undefined') {
+			node.visible = hasVisibleDescendants;
+		}
 
-			return node.revealedCount;
-		};
+		if (!node.visible) {
+			node.revealedCount = 0;
 
-		recurse(node);
+			if (revealed) {
+				result.pop();
+			}
+		} else if (!node.collapsed) {
+			node.revealedCount += result.length - resultStartLength;
+		}
 
-		const revealedCountDiff = result.length - previousRevealedCount;
+		return node.visible;
+	}
 
-		if (revealedCountDiff === 0) {
-			return result;
+	private _updateParentRevealedCount(node: IMutableTreeNode<T, TFilterData>, diff: number): void {
+		if (diff === 0) {
+			return;
 		}
 
-		node = node.parent;
-
 		while (node) {
-			node.revealedCount += revealedCountDiff;
+			node.revealedCount += diff;
 			node = node.parent;
 		}
-
-		return result;
 	}
 
 	private updateNodeFilterState(node: IMutableTreeNode<T, TFilterData>): void {

src/vs/base/test/browser/ui/tree/treeModel.test.ts

@@ -366,7 +366,7 @@ suite('TreeModel2', function () {
 		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
 	});
 
-	test('collapse & expand should refilter', function () {
+	test('refilter', function () {
 		const list = [] as ITreeNode<number>[];
 		let shouldFilter = false;
 		const filter = new class implements ITreeFilter<number> {
@@ -393,54 +393,65 @@ suite('TreeModel2', function () {
 
 		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
 
-		model.setCollapsed([0], true);
-		assert.deepEqual(toArray(list), [0]);
+		model.refilter();
+		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
 
 		shouldFilter = true;
-		model.setCollapsed([0], false);
+		model.refilter();
 		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
+
+		shouldFilter = false;
+		model.refilter();
+		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
 	});
 
-	test('refilter', function () {
-		const list = [] as ITreeNode<number>[];
-		let shouldFilter = false;
-		const filter = new class implements ITreeFilter<number> {
-			filter(element: number): Visibility {
-				return (!shouldFilter || element % 2 === 0) ? Visibility.Visible : Visibility.Hidden;
+	test('recursive filter', function () {
+		const list = [] as ITreeNode<string>[];
+		let query = new RegExp('');
+		const filter = new class implements ITreeFilter<string> {
+			filter(element: string): Visibility {
+				return query.test(element) ? Visibility.Visible : Visibility.Recurse;
 			}
 		};
 
-		const model = new TreeModel<number>(toSpliceable(list), { filter });
+		const model = new TreeModel<string>(toSpliceable(list), { filter });
 
 		model.splice([0], 0, Iterator.fromArray([
 			{
-				element: 0, children: [
-					{ element: 1 },
-					{ element: 2 },
-					{ element: 3 },
-					{ element: 4 },
-					{ element: 5 },
-					{ element: 6 },
-					{ element: 7 }
+				element: 'vscode', children: [
+					{ element: '.build' },
+					{ element: 'git' },
+					{
+						element: 'github', children: [
+							{ element: 'calendar.yml' },
+							{ element: 'endgame' },
+							{ element: 'build.js' },
+						]
+					},
+					{
+						element: 'build', children: [
+							{ element: 'lib' },
+							{ element: 'gulpfile.js' }
+						]
+					}
 				]
 			},
 		]));
 
-		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
+		assert.deepEqual(list.length, 10);
 
+		query = /build/;
 		model.refilter();
-		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
+		assert.deepEqual(toArray(list), ['vscode', '.build', 'github', 'build.js', 'build']);
 
-		shouldFilter = true;
-		model.refilter();
-		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
+		model.setCollapsed([0], true);
+		assert.deepEqual(toArray(list), ['vscode']);
 
-		shouldFilter = false;
-		model.refilter();
-		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
+		model.setCollapsed([0], false);
+		assert.deepEqual(toArray(list), ['vscode', '.build', 'github', 'build.js', 'build']);
 	});
 
-	test('recursive filter', function () {
+	test('recursive filter with collapse', function () {
 		const list = [] as ITreeNode<string>[];
 		let query = new RegExp('');
 		const filter = new class implements ITreeFilter<string> {
@@ -475,15 +486,58 @@ suite('TreeModel2', function () {
 
 		assert.deepEqual(list.length, 10);
 
-		query = /build/;
+		query = /gulp/;
 		model.refilter();
-		assert.deepEqual(toArray(list), ['vscode', '.build', 'github', 'build.js', 'build']);
+		assert.deepEqual(toArray(list), ['vscode', 'build', 'gulpfile.js']);
+
+		model.setCollapsed([0, 3], true);
+		assert.deepEqual(toArray(list), ['vscode', 'build']);
 
 		model.setCollapsed([0], true);
 		assert.deepEqual(toArray(list), ['vscode']);
+	});
+
+	test('recursive filter while collapsed', function () {
+		const list = [] as ITreeNode<string>[];
+		let query = new RegExp('');
+		const filter = new class implements ITreeFilter<string> {
+			filter(element: string): Visibility {
+				return query.test(element) ? Visibility.Visible : Visibility.Recurse;
+			}
+		};
+
+		const model = new TreeModel<string>(toSpliceable(list), { filter });
+
+		model.splice([0], 0, Iterator.fromArray([
+			{
+				element: 'vscode', collapsed: true, children: [
+					{ element: '.build' },
+					{ element: 'git' },
+					{
+						element: 'github', children: [
+							{ element: 'calendar.yml' },
+							{ element: 'endgame' },
+							{ element: 'build.js' },
+						]
+					},
+					{
+						element: 'build', children: [
+							{ element: 'lib' },
+							{ element: 'gulpfile.js' }
+						]
+					}
+				]
+			},
+		]));
+
+		assert.deepEqual(toArray(list), ['vscode']);
+
+		query = /gulp/;
+		model.refilter();
+		assert.deepEqual(toArray(list), ['vscode']);
 
 		model.setCollapsed([0], false);
-		assert.deepEqual(toArray(list), ['vscode', '.build', 'github', 'build.js', 'build']);
+		assert.deepEqual(toArray(list), ['vscode', 'build', 'gulpfile.js']);
 	});
 
 	suite('getNodeLocation', function () {