From e2c0dc0eb2bd28d441b21e1f7ff43a4eb803f01f Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Tue, 11 Jan 2022 16:01:43 -0800
Subject: [PATCH] testing: persist collapse state more correctly

---
 src/vs/base/browser/ui/tree/abstractTree.ts   | 22 +++++++++++++------
 src/vs/base/browser/ui/tree/dataTree.ts       |  2 +-
 .../hierarchalByLocation.ts                   |  4 +++-
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/src/vs/base/browser/ui/tree/abstractTree.ts b/src/vs/base/browser/ui/tree/abstractTree.ts
index c8846433dc4..559225d9e71 100644
--- a/src/vs/base/browser/ui/tree/abstractTree.ts
+++ b/src/vs/base/browser/ui/tree/abstractTree.ts
@@ -237,14 +237,14 @@ interface ITreeListTemplateData<T> {
 export interface IAbstractTreeViewState {
 	readonly focus: Iterable<string>;
 	readonly selection: Iterable<string>;
-	readonly expanded: Iterable<string>;
+	readonly expanded: { [id: string]: 1 | 0 };
 	readonly scrollTop: number;
 }
 
 export class AbstractTreeViewState implements IAbstractTreeViewState {
 	public readonly focus: Set<string>;
 	public readonly selection: Set<string>;
-	public readonly expanded: Set<string>;
+	public readonly expanded: { [id: string]: 1 | 0 };
 	public scrollTop: number;
 
 	public static lift(state: IAbstractTreeViewState) {
@@ -255,7 +255,7 @@ export class AbstractTreeViewState implements IAbstractTreeViewState {
 		return new AbstractTreeViewState({
 			focus: [],
 			selection: [],
-			expanded: [],
+			expanded: Object.create(null),
 			scrollTop,
 		});
 	}
@@ -263,7 +263,15 @@ export class AbstractTreeViewState implements IAbstractTreeViewState {
 	protected constructor(state: IAbstractTreeViewState) {
 		this.focus = new Set(state.focus);
 		this.selection = new Set(state.selection);
-		this.expanded = new Set(state.expanded);
+		if (state.expanded instanceof Array) { // old format
+			this.expanded = Object.create(null);
+			for (const id of state.expanded as string[]) {
+				this.expanded[id] = 1;
+			}
+		} else {
+			this.expanded = state.expanded;
+		}
+		this.expanded = state.expanded;
 		this.scrollTop = state.scrollTop;
 	}
 
@@ -271,7 +279,7 @@ export class AbstractTreeViewState implements IAbstractTreeViewState {
 		return {
 			focus: Array.from(this.focus),
 			selection: Array.from(this.selection),
-			expanded: Array.from(this.expanded),
+			expanded: this.expanded,
 			scrollTop: this.scrollTop,
 		};
 	}
@@ -1747,8 +1755,8 @@ export abstract class AbstractTree<T, TFilterData, TRef> implements IDisposable
 		while (queue.length > 0) {
 			const node = queue.shift()!;
 
-			if (node !== root && node.collapsible && !node.collapsed) {
-				state.expanded.add(getId(node.element!));
+			if (node !== root && node.collapsible) {
+				state.expanded[getId(node.element!)] = node.collapsed ? 0 : 1;
 			}
 
 			queue.push(...node.children);
diff --git a/src/vs/base/browser/ui/tree/dataTree.ts b/src/vs/base/browser/ui/tree/dataTree.ts
index 2fdf125b7d3..262313f2a2e 100644
--- a/src/vs/base/browser/ui/tree/dataTree.ts
+++ b/src/vs/base/browser/ui/tree/dataTree.ts
@@ -63,7 +63,7 @@ export class DataTree<TInput, T, TFilterData = void> extends AbstractTree<T | nu
 
 		const isCollapsed = (element: T) => {
 			const id = this.identityProvider!.getId(element).toString();
-			return !viewState.expanded.has(id);
+			return !viewState.expanded[id];
 		};
 
 		const onDidCreateNode = (node: ITreeNode<T, TFilterData>) => {
diff --git a/src/vs/workbench/contrib/testing/browser/explorerProjections/hierarchalByLocation.ts b/src/vs/workbench/contrib/testing/browser/explorerProjections/hierarchalByLocation.ts
index 718fe06c4e6..ee6b60a369f 100644
--- a/src/vs/workbench/contrib/testing/browser/explorerProjections/hierarchalByLocation.ts
+++ b/src/vs/workbench/contrib/testing/browser/explorerProjections/hierarchalByLocation.ts
@@ -235,7 +235,9 @@ export class HierarchicalByLocationProjection extends Disposable implements ITes
 		return {
 			element: node,
 			collapsible: node.test.expand !== TestItemExpandState.NotExpandable,
-			collapsed: !this.lastState.expanded.has(node.treeId) && node.test.expand === TestItemExpandState.Expandable ? true : undefined,
+			collapsed: this.lastState.expanded[node.treeId] !== undefined
+				? !this.lastState.expanded[node.treeId]
+				: node.test.expand === TestItemExpandState.Expandable ? true : undefined,
 			children: recurse(node.children),
 		};
 	};