tree: fix expand/collapse

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

@@ -60,6 +60,7 @@ class TreeDelegate<T> implements IVirtualDelegate<ITreeNode<T>> {
 
 interface ITreeListTemplateData<T> {
 	twistie: HTMLElement;
+	count: HTMLElement;
 	templateData: T;
 }
 
@@ -89,9 +90,10 @@ class TreeRenderer<T, TTemplateData> implements IRenderer<ITreeNode<T>, ITreeLis
 		const el = append(container, $('.monaco-tl-row'));
 		const twistie = append(el, $('.tl-twistie'));
 		const contents = append(el, $('.tl-contents'));
+		const count = append(el, $('.tl-count'));
 		const templateData = this.renderer.renderTemplate(contents);
 
-		return { twistie, templateData };
+		return { twistie, count, templateData };
 	}
 
 	renderElement(node: ITreeNode<T>, index: number, templateData: ITreeListTemplateData<TTemplateData>): void {
@@ -99,6 +101,7 @@ class TreeRenderer<T, TTemplateData> implements IRenderer<ITreeNode<T>, ITreeLis
 
 		templateData.twistie.style.width = `${10 + node.depth * 10}px`;
 		renderTwistie(node, templateData.twistie);
+		templateData.count.textContent = `${node.visibleCount}`;
 
 		this.renderer.renderElement(node.element, index, templateData.templateData);
 	}
@@ -119,6 +122,7 @@ class TreeRenderer<T, TTemplateData> implements IRenderer<ITreeNode<T>, ITreeLis
 		}
 
 		renderTwistie(node, templateData.twistie);
+		templateData.count.textContent = `${node.visibleCount}`;
 	}
 
 	dispose(): void {

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

@@ -39,13 +39,35 @@ function getVisibleCount<T>(nodes: IMutableTreeNode<T>[]): number {
 	return nodes.reduce(visibleCountReducer, 0);
 }
 
-function getVisibleNodes<T>(nodes: IMutableTreeNode<T>[], result: ITreeNode<T>[] = []): ITreeNode<T>[] {
-	for (const node of nodes) {
+/**
+ * Recursively updates the visibleCount of a subtree, while collecting
+ * all the visible nodes in an array.
+ */
+function updateVisibleCount<T>(node: IMutableTreeNode<T>): ITreeNode<T>[] {
+	const previousVisibleCount = node.visibleCount;
+	const result: ITreeNode<T>[] = [];
+
+	function _updateVisibleCount(node: IMutableTreeNode<T>): number {
 		result.push(node);
+		node.visibleCount = 1;
 
 		if (!node.collapsed) {
-			getVisibleNodes(node.children, result);
+			for (const child of node.children) {
+				node.visibleCount += _updateVisibleCount(child);
+			}
 		}
+
+		return node.visibleCount;
+	}
+
+	_updateVisibleCount(node);
+
+	const visibleCountDiff = result.length - previousVisibleCount;
+	node = node.parent;
+
+	while (node) {
+		node.visibleCount += visibleCountDiff;
+		node = node.parent;
 	}
 
 	return result;
@@ -108,6 +130,7 @@ export class TreeModel<T> {
 		visibleCount: 1
 	};
 
+	// TODO@joao can't we do without this?
 	private _onDidChangeCollapseState = new Emitter<ITreeNode<T>>();
 	readonly onDidChangeCollapseState: Event<ITreeNode<T>> = this._onDidChangeCollapseState.event;
 
@@ -165,30 +188,14 @@ export class TreeModel<T> {
 		node.collapsed = collapsed;
 
 		if (visible) {
-			this._onDidChangeCollapseState.fire(node);
+			const previousVisibleCount = node.visibleCount;
+			const toInsert = updateVisibleCount(node);
 
-			let visibleCountDiff: number;
-
-			if (collapsed) {
-				const deleteCount = getVisibleCount(node.children);
-
-				this.list.splice(listIndex + 1, deleteCount, []);
-				visibleCountDiff = -deleteCount;
-			} else {
-				const toInsert = getVisibleNodes(node.children);
-
-				this.list.splice(listIndex + 1, 0, toInsert);
-				visibleCountDiff = toInsert.length;
-			}
-
-			let mutableNode = node;
-
-			while (mutableNode) {
-				mutableNode.visibleCount += visibleCountDiff;
-				mutableNode = mutableNode.parent;
-			}
+			this.list.splice(listIndex + 1, previousVisibleCount - 1, toInsert.slice(1));
 		}
 
+		this._onDidChangeCollapseState.fire(node);
+
 		return true;
 	}
 

test/tree/public/index.html

@@ -13,6 +13,10 @@
 		.monaco-scrollable-element>.scrollbar>.slider {
 			background: rgba(100, 100, 100, .4);
 		}
+
+		.tl-contents {
+			flex: 1;
+		}
 	</style>
 </head>
 

test/tree/server.js

@@ -17,7 +17,7 @@ async function getTree(fsPath, level) {
 	const element = path.basename(fsPath);
 	const stat = await fs.stat(fsPath);
 
-	if (!stat.isDirectory() || element === '.git' || element === '.build' || level >= 3) {
+	if (!stat.isDirectory() || element === '.git' || element === '.build' || level >= 4) {
 		return { element };
 	}