feat(cdk/tree): add trackBy to NestedTreeControl (#19602)

Author: tjaneczko

src/cdk/tree/control/nested-tree-control.spec.ts

@@ -4,14 +4,15 @@ import {NestedTreeControl} from './nested-tree-control';
 
 
 describe('CdkNestedTreeControl', () => {
-  let treeControl: NestedTreeControl<TestData>;
   let getChildren = (node: TestData) => observableOf(node.children);
 
-  beforeEach(() => {
-    treeControl = new NestedTreeControl<TestData>(getChildren);
-  });
-
   describe('base tree control actions', () => {
+    let treeControl: NestedTreeControl<TestData>;
+
+    beforeEach(() => {
+      treeControl = new NestedTreeControl<TestData>(getChildren);
+    });
+
     it('should be able to expand and collapse dataNodes', () => {
       const nodes = generateData(10, 4);
       const node = nodes[1];
@@ -196,6 +197,24 @@ describe('CdkNestedTreeControl', () => {
       });
     });
   });
+
+  it('maintains node expansion state based on trackBy function, if provided', () => {
+    const treeControl = new NestedTreeControl<TestData, string>(
+        getChildren, {trackBy: (node: TestData) => `${node.a} ${node.b} ${node.c}`});
+
+    const nodes = generateData(2, 2);
+    const secondNode = nodes[1];
+    treeControl.dataNodes = nodes;
+
+    treeControl.expand(secondNode);
+    expect(treeControl.isExpanded(secondNode)).toBeTruthy('Expect second node to be expanded');
+
+    // Replace the second node with a brand new instance with same hash
+    nodes[1] = new TestData(
+        secondNode.a, secondNode.b, secondNode.c, secondNode.level, secondNode.children);
+    expect(treeControl.isExpanded(nodes[1])).toBeTruthy('Expect second node to still be expanded');
+  });
+
 });
 
 export class TestData {

src/cdk/tree/control/nested-tree-control.ts

@@ -9,12 +9,22 @@ import {Observable, isObservable} from 'rxjs';
 import {take, filter} from 'rxjs/operators';
 import {BaseTreeControl} from './base-tree-control';
 
-/** Nested tree control. Able to expand/collapse a subtree recursively for NestedNode type. */
-export class NestedTreeControl<T> extends BaseTreeControl<T> {
+/** Optional set of configuration that can be provided to the NestedTreeControl. */
+export interface NestedTreeControlOptions<T, K> {
+  trackBy?: (dataNode: T) => K;
+}
 
+/** Nested tree control. Able to expand/collapse a subtree recursively for NestedNode type. */
+export class NestedTreeControl<T, K = T> extends BaseTreeControl<T, K> {
   /** Construct with nested tree function getChildren. */
-  constructor(public getChildren: (dataNode: T) => (Observable<T[]> | T[] | undefined | null)) {
+  constructor(
+      public getChildren: (dataNode: T) => (Observable<T[]>| T[] | undefined | null),
+      public options?: NestedTreeControlOptions<T, K>) {
     super();
+
+    if (this.options) {
+      this.trackBy = this.options.trackBy;
+    }
   }
 
   /**
@@ -27,7 +37,7 @@ export class NestedTreeControl<T> extends BaseTreeControl<T> {
     this.expansionModel.clear();
     const allNodes = this.dataNodes.reduce((accumulator: T[], dataNode) =>
         [...accumulator, ...this.getDescendants(dataNode), dataNode], []);
-    this.expansionModel.select(...allNodes);
+    this.expansionModel.select(...allNodes.map(node => this._trackByValue(node)));
   }
 
   /** Gets a list of descendant dataNodes of a subtree rooted at given data node recursively. */

tools/public_api_guard/cdk/tree.d.ts

@@ -161,14 +161,19 @@ export declare function getTreeMultipleDefaultNodeDefsError(): Error;
 
 export declare function getTreeNoValidDataSourceError(): Error;
 
-export declare class NestedTreeControl<T> extends BaseTreeControl<T> {
+export declare class NestedTreeControl<T, K = T> extends BaseTreeControl<T, K> {
     getChildren: (dataNode: T) => (Observable<T[]> | T[] | undefined | null);
-    constructor(getChildren: (dataNode: T) => (Observable<T[]> | T[] | undefined | null));
+    options?: NestedTreeControlOptions<T, K> | undefined;
+    constructor(getChildren: (dataNode: T) => (Observable<T[]> | T[] | undefined | null), options?: NestedTreeControlOptions<T, K> | undefined);
     protected _getDescendants(descendants: T[], dataNode: T): void;
     expandAll(): void;
     getDescendants(dataNode: T): T[];
 }
 
+export interface NestedTreeControlOptions<T, K> {
+    trackBy?: (dataNode: T) => K;
+}
+
 export interface TreeControl<T, K = T> {
     dataNodes: T[];
     expansionModel: SelectionModel<K>;