object tree model: preserve collapse with identity provider

fixes #63149

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

@@ -4,18 +4,24 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { ISpliceable } from 'vs/base/common/sequence';
-import { Iterator, ISequence } from 'vs/base/common/iterator';
+import { Iterator, ISequence, getSequenceIterator } from 'vs/base/common/iterator';
 import { IndexTreeModel, IIndexTreeModelOptions } from 'vs/base/browser/ui/tree/indexTreeModel';
 import { Event } from 'vs/base/common/event';
 import { ITreeModel, ITreeNode, ITreeElement } from 'vs/base/browser/ui/tree/tree';
+import { IIdentityProvider } from 'vs/base/browser/ui/list/list';
 
-export interface IObjectTreeModelOptions<T, TFilterData> extends IIndexTreeModelOptions<T, TFilterData> { }
+export interface IObjectTreeModelOptions<T, TFilterData> extends IIndexTreeModelOptions<T, TFilterData> {
+	identityProvider?: IIdentityProvider<T>;
+}
 
 export class ObjectTreeModel<T extends NonNullable<any>, TFilterData = void> implements ITreeModel<T, TFilterData, T> {
 
 	private model: IndexTreeModel<T, TFilterData>;
 	private nodes = new Map<T, ITreeNode<T, TFilterData>>();
 
+	private identityProvider: IIdentityProvider<T> | undefined = undefined;
+	private nodesByIdentity: Map<string, ITreeNode<T, TFilterData>> | undefined = undefined;
+
 	readonly onDidChangeCollapseState: Event<ITreeNode<T, TFilterData>>;
 	readonly onDidChangeRenderNodeCount: Event<ITreeNode<T, TFilterData>>;
 
@@ -23,6 +29,12 @@ export class ObjectTreeModel<T extends NonNullable<any>, TFilterData = void> imp
 
 	constructor(list: ISpliceable<ITreeNode<T, TFilterData>>, options: IObjectTreeModelOptions<T, TFilterData> = {}) {
 		this.model = new IndexTreeModel(list, options);
+		this.identityProvider = options.identityProvider;
+
+		if (this.identityProvider) {
+			this.nodesByIdentity = new Map();
+		}
+
 		this.onDidChangeCollapseState = this.model.onDidChangeCollapseState;
 		this.onDidChangeRenderNodeCount = this.model.onDidChangeRenderNodeCount;
 	}
@@ -35,11 +47,20 @@ export class ObjectTreeModel<T extends NonNullable<any>, TFilterData = void> imp
 	): Iterator<ITreeElement<T>> {
 		const location = this.getElementLocation(element);
 		const insertedElements = new Set<T>();
+		let createdNodesByIdentity: Map<string, ITreeNode<T, TFilterData>> | undefined;
+
+		if (this.identityProvider) {
+			createdNodesByIdentity = new Map();
+		}
 
 		const _onDidCreateNode = (node: ITreeNode<T, TFilterData>) => {
 			insertedElements.add(node.element);
 			this.nodes.set(node.element, node);
 
+			if (createdNodesByIdentity) {
+				createdNodesByIdentity.set(this.identityProvider!.getId(node.element).toString(), node);
+			}
+
 			if (onDidCreateNode) {
 				onDidCreateNode(node);
 			}
@@ -50,12 +71,43 @@ export class ObjectTreeModel<T extends NonNullable<any>, TFilterData = void> imp
 				this.nodes.delete(node.element);
 			}
 
+			if (this.nodesByIdentity) {
+				this.nodesByIdentity.delete(this.identityProvider!.getId(node.element).toString());
+			}
+
 			if (onDidDeleteNode) {
 				onDidDeleteNode(node);
 			}
 		};
 
-		return this.model.splice([...location, 0], Number.MAX_VALUE, children, _onDidCreateNode, _onDidDeleteNode);
+		const preserveCollapseState = (elements: ISequence<ITreeElement<T>>): ISequence<ITreeElement<T>> => {
+			const iterator = getSequenceIterator(elements);
+
+			return Iterator.map(iterator, treeElement => {
+				const identity = this.identityProvider!.getId(treeElement.element).toString();
+				const node = this.nodesByIdentity!.get(identity);
+				const collapsible = typeof treeElement.collapsible === 'boolean' ? treeElement.collapsible : (typeof treeElement.collapsed !== 'undefined');
+
+				return {
+					...treeElement,
+					collapsible,
+					collapsed: typeof treeElement.collapsed === 'undefined' ? (node ? (collapsible && node.collapsed) : undefined) : treeElement.collapsed,
+					children: treeElement.children && preserveCollapseState(treeElement.children)
+				};
+			});
+		};
+
+		if (children && this.identityProvider) {
+			children = preserveCollapseState(children);
+		}
+
+		const result = this.model.splice([...location, 0], Number.MAX_VALUE, children, _onDidCreateNode, _onDidDeleteNode);
+
+		if (createdNodesByIdentity) {
+			createdNodesByIdentity.forEach((node, identity) => this.nodesByIdentity!.set(identity, node));
+		}
+
+		return result;
 	}
 
 	getParentElement(ref: T | null = null): T | null {

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

@@ -140,4 +140,35 @@ suite('ObjectTreeModel', function () {
 		model.setCollapsed(1, false);
 		assert.deepEqual(toArray(list), [1, 11, 111, 112, 2]);
 	});
+
+	test('collapse state is lost without an identity provider', () => {
+		const list: ITreeNode<string>[] = [];
+		const model = new ObjectTreeModel<string>(toSpliceable(list), { collapseByDefault: true });
+		const data = [{ element: 'father', children: [{ element: 'child' }] }];
+
+		model.setChildren(null, data);
+		assert.deepEqual(toArray(list), ['father']);
+
+		model.setCollapsed('father', false);
+		assert.deepEqual(toArray(list), ['father', 'child']);
+
+		model.setChildren(null, data);
+		assert.deepEqual(toArray(list), ['father']);
+	});
+
+	test('collapse state is preserved with an identity provider', () => {
+		const list: ITreeNode<string>[] = [];
+		const identityProvider = { getId: (name: string) => name };
+		const model = new ObjectTreeModel<string>(toSpliceable(list), { collapseByDefault: true, identityProvider });
+		const data = [{ element: 'father', children: [{ element: 'child' }] }];
+
+		model.setChildren(null, data);
+		assert.deepEqual(toArray(list), ['father']);
+
+		model.setCollapsed('father', false);
+		assert.deepEqual(toArray(list), ['father', 'child']);
+
+		model.setChildren(null, data);
+		assert.deepEqual(toArray(list), ['father', 'child']);
+	});
 });
\ No newline at end of file