treeModel: setCollapsed

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

@@ -37,14 +37,26 @@ function getVisibleCount<T>(nodes: ITreeNode<T>[]): number {
 	return nodes.reduce(visibleCountReducer, 0);
 }
 
-function treeElementToNode<T>(treeElement: ITreeElement<T>, depth: number, visible: boolean, visibleElements: ITreeListElement<T>[]): ITreeNode<T> {
+function getVisibleNodes<T>(nodes: ITreeNode<T>[], result: ITreeListElement<T>[] = []): ITreeListElement<T>[] {
+	for (const node of nodes) {
+		result.push(node);
+
+		if (!node.collapsed) {
+			getVisibleNodes(node.children, result);
+		}
+	}
+
+	return result;
+}
+
+function treeElementToNode<T>(treeElement: ITreeElement<T>, depth: number, visible: boolean, treeListElements: ITreeListElement<T>[]): ITreeNode<T> {
 	const { element, collapsed } = treeElement;
 
 	if (visible) {
-		visibleElements.push({ element, collapsed, depth });
+		treeListElements.push({ element, collapsed, depth });
 	}
 
-	const children = collect(map(treeElement.children, el => treeElementToNode(el, depth + 1, visible && !treeElement.collapsed, visibleElements)));
+	const children = collect(map(treeElement.children, el => treeElementToNode(el, depth + 1, visible && !treeElement.collapsed, treeListElements)));
 	const visibleCount = 1 + getVisibleCount(children);
 
 	return { element, children, depth, collapsed, visibleCount };
@@ -75,20 +87,61 @@ export class TreeModel<T> {
 		}
 
 		const { parentNode, listIndex, visible } = this.findParentNode(location);
-		const listToInsert: ITreeListElement<T>[] = [];
-		const nodesToInsert = collect(map(toInsert, el => treeElementToNode(el, parentNode.depth + 1, visible, listToInsert)));
+		const treeListElementsToInsert: ITreeListElement<T>[] = [];
+		const nodesToInsert = collect(map(toInsert, el => treeElementToNode(el, parentNode.depth + 1, visible, treeListElementsToInsert)));
 		const deletedNodes = parentNode.children.splice(last(location), deleteCount, ...nodesToInsert);
 		const visibleDeleteCount = getVisibleCount(deletedNodes);
 
 		parentNode.visibleCount += getVisibleCount(nodesToInsert) - visibleDeleteCount;
 
 		if (visible) {
-			this.list.splice(listIndex, visibleDeleteCount, listToInsert);
+			this.list.splice(listIndex, visibleDeleteCount, treeListElementsToInsert);
 		}
 
 		return map(iter(deletedNodes), treeNodeToElement);
 	}
 
+	setCollapsed(location: number[], collapsed: boolean): void {
+		const { node, listIndex, visible } = this.findNode(location);
+
+		if (node.collapsed === collapsed) {
+			return;
+		}
+
+		node.collapsed = collapsed;
+
+		if (visible) {
+			if (collapsed) {
+				const deleteCount = getVisibleCount(node.children);
+				const { element, depth } = node;
+
+				this.list.splice(listIndex, 1 + deleteCount, [{ element, collapsed, depth }]);
+			} else {
+				const toInsert = [node, ...getVisibleNodes(node.children)];
+
+				this.list.splice(listIndex, 1, toInsert);
+			}
+		}
+	}
+
+	isCollapsed(location: number[]): boolean {
+		const { node } = this.findNode(location);
+		return node.collapsed;
+	}
+
+	private findNode(location: number[]): { node: ITreeNode<T>, listIndex: number, visible: boolean } {
+		const { parentNode, listIndex, visible } = this.findParentNode(location);
+		const index = last(location);
+
+		if (index < 0 || index > parentNode.children.length) {
+			throw new Error('Invalid tree location');
+		}
+
+		const node = parentNode.children[index];
+
+		return { node, listIndex, visible };
+	}
+
 	private findParentNode(location: number[], node: ITreeNode<T> = this.root, listIndex: number = 0, visible = true): { parentNode: ITreeNode<T>; listIndex: number; visible: boolean; } {
 		const [index, ...rest] = location;
 

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

@@ -223,4 +223,75 @@ suite('TreeModel2', () => {
 		model.splice([0, 0], 2);
 		assert.deepEqual(list.length, 3);
 	});
+
+	test('collapse', () => {
+		const list = [] as ITreeListElement<number>[];
+		const model = new TreeModel<number>(toSpliceable(list));
+
+		model.splice([0], 0, iter([
+			{
+				element: 0, collapsed: false, children: iter([
+					{ element: 10, collapsed: false, children: iter([]) },
+					{ element: 11, collapsed: false, children: iter([]) },
+					{ element: 12, collapsed: false, children: iter([]) },
+				])
+			},
+			{ element: 1, collapsed: false, children: iter([]) },
+			{ element: 2, collapsed: false, children: iter([]) }
+		]));
+
+		assert.deepEqual(list.length, 6);
+
+		model.setCollapsed([0], true);
+		assert.deepEqual(list.length, 3);
+		assert.deepEqual(list[0].element, 0);
+		assert.deepEqual(list[0].collapsed, true);
+		assert.deepEqual(list[0].depth, 1);
+		assert.deepEqual(list[1].element, 1);
+		assert.deepEqual(list[1].collapsed, false);
+		assert.deepEqual(list[1].depth, 1);
+		assert.deepEqual(list[2].element, 2);
+		assert.deepEqual(list[2].collapsed, false);
+		assert.deepEqual(list[2].depth, 1);
+	});
+
+	test('expand', () => {
+		const list = [] as ITreeListElement<number>[];
+		const model = new TreeModel<number>(toSpliceable(list));
+
+		model.splice([0], 0, iter([
+			{
+				element: 0, collapsed: true, children: iter([
+					{ element: 10, collapsed: false, children: iter([]) },
+					{ element: 11, collapsed: false, children: iter([]) },
+					{ element: 12, collapsed: false, children: iter([]) },
+				])
+			},
+			{ element: 1, collapsed: false, children: iter([]) },
+			{ element: 2, collapsed: false, children: iter([]) }
+		]));
+
+		assert.deepEqual(list.length, 3);
+
+		model.setCollapsed([0], false);
+		assert.deepEqual(list.length, 6);
+		assert.deepEqual(list[0].element, 0);
+		assert.deepEqual(list[0].collapsed, false);
+		assert.deepEqual(list[0].depth, 1);
+		assert.deepEqual(list[1].element, 10);
+		assert.deepEqual(list[1].collapsed, false);
+		assert.deepEqual(list[1].depth, 2);
+		assert.deepEqual(list[2].element, 11);
+		assert.deepEqual(list[2].collapsed, false);
+		assert.deepEqual(list[2].depth, 2);
+		assert.deepEqual(list[3].element, 12);
+		assert.deepEqual(list[3].collapsed, false);
+		assert.deepEqual(list[3].depth, 2);
+		assert.deepEqual(list[4].element, 1);
+		assert.deepEqual(list[4].collapsed, false);
+		assert.deepEqual(list[4].depth, 1);
+		assert.deepEqual(list[5].element, 2);
+		assert.deepEqual(list[5].collapsed, false);
+		assert.deepEqual(list[5].depth, 1);
+	});
 });