tree: collapse and expand should refilter

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

@@ -29,8 +29,9 @@ export interface ITreeNode<T, TFilterData = void> {
 }
 
 interface IMutableTreeNode<T, TFilterData> extends ITreeNode<T, TFilterData> {
-	readonly parent: IMutableTreeNode<T, TFilterData> | undefined;
-	readonly children: IMutableTreeNode<T, TFilterData>[];
+	parent: IMutableTreeNode<T, TFilterData> | undefined;
+	children: IMutableTreeNode<T, TFilterData>[];
+	collapsible: boolean;
 	collapsed: boolean;
 	revealedCount: number;
 	visible: boolean;
@@ -64,46 +65,6 @@ function getVisibleCount<T>(nodes: IMutableTreeNode<T, any>[]): number {
 	return nodes.reduce(visibleCountReducer, 0);
 }
 
-/**
- * Recursively updates the visibleCount of a subtree, while collecting
- * all the visible nodes in an array.
- */
-function updateVisibleCount<T, TFilterData>(node: IMutableTreeNode<T, TFilterData>): ITreeNode<T, TFilterData>[] {
-	const previousVisibleCount = node.revealedCount;
-	const result: ITreeNode<T, TFilterData>[] = [];
-	let first = true;
-
-	function _updateVisibleCount(node: IMutableTreeNode<T, TFilterData>): number {
-		if (!first && !node.visible) {
-			return 0;
-		}
-
-		first = false;
-		result.push(node);
-		node.revealedCount = 1;
-
-		if (!node.collapsed) {
-			for (const child of node.children) {
-				node.revealedCount += _updateVisibleCount(child);
-			}
-		}
-
-		return node.revealedCount;
-	}
-
-	_updateVisibleCount(node);
-
-	const visibleCountDiff = result.length - previousVisibleCount;
-	node = node.parent;
-
-	while (node) {
-		node.revealedCount += visibleCountDiff;
-		node = node.parent;
-	}
-
-	return result;
-}
-
 function getTreeElementIterator<T>(elements: Iterator<ITreeElement<T>> | ITreeElement<T>[] | undefined): Iterator<ITreeElement<T>> {
 	if (!elements) {
 		return Iterator.empty();
@@ -181,38 +142,6 @@ export class TreeModel<T, TFilterData = void> {
 		return Iterator.map(Iterator.fromArray(deletedNodes), treeNodeToElement);
 	}
 
-	private createTreeNode(treeElement: ITreeElement<T>, parent: IMutableTreeNode<T, TFilterData>, revealed: boolean, treeListElements: ITreeNode<T, TFilterData>[]): IMutableTreeNode<T, TFilterData> {
-		const depth = parent.depth + 1;
-		const { element, collapsible, collapsed } = treeElement;
-		const visibility = this.filter ? this.filter.getVisibility(element) : Visibility.Visible;
-
-		let visible = true;
-		let filterData: TFilterData | undefined = undefined;
-
-		if (isFilterResult(visibility)) {
-			visible = visibility.visibility === Visibility.Visible;
-			filterData = visibility.data;
-		} else {
-			visible = visibility === Visibility.Visible;
-		}
-
-		const node = { parent, element, children: [], depth, collapsible: !!collapsible, collapsed: !!collapsed, revealedCount: 1, visible, filterData };
-
-		if (revealed && visible) {
-			treeListElements.push(node);
-		}
-
-		const children = getTreeElementIterator(treeElement.children);
-		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 (!collapsed) {
-			node.revealedCount += getVisibleCount(node.children);
-		}
-
-		return node;
-	}
-
 	getListIndex(location: number[]): number {
 		return this.findNode(location).listIndex;
 	}
@@ -227,6 +156,27 @@ export class TreeModel<T, TFilterData = void> {
 		this._setCollapsed(node, listIndex, revealed);
 	}
 
+	// TODO@joao cleanup
+	setCollapsedAll(collapsed: boolean): void {
+		if (collapsed) {
+			const queue = [...this.root.children]; // TODO@joao use a linked list
+			let listIndex = 0;
+
+			while (queue.length > 0) {
+				const node = queue.shift();
+				const revealed = listIndex < this.root.children.length;
+				this._setCollapsed(node, listIndex, revealed, collapsed);
+
+				queue.push(...node.children);
+				listIndex++;
+			}
+		}
+	}
+
+	isCollapsed(location: number[]): boolean {
+		return this.findNode(location).node.collapsed;
+	}
+
 	private _setCollapsed(node: IMutableTreeNode<T, TFilterData>, listIndex: number, revealed: boolean, collapsed?: boolean | undefined): boolean {
 		if (!node.collapsible) {
 			return false;
@@ -244,7 +194,7 @@ export class TreeModel<T, TFilterData = void> {
 
 		if (revealed) {
 			const previousVisibleCount = node.revealedCount;
-			const toInsert = updateVisibleCount(node);
+			const toInsert = this.updateVisibleCount(node);
 
 			this.list.splice(listIndex + 1, previousVisibleCount - 1, toInsert.slice(1));
 			this._onDidChangeCollapseState.fire(node);
@@ -253,25 +203,82 @@ export class TreeModel<T, TFilterData = void> {
 		return true;
 	}
 
-	// TODO@joao cleanup
-	setCollapsedAll(collapsed: boolean): void {
-		if (collapsed) {
-			const queue = [...this.root.children]; // TODO@joao use a linked list
-			let listIndex = 0;
+	private createTreeNode(treeElement: ITreeElement<T>, parent: IMutableTreeNode<T, TFilterData>, revealed: boolean, treeListElements: ITreeNode<T, TFilterData>[]): IMutableTreeNode<T, TFilterData> {
+		const depth = parent.depth + 1;
+		const { element, collapsible, collapsed } = treeElement;
+		const node: IMutableTreeNode<T, TFilterData> = { parent, element, children: [], depth, collapsible: !!collapsible, collapsed: !!collapsed, revealedCount: 1, visible: true, filterData: undefined };
 
-			while (queue.length > 0) {
-				const node = queue.shift();
-				const revealed = listIndex < this.root.children.length;
-				this._setCollapsed(node, listIndex, revealed, collapsed);
+		this.filterNode(node);
 
-				queue.push(...node.children);
-				listIndex++;
+		if (revealed && node.visible) {
+			treeListElements.push(node);
+		}
+
+		const children = getTreeElementIterator(treeElement.children);
+		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 (!collapsed) {
+			node.revealedCount += getVisibleCount(node.children);
+		}
+
+		return node;
+	}
+
+	/**
+	 * Recursively updates the visibleCount of a subtree, while collecting
+	 * all the visible nodes in an array. Used in expanding/collapsing.
+	 */
+	private updateVisibleCount(node: IMutableTreeNode<T, TFilterData>): ITreeNode<T, TFilterData>[] {
+		const previousVisibleCount = node.revealedCount;
+		const result: ITreeNode<T, TFilterData>[] = [];
+		let first = true;
+
+		const recurse = (node: IMutableTreeNode<T, TFilterData>): number => {
+			if (!first) {
+				this.filterNode(node);
+			}
+
+			if (!first && !node.visible) {
+				return 0;
+			}
+
+			first = false;
+			result.push(node);
+			node.revealedCount = 1;
+
+			if (!node.collapsed) {
+				for (const child of node.children) {
+					node.revealedCount += recurse(child);
+				}
 			}
+
+			return node.revealedCount;
+		};
+
+		recurse(node);
+
+		const visibleCountDiff = result.length - previousVisibleCount;
+		node = node.parent;
+
+		while (node) {
+			node.revealedCount += visibleCountDiff;
+			node = node.parent;
 		}
+
+		return result;
 	}
 
-	isCollapsed(location: number[]): boolean {
-		return this.findNode(location).node.collapsed;
+	private filterNode(node: IMutableTreeNode<T, TFilterData>): void {
+		const visibility = this.filter ? this.filter.getVisibility(node.element) : Visibility.Visible;
+
+		if (isFilterResult(visibility)) {
+			node.visible = visibility.visibility === Visibility.Visible;
+			node.filterData = visibility.data;
+		} else {
+			node.visible = visibility === Visibility.Visible;
+			node.filterData = undefined;
+		}
 	}
 
 	private findNode(location: number[]): { node: IMutableTreeNode<T, TFilterData>, listIndex: number, revealed: boolean } {

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

@@ -343,28 +343,61 @@ suite('TreeModel2', function () {
 		const model = new TreeModel<number>(toSpliceable(list), { filter });
 
 		model.splice([0], 0, Iterator.fromArray([
-			{ element: 0 },
-			{ element: 1 },
-			{ element: 2 },
-			{ element: 3 },
-			{ element: 4 },
-			{ element: 5 },
-			{ element: 6 },
-			{ element: 7 }
+			{
+				element: 0, children: [
+					{ element: 1 },
+					{ element: 2 },
+					{ element: 3 },
+					{ element: 4 },
+					{ element: 5 },
+					{ element: 6 },
+					{ element: 7 }
+				]
+			}
 		]));
 
 		assert.deepEqual(list.length, 4);
-		assert.deepEqual(list[0].element, 0);
-		assert.deepEqual(list[0].collapsed, false);
-		assert.deepEqual(list[0].depth, 1);
-		assert.deepEqual(list[1].element, 2);
-		assert.deepEqual(list[1].collapsed, false);
-		assert.deepEqual(list[1].depth, 1);
-		assert.deepEqual(list[2].element, 4);
-		assert.deepEqual(list[2].collapsed, false);
-		assert.deepEqual(list[2].depth, 1);
-		assert.deepEqual(list[3].element, 6);
-		assert.deepEqual(list[3].collapsed, false);
-		assert.deepEqual(list[3].depth, 1);
+		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
+
+		model.setCollapsed([0], true);
+		assert.deepEqual(toArray(list), [0]);
+
+		model.setCollapsed([0], false);
+		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
+	});
+
+	test('collapse & expand should refilter', function () {
+		const list = [] as ITreeNode<number>[];
+		let shouldFilter = false;
+		const filter = new class implements ITreeFilter<number> {
+			getVisibility(element: number): Visibility {
+				return (!shouldFilter || element % 2 === 0) ? Visibility.Visible : Visibility.Hidden;
+			}
+		};
+
+		const model = new TreeModel<number>(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 }
+				]
+			},
+		]));
+
+		assert.deepEqual(toArray(list), [0, 1, 2, 3, 4, 5, 6, 7]);
+
+		model.setCollapsed([0], true);
+		assert.deepEqual(toArray(list), [0]);
+
+		shouldFilter = true;
+		model.setCollapsed([0], false);
+		assert.deepEqual(toArray(list), [0, 2, 4, 6]);
 	});
 });
\ No newline at end of file