objectTreeModel preserves collapse state on strict identity

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

@@ -8,20 +8,14 @@ import { Iterator, ISequence, getSequenceIterator } from 'vs/base/common/iterato
 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> {
-	identityProvider?: IIdentityProvider<T>;
-}
+export interface IObjectTreeModelOptions<T, TFilterData> extends IIndexTreeModelOptions<T, TFilterData> { }
 
 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>>;
 
@@ -29,39 +23,23 @@ 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;
 	}
 
 	setChildren(
 		element: T | null,
-		// TODO@joao also use identity provider to preserve children when `children` is undefined!!!
 		children: ISequence<ITreeElement<T>> | undefined,
 		onDidCreateNode?: (node: ITreeNode<T, TFilterData>) => void,
 		onDidDeleteNode?: (node: ITreeNode<T, TFilterData>) => void
 	): 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);
 			}
@@ -72,43 +50,40 @@ 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);
 			}
 		};
 
-		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)
-				};
-			});
-		};
+		return this.model.splice(
+			[...location, 0],
+			Number.MAX_VALUE,
+			this.preserveCollapseState(children),
+			_onDidCreateNode,
+			_onDidDeleteNode
+		);
+	}
 
-		if (children && this.identityProvider) {
-			children = preserveCollapseState(children);
-		}
+	private preserveCollapseState(elements: ISequence<ITreeElement<T>> | undefined): ISequence<ITreeElement<T>> {
+		const iterator = elements ? getSequenceIterator(elements) : Iterator.empty<ITreeElement<T>>();
 
-		const result = this.model.splice([...location, 0], Number.MAX_VALUE, children, _onDidCreateNode, _onDidDeleteNode);
+		return Iterator.map(iterator, treeElement => {
+			const node = this.nodes.get(treeElement.element);
 
-		if (createdNodesByIdentity) {
-			createdNodesByIdentity.forEach((node, identity) => this.nodesByIdentity!.set(identity, node));
-		}
+			if (!node) {
+				return treeElement;
+			}
+
+			const collapsible = typeof treeElement.collapsible === 'boolean' ? treeElement.collapsible : node.collapsible;
+			const collapsed = typeof treeElement.collapsed !== 'undefined' ? treeElement.collapsed : (collapsible && node.collapsed);
 
-		return result;
+			return {
+				...treeElement,
+				collapsible,
+				collapsed,
+				children: this.preserveCollapseState(treeElement.children)
+			};
+		});
 	}
 
 	getParentElement(ref: T | null = null): T | null {

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

@@ -141,7 +141,7 @@ suite('ObjectTreeModel', function () {
 		assert.deepEqual(toArray(list), [1, 11, 111, 112, 2]);
 	});
 
-	test('collapse state is lost without an identity provider', () => {
+	test('collapse state is preserved with strict identity', () => {
 		const list: ITreeNode<string>[] = [];
 		const model = new ObjectTreeModel<string>(toSpliceable(list), { collapseByDefault: true });
 		const data = [{ element: 'father', children: [{ element: 'child' }] }];
@@ -153,22 +153,19 @@ suite('ObjectTreeModel', function () {
 		assert.deepEqual(toArray(list), ['father', 'child']);
 
 		model.setChildren(null, data);
-		assert.deepEqual(toArray(list), ['father']);
-	});
+		assert.deepEqual(toArray(list), ['father', 'child']);
 
-	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' }] }];
+		const data2 = [{ element: 'father', children: [{ element: 'child' }] }, { element: 'uncle' }];
+		model.setChildren(null, data2);
+		assert.deepEqual(toArray(list), ['father', 'child', 'uncle']);
 
-		model.setChildren(null, data);
-		assert.deepEqual(toArray(list), ['father']);
+		model.setChildren(null, [{ element: 'uncle' }]);
+		assert.deepEqual(toArray(list), ['uncle']);
 
-		model.setCollapsed('father', false);
-		assert.deepEqual(toArray(list), ['father', 'child']);
+		model.setChildren(null, data2);
+		assert.deepEqual(toArray(list), ['father', 'uncle']);
 
 		model.setChildren(null, data);
-		assert.deepEqual(toArray(list), ['father', 'child']);
+		assert.deepEqual(toArray(list), ['father']);
 	});
 });
\ No newline at end of file