feat(tree): add trackBy to tree controls to keep the expansion of rebuilt nodes (#DS-5785)

packages/components/tree-select/tree-select.component.spec.ts

@@ -9,6 +9,7 @@ import {
     InputSignalWithTransform,
     OnInit,
     Provider,
+    Signal,
     Type,
     effect,
     signal,
@@ -5879,7 +5880,8 @@ describe('KbqTreeSelect', () => {
         });
 
         // The types are checked by `check-spec-types`: a query that names the data node gets the model typed by it,
-        // the tree's included, one that does not keeps `any`.
+        // the tree's included, one that does not keeps `any`. One that names the key of the tree control gets the tree
+        // typed by it too.
         it('should share a selection model typed by the data node with the tree', () => {
             const select = fixture.componentInstance.select();
             const tree = select.tree()!;
@@ -5892,6 +5894,12 @@ describe('KbqTreeSelect', () => {
             expectTypeOf(select.selectionModel).toEqualTypeOf<SelectionModel<FileFlatNode>>();
             expectTypeOf(select.selected).toEqualTypeOf<FileFlatNode | FileFlatNode[]>();
             expectTypeOf(tree.selectionModel).toEqualTypeOf<SelectionModel<FileFlatNode>>();
+            expectTypeOf<KbqTreeSelect<FileFlatNode>['tree']>().toEqualTypeOf<
+                Signal<KbqTreeSelection<FileFlatNode, FileFlatNode> | undefined>
+            >();
+            expectTypeOf<KbqTreeSelect<FileFlatNode, string>['tree']>().toEqualTypeOf<
+                Signal<KbqTreeSelection<FileFlatNode, string> | undefined>
+            >();
             expect(tree.selectionModel).toBe(select.selectionModel);
             expect(select.selectionModel.selected.map(({ name }) => name).sort()).toEqual(['Downloads', 'rootNode_1']);
         });

packages/components/tree-select/tree-select.component.ts

@@ -259,7 +259,7 @@ export class KbqTreeSelectChange<T = any> {
     ],
     exportAs: 'kbqTreeSelect'
 })
-export class KbqTreeSelect<T = any>
+export class KbqTreeSelect<T = any, K = T>
     extends KbqAbstractSelect
     implements
         AfterViewChecked,
@@ -319,7 +319,8 @@ export class KbqTreeSelect<T = any>
      * Deals with the selection logic. The projected `kbq-tree-selection` shares the same instance.
      *
      * `T` is the type of the data nodes, the same as the projected tree's. A template infers it from `[clearPredicate]`
-     * alone, apart from the tree; a query names it: `viewChild.required<KbqTreeSelect<Node>>(KbqTreeSelect)`.
+     * alone, apart from the tree; a query names it: `viewChild.required<KbqTreeSelect<Node>>(KbqTreeSelect)`, and a
+     * select whose tree control has a `trackBy` names the key type too: `KbqTreeSelect<Node, string>`.
      */
     selectionModel: SelectionModel<T>;
 
@@ -405,7 +406,7 @@ export class KbqTreeSelect<T = any>
 
     readonly customTagTemplateRef = contentChild('kbqSelectTagContent', { read: TemplateRef });
 
-    readonly tree = contentChild<KbqTreeSelection<T>>(KbqTreeSelection);
+    readonly tree = contentChild<KbqTreeSelection<T, K>>(KbqTreeSelection);
 
     readonly search = contentChild(KbqSelectSearch);
 
@@ -467,7 +468,7 @@ export class KbqTreeSelect<T = any>
      * Event emitted whenever the tree's `selectAllOptions()` runs — a click on the `selectAll` master
      * checkbox, the Ctrl/Cmd + A shortcut, or a custom `selectAllHandler` that delegates to it.
      */
-    readonly onSelectAll = output<KbqSelectAllEvent<KbqTreeOption<T>, KbqTreeSelect<T>>>();
+    readonly onSelectAll = output<KbqSelectAllEvent<KbqTreeOption<T>, KbqTreeSelect<T, K>>>();
 
     /** Classes to be passed to the select panel. Supports the same syntax as the `[class]` binding. */
     readonly panelClass = input<
@@ -667,8 +668,8 @@ export class KbqTreeSelect<T = any>
 
     /** @docs-private */
     readonly selectAllHandlerInput = input<
-        ((event: KeyboardEvent, select: KbqTreeSelect<T>) => void) | undefined,
-        ((event: KeyboardEvent, select: KbqTreeSelect<T>) => void) | undefined
+        ((event: KeyboardEvent, select: KbqTreeSelect<T, K>) => void) | undefined,
+        ((event: KeyboardEvent, select: KbqTreeSelect<T, K>) => void) | undefined
     >(undefined, {
         alias: 'selectAllHandler',
         transform: kbqSelectAllHandlerTransform
@@ -677,12 +678,12 @@ export class KbqTreeSelect<T = any>
     /**
      * Function for handling the combination Ctrl + A (select all). By default, the internal handler is used.
      */
-    get selectAllHandler(): (event: KeyboardEvent, select: KbqTreeSelect<T>) => void {
+    get selectAllHandler(): (event: KeyboardEvent, select: KbqTreeSelect<T, K>) => void {
         return this.selectAllHandlerInput() ?? this.defaultSelectAllHandler;
     }
 
     /** Function for handling the combination Ctrl + A (select all). By default, the internal handler is used. */
-    private defaultSelectAllHandler(event: KeyboardEvent, select: KbqTreeSelect<T>): void {
+    private defaultSelectAllHandler(event: KeyboardEvent, select: KbqTreeSelect<T, K>): void {
         const searchInput = isInput(event) ? (event.target as HTMLInputElement) : null;
 
         if (shouldSelectSearchText(searchInput)) {

packages/components/tree/control/base-tree-control.ts

@@ -3,11 +3,19 @@ import { BehaviorSubject, Observable } from 'rxjs';
 import { TreeControl } from './tree-control';
 
 /** Base tree control. It has basic toggle/expand/collapse operations on a single data node. */
-export abstract class BaseTreeControl<T> implements TreeControl<T> {
+export abstract class BaseTreeControl<T, K = T> implements TreeControl<T, K> {
     dataNodes: T[];
 
-    /** A selection model with multi-selection to track expansion status. */
-    expansionModel: SelectionModel<T> = new SelectionModel<T>(true);
+    /** A selection model with multi-selection to track expansion status: it holds the key of every expanded node. */
+    expansionModel: SelectionModel<K> = new SelectionModel<K>(true);
+
+    /**
+     * Gives the key a data node is expanded by. Without it the node object is its own key, so a node
+     * rebuilt as a new object — as a flattener does whenever its data changes — comes back collapsed.
+     *
+     * Set it before anything is expanded: keys stored under another function no longer match.
+     */
+    trackBy?: (dataNode: T) => K;
 
     filterModel: SelectionModel<T> = new SelectionModel<T>(true);
 
@@ -37,7 +45,7 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
             return;
         }
 
-        this.expansionModel.toggle(dataNode);
+        this.expansionModel.toggle(this.expansionKey(dataNode));
     }
 
     /** Expands one single data node. */
@@ -46,7 +54,7 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
             return;
         }
 
-        this.expansionModel.select(dataNode);
+        this.expansionModel.select(this.expansionKey(dataNode));
     }
 
     /** Collapses one single data node. */
@@ -55,17 +63,17 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
             return;
         }
 
-        this.expansionModel.deselect(dataNode);
+        this.expansionModel.deselect(this.expansionKey(dataNode));
     }
 
     /** Whether a given data node is expanded or not. Returns true if the data node is expanded. */
     isExpanded(dataNode: T): boolean {
-        return this.expansionModel.isSelected(dataNode);
+        return this.expansionModel.isSelected(this.expansionKey(dataNode));
     }
 
     /** Toggles a subtree rooted at `node` recursively. */
     toggleDescendants(dataNode: T): void {
-        if (this.expansionModel.isSelected(dataNode)) {
+        if (this.isExpanded(dataNode)) {
             this.collapseDescendants(dataNode);
         } else {
             this.expandDescendants(dataNode);
@@ -82,14 +90,19 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
         const toBeProcessed = [dataNode];
 
         toBeProcessed.push(...this.getDescendants(dataNode));
-        this.expansionModel.select(...toBeProcessed);
+        this.expansionModel.select(...toBeProcessed.map((node) => this.expansionKey(node)));
     }
 
     /** Collapses a subtree rooted at given data node recursively. */
     collapseDescendants(dataNode: T): void {
         const toBeProcessed = [dataNode];
 
         toBeProcessed.push(...this.getDescendants(dataNode));
-        this.expansionModel.deselect(...toBeProcessed);
+        this.expansionModel.deselect(...toBeProcessed.map((node) => this.expansionKey(node)));
+    }
+
+    /** The key `expansionModel` holds for a data node. */
+    protected expansionKey(dataNode: T): K {
+        return this.trackBy ? this.trackBy(dataNode) : (dataNode as unknown as K);
     }
 }

packages/components/tree/control/flat-tree-control.filters.ts

@@ -12,10 +12,10 @@ export interface FlatTreeControlFilter<T> {
  * Default first stage of the filter pipeline: keeps the nodes whose view value matches the query.
  * @docs-private
  */
-export class FilterByViewValue<T> implements FlatTreeControlFilter<T> {
+export class FilterByViewValue<T, K = T> implements FlatTreeControlFilter<T> {
     result: T[];
 
-    constructor(private control: FlatTreeControl<T>) {}
+    constructor(private control: FlatTreeControl<T, K>) {}
 
     handle(value: string | null): T[] {
         const viewValue = value || null;
@@ -39,10 +39,10 @@ export class FilterByViewValue<T> implements FlatTreeControlFilter<T> {
  * matches stay reachable in the rendered hierarchy.
  * @docs-private
  */
-export class FilterParentsForNodes<T> implements FlatTreeControlFilter<T> {
+export class FilterParentsForNodes<T, K = T> implements FlatTreeControlFilter<T> {
     result: T[];
 
-    constructor(private control: FlatTreeControl<T>) {}
+    constructor(private control: FlatTreeControl<T, K>) {}
 
     handle(_, prevFilter: FlatTreeControlFilter<T>): T[] {
         const result: Set<T> = new Set();
@@ -61,12 +61,12 @@ export class FilterParentsForNodes<T> implements FlatTreeControlFilter<T> {
     }
 }
 
-export class FilterByValues<T> implements FlatTreeControlFilter<T> {
+export class FilterByValues<T, K = T> implements FlatTreeControlFilter<T> {
     result: T[];
 
     private values: string[] = [];
 
-    constructor(private control: FlatTreeControl<T>) {}
+    constructor(private control: FlatTreeControl<T, K>) {}
 
     handle(_, prevFilter?: FlatTreeControlFilter<T>): T[] {
         const found = this.control.dataNodes.filter((node) => this.values.includes(this.control.getValue(node)));

packages/components/tree/control/flat-tree-control.spec.ts

@@ -253,6 +253,121 @@ describe('FlatTreeControl', () => {
             });
         });
     });
+
+    describe('with trackBy', () => {
+        let control: FlatTreeControl<FlatNode, string>;
+        let nodes: FlatNode[];
+
+        beforeEach(() => {
+            nodes = buildFlatNodes();
+            control = new FlatTreeControl<FlatNode, string>(
+                (node) => node.level,
+                (node) => node.expandable,
+                (node) => node.name,
+                (node) => node.name,
+                undefined,
+                undefined,
+                undefined,
+                { trackBy: (node) => node.name }
+            );
+            control.dataNodes = nodes;
+        });
+
+        it('should keep a node expanded once the node is rebuilt as a new object', () => {
+            control.expand(byName(nodes, 'documents'));
+
+            control.dataNodes = buildFlatNodes();
+
+            expect(control.isExpanded(byName(control.dataNodes, 'documents'))).toBe(true);
+        });
+
+        it('should collapse a node through an object rebuilt after it was expanded', () => {
+            control.expand(byName(nodes, 'documents'));
+
+            control.collapse(byName(buildFlatNodes(), 'documents'));
+
+            expect(control.expansionModel.selected).toEqual([]);
+        });
+
+        it('should hold the keys rather than the nodes in the expansion model', () => {
+            control.toggle(byName(nodes, 'root'));
+
+            expect(control.expansionModel.selected).toEqual(['root']);
+        });
+
+        it('should expand and collapse all the nodes by key', () => {
+            control.expandAll();
+
+            expect(control.expansionModel.selected).toEqual(nodes.map((node) => node.name));
+
+            control.collapseAll();
+
+            expect(control.expansionModel.selected).toEqual([]);
+        });
+
+        it('should toggle a subtree by key, also once the subtree is rebuilt', () => {
+            control.toggleDescendants(byName(nodes, 'documents'));
+
+            expect(control.expansionModel.selected).toEqual(['documents', 'specs', 'draft', 'readme']);
+
+            control.dataNodes = buildFlatNodes();
+            control.toggleDescendants(byName(control.dataNodes, 'documents'));
+
+            expect(control.expansionModel.selected).toEqual([]);
+        });
+
+        it('should find the descendants and parents of a node through an object it was rebuilt from', () => {
+            const stale = byName(nodes, 'documents');
+
+            control.dataNodes = buildFlatNodes();
+
+            expect(control.getDescendants(stale)).toEqual(
+                ['specs', 'draft', 'readme'].map((name) => byName(control.dataNodes, name))
+            );
+            expect(control.getParents(byName(nodes, 'draft'))).toEqual(
+                ['root', 'documents', 'specs'].map((name) => byName(control.dataNodes, name))
+            );
+
+            control.expandDescendants(stale);
+
+            expect(control.expansionModel.selected).toEqual(['documents', 'specs', 'draft', 'readme']);
+        });
+
+        it('should expand the expandable nodes a filter keeps by key', () => {
+            control.filterNodes('draft');
+
+            expect(control.expansionModel.selected).toEqual(['root', 'documents', 'specs']);
+        });
+
+        it('should restore the keys expanded before a filter once the query is dropped', async () => {
+            control.expand(byName(nodes, 'documents'));
+
+            control.filterNodes('draft');
+            await Promise.resolve();
+
+            control.filterNodes('');
+            await Promise.resolve();
+
+            expect(control.expansionModel.selected).toEqual(['documents']);
+        });
+
+        it('should not carry the expansion over to a rebuilt node without trackBy', () => {
+            const plain = new FlatTreeControl<FlatNode>(
+                (node) => node.level,
+                (node) => node.expandable,
+                (node) => node.name,
+                (node) => node.name
+            );
+
+            plain.dataNodes = nodes;
+            plain.expand(byName(nodes, 'documents'));
+
+            plain.dataNodes = buildFlatNodes();
+
+            expect(plain.isExpanded(byName(plain.dataNodes, 'documents'))).toBe(false);
+            expect(plain.expansionModel.selected).toEqual([byName(nodes, 'documents')]);
+        });
+    });
 });
 
 /** Flat node in the depth-first order the flattener produces, with 0-based levels. */