wip: propagate filter result to tree nodes

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

@@ -101,7 +101,7 @@ class TreeRenderer<T, TFilterData, TTemplateData> implements IRenderer<ITreeNode
 
 		templateData.twistie.style.width = `${10 + node.depth * 10}px`;
 		renderTwistie(node, templateData.twistie);
-		templateData.count.textContent = `${node.visibleCount}`;
+		templateData.count.textContent = `${node.revealedCount}`;
 
 		this.renderer.renderElement(node.element, index, templateData.templateData);
 	}
@@ -122,7 +122,7 @@ class TreeRenderer<T, TFilterData, TTemplateData> implements IRenderer<ITreeNode
 		}
 
 		renderTwistie(node, templateData.twistie);
-		templateData.count.textContent = `${node.visibleCount}`;
+		templateData.count.textContent = `${node.revealedCount}`;
 	}
 
 	dispose(): void {

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

@@ -23,7 +23,7 @@ export interface ITreeNode<T, TFilterData = void> {
 	readonly depth: number;
 	readonly collapsible: boolean;
 	readonly collapsed: boolean;
-	readonly visibleCount: number;
+	readonly revealedCount: number;
 	readonly visible: boolean;
 	readonly filterData: TFilterData | undefined;
 }
@@ -32,7 +32,7 @@ interface IMutableTreeNode<T, TFilterData> extends ITreeNode<T, TFilterData> {
 	readonly parent: IMutableTreeNode<T, TFilterData> | undefined;
 	readonly children: IMutableTreeNode<T, TFilterData>[];
 	collapsed: boolean;
-	visibleCount: number;
+	revealedCount: number;
 	visible: boolean;
 	filterData: TFilterData | undefined;
 }
@@ -40,7 +40,7 @@ interface IMutableTreeNode<T, TFilterData> extends ITreeNode<T, TFilterData> {
 export const enum Visibility {
 	Hidden,
 	Visible,
-	Recurse // TODO@joao come up with a better name
+	// Recurse // TODO@joao come up with a better name
 }
 
 export interface IFilterResult<TFilterData> {
@@ -48,12 +48,16 @@ export interface IFilterResult<TFilterData> {
 	data: TFilterData;
 }
 
+function isFilterResult<T>(obj: any): obj is IFilterResult<T> {
+	return typeof obj === 'object' && 'visibility' in obj && 'data' in obj;
+}
+
 export interface IFilter<T, TFilterData> {
-	filter(element: T): Visibility | IFilterResult<TFilterData>;
+	getVisibility(element: T): Visibility | IFilterResult<TFilterData>;
 }
 
 function visibleCountReducer<T>(result: number, node: IMutableTreeNode<T, any>): number {
-	return result + (node.collapsed ? 1 : node.visibleCount);
+	return result + (node.collapsed ? 1 : node.revealedCount);
 }
 
 function getVisibleCount<T>(nodes: IMutableTreeNode<T, any>[]): number {
@@ -65,20 +69,20 @@ function getVisibleCount<T>(nodes: IMutableTreeNode<T, any>[]): number {
  * all the visible nodes in an array.
  */
 function updateVisibleCount<T, TFilterData>(node: IMutableTreeNode<T, TFilterData>): ITreeNode<T, TFilterData>[] {
-	const previousVisibleCount = node.visibleCount;
+	const previousVisibleCount = node.revealedCount;
 	const result: ITreeNode<T, TFilterData>[] = [];
 
 	function _updateVisibleCount(node: IMutableTreeNode<T, TFilterData>): number {
 		result.push(node);
-		node.visibleCount = 1;
+		node.revealedCount = 1;
 
 		if (!node.collapsed) {
 			for (const child of node.children) {
-				node.visibleCount += _updateVisibleCount(child);
+				node.revealedCount += _updateVisibleCount(child);
 			}
 		}
 
-		return node.visibleCount;
+		return node.revealedCount;
 	}
 
 	_updateVisibleCount(node);
@@ -87,7 +91,7 @@ function updateVisibleCount<T, TFilterData>(node: IMutableTreeNode<T, TFilterDat
 	node = node.parent;
 
 	while (node) {
-		node.visibleCount += visibleCountDiff;
+		node.revealedCount += visibleCountDiff;
 		node = node.parent;
 	}
 
@@ -104,26 +108,6 @@ function getTreeElementIterator<T>(elements: Iterator<ITreeElement<T>> | ITreeEl
 	}
 }
 
-function treeElementToNode<T, TFilterData>(treeElement: ITreeElement<T>, parent: IMutableTreeNode<T, TFilterData>, visible: boolean, treeListElements: ITreeNode<T, TFilterData>[]): IMutableTreeNode<T, TFilterData> {
-	const depth = parent.depth + 1;
-	const { element, collapsible, collapsed } = treeElement;
-	const node = { parent, element, children: [], depth, collapsible: !!collapsible, collapsed: !!collapsed, visibleCount: 1, visible: true, filterData: undefined };
-
-	if (visible) {
-		treeListElements.push(node);
-	}
-
-	const children = getTreeElementIterator(treeElement.children);
-	node.children = Iterator.collect(Iterator.map(children, el => treeElementToNode(el, node, visible && !treeElement.collapsed, treeListElements)));
-	node.collapsible = node.collapsible || node.children.length > 0;
-
-	if (!collapsed) {
-		node.visibleCount += getVisibleCount(node.children);
-	}
-
-	return node;
-}
-
 function treeNodeToElement<T>(node: IMutableTreeNode<T, any>): ITreeElement<T> {
 	const { element, collapsed } = node;
 	const children = Iterator.map(Iterator.fromArray(node.children), treeNodeToElement);
@@ -142,6 +126,10 @@ export function getNodeLocation<T>(node: ITreeNode<T, any>): number[] {
 	return location.reverse();
 }
 
+export interface ITreeOptions<T, TFilterData = void> {
+	filter?: IFilter<T, TFilterData>;
+}
+
 export class TreeModel<T, TFilterData = void> {
 
 	private root: IMutableTreeNode<T, TFilterData> = {
@@ -151,7 +139,7 @@ export class TreeModel<T, TFilterData = void> {
 		depth: 0,
 		collapsible: false,
 		collapsed: false,
-		visibleCount: 1,
+		revealedCount: 1,
 		visible: true,
 		filterData: undefined
 	};
@@ -159,45 +147,79 @@ export class TreeModel<T, TFilterData = void> {
 	private _onDidChangeCollapseState = new Emitter<ITreeNode<T, TFilterData>>();
 	readonly onDidChangeCollapseState: Event<ITreeNode<T, TFilterData>> = this._onDidChangeCollapseState.event;
 
-	constructor(private list: ISpliceable<ITreeNode<T, TFilterData>>) { }
+	private filter?: IFilter<T, TFilterData>;
+
+	constructor(private list: ISpliceable<ITreeNode<T, TFilterData>>, options: ITreeOptions<T, TFilterData> = {}) {
+		this.filter = options.filter;
+	}
 
 	splice(location: number[], deleteCount: number, toInsert?: ISequence<ITreeElement<T>>): Iterator<ITreeElement<T>> {
 		if (location.length === 0) {
 			throw new Error('Invalid tree location');
 		}
 
-		const { parentNode, listIndex, visible } = this.findParentNode(location);
+		const { parentNode, listIndex, revealed } = this.findParentNode(location);
 		const treeListElementsToInsert: ITreeNode<T, TFilterData>[] = [];
 		const elementsToInsert = getTreeElementIterator(toInsert);
-		const nodesToInsert = Iterator.collect(Iterator.map(elementsToInsert, el => treeElementToNode(el, parentNode, visible, treeListElementsToInsert)));
+		const nodesToInsert = Iterator.collect(Iterator.map(elementsToInsert, el => this.createTreeNode(el, parentNode, revealed, treeListElementsToInsert)));
 		const lastIndex = location[location.length - 1];
 		const deletedNodes = parentNode.children.splice(lastIndex, deleteCount, ...nodesToInsert);
 		const visibleDeleteCount = getVisibleCount(deletedNodes);
 
-		parentNode.visibleCount += getVisibleCount(nodesToInsert) - visibleDeleteCount;
+		parentNode.revealedCount += getVisibleCount(nodesToInsert) - visibleDeleteCount;
 
-		if (visible) {
+		if (revealed) {
 			this.list.splice(listIndex, visibleDeleteCount, treeListElementsToInsert);
 		}
 
 		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;
+		}
+
+		const node = { parent, element, children: [], depth, collapsible: !!collapsible, collapsed: !!collapsed, revealedCount: 1, visible, filterData };
+
+		if (revealed) {
+			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;
 	}
 
 	setCollapsed(location: number[], collapsed: boolean): boolean {
-		const { node, listIndex, visible } = this.findNode(location);
-		return this._setCollapsed(node, listIndex, visible, collapsed);
+		const { node, listIndex, revealed } = this.findNode(location);
+		return this._setCollapsed(node, listIndex, revealed, collapsed);
 	}
 
 	toggleCollapsed(location: number[]): void {
-		const { node, listIndex, visible } = this.findNode(location);
-		this._setCollapsed(node, listIndex, visible);
+		const { node, listIndex, revealed } = this.findNode(location);
+		this._setCollapsed(node, listIndex, revealed);
 	}
 
-	private _setCollapsed(node: IMutableTreeNode<T, TFilterData>, listIndex: number, visible: boolean, collapsed?: boolean | undefined): boolean {
+	private _setCollapsed(node: IMutableTreeNode<T, TFilterData>, listIndex: number, revealed: boolean, collapsed?: boolean | undefined): boolean {
 		if (!node.collapsible) {
 			return false;
 		}
@@ -212,8 +234,8 @@ export class TreeModel<T, TFilterData = void> {
 
 		node.collapsed = collapsed;
 
-		if (visible) {
-			const previousVisibleCount = node.visibleCount;
+		if (revealed) {
+			const previousVisibleCount = node.revealedCount;
 			const toInsert = updateVisibleCount(node);
 
 			this.list.splice(listIndex + 1, previousVisibleCount - 1, toInsert.slice(1));
@@ -231,8 +253,8 @@ export class TreeModel<T, TFilterData = void> {
 
 			while (queue.length > 0) {
 				const node = queue.shift();
-				const visible = listIndex < this.root.children.length;
-				this._setCollapsed(node, listIndex, visible, collapsed);
+				const revealed = listIndex < this.root.children.length;
+				this._setCollapsed(node, listIndex, revealed, collapsed);
 
 				queue.push(...node.children);
 				listIndex++;
@@ -244,8 +266,8 @@ export class TreeModel<T, TFilterData = void> {
 		return this.findNode(location).node.collapsed;
 	}
 
-	private findNode(location: number[]): { node: IMutableTreeNode<T, TFilterData>, listIndex: number, visible: boolean } {
-		const { parentNode, listIndex, visible } = this.findParentNode(location);
+	private findNode(location: number[]): { node: IMutableTreeNode<T, TFilterData>, listIndex: number, revealed: boolean } {
+		const { parentNode, listIndex, revealed } = this.findParentNode(location);
 		const index = location[location.length - 1];
 
 		if (index < 0 || index > parentNode.children.length) {
@@ -254,10 +276,10 @@ export class TreeModel<T, TFilterData = void> {
 
 		const node = parentNode.children[index];
 
-		return { node, listIndex, visible };
+		return { node, listIndex, revealed };
 	}
 
-	private findParentNode(location: number[], node: IMutableTreeNode<T, TFilterData> = this.root, listIndex: number = 0, visible = true): { parentNode: IMutableTreeNode<T, TFilterData>; listIndex: number; visible: boolean; } {
+	private findParentNode(location: number[], node: IMutableTreeNode<T, TFilterData> = this.root, listIndex: number = 0, revealed = true): { parentNode: IMutableTreeNode<T, TFilterData>; listIndex: number; revealed: boolean; } {
 		const [index, ...rest] = location;
 
 		if (index < 0 || index > node.children.length) {
@@ -266,15 +288,15 @@ export class TreeModel<T, TFilterData = void> {
 
 		// TODO@joao perf!
 		for (let i = 0; i < index; i++) {
-			listIndex += node.children[i].visibleCount;
+			listIndex += node.children[i].revealedCount;
 		}
 
-		visible = visible && !node.collapsed;
+		revealed = revealed && !node.collapsed;
 
 		if (rest.length === 0) {
-			return { parentNode: node, listIndex, visible };
+			return { parentNode: node, listIndex, revealed };
 		}
 
-		return this.findParentNode(rest, node.children[index], listIndex + 1, visible);
+		return this.findParentNode(rest, node.children[index], listIndex + 1, revealed);
 	}
 }
\ No newline at end of file