feat(cdk/tree): support optional trackBy in FlatTreeControl (#18708)

This is similar to trackBy in *ngFor and is important for those who
use the redux pattern (i.e. NGRX) since node references are frequently replaced with modifies instances from the Store (as opposed to changing the state within the node objects
themselves)

Author: arielements

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

@@ -10,7 +10,7 @@ import {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> {
 
   /** Gets a list of descendent data nodes of a subtree rooted at given data node recursively. */
   abstract getDescendants(dataNode: T): T[];
@@ -22,7 +22,15 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
   dataNodes: T[];
 
   /** A selection model with multi-selection to track expansion status. */
-  expansionModel: SelectionModel<T> = new SelectionModel<T>(true);
+  expansionModel: SelectionModel<K> = new SelectionModel<K>(true);
+
+  /**
+   * Returns the identifier by which a dataNode should be tracked, should its
+   * reference change.
+   *
+   * Similar to trackBy for *ngFor
+   */
+  trackBy?: (dataNode: T) => K;
 
   /** Get depth of a given data node, return the level number. This is for flat tree node. */
   getLevel: (dataNode: T) => number;
@@ -38,29 +46,29 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
 
   /** Toggles one single data node's expanded/collapsed state. */
   toggle(dataNode: T): void {
-    this.expansionModel.toggle(dataNode);
+    this.expansionModel.toggle(this._trackByValue(dataNode));
   }
 
   /** Expands one single data node. */
   expand(dataNode: T): void {
-    this.expansionModel.select(dataNode);
+    this.expansionModel.select(this._trackByValue(dataNode));
   }
 
   /** Collapses one single data node. */
   collapse(dataNode: T): void {
-    this.expansionModel.deselect(dataNode);
+    this.expansionModel.deselect(this._trackByValue(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._trackByValue(dataNode));
   }
 
   /** Toggles a subtree rooted at `node` recursively. */
   toggleDescendants(dataNode: T): void {
-    this.expansionModel.isSelected(dataNode)
-        ? this.collapseDescendants(dataNode)
-        : this.expandDescendants(dataNode);
+    this.expansionModel.isSelected(this._trackByValue(dataNode)) ?
+        this.collapseDescendants(dataNode) :
+        this.expandDescendants(dataNode);
   }
 
   /** Collapse all dataNodes in the tree. */
@@ -72,13 +80,17 @@ export abstract class BaseTreeControl<T> implements TreeControl<T> {
   expandDescendants(dataNode: T): void {
     let toBeProcessed = [dataNode];
     toBeProcessed.push(...this.getDescendants(dataNode));
-    this.expansionModel.select(...toBeProcessed);
+    this.expansionModel.select(...toBeProcessed.map(value => this._trackByValue(value)));
   }
 
   /** Collapses a subtree rooted at given data node recursively. */
   collapseDescendants(dataNode: T): void {
     let toBeProcessed = [dataNode];
     toBeProcessed.push(...this.getDescendants(dataNode));
-    this.expansionModel.deselect(...toBeProcessed);
+    this.expansionModel.deselect(...toBeProcessed.map(value => this._trackByValue(value)));
+  }
+
+  protected _trackByValue(value: T|K): K {
+    return this.trackBy ? this.trackBy(value as T) : value as K;
   }
 }

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

@@ -1,15 +1,16 @@
 import {FlatTreeControl} from './flat-tree-control';
 
 describe('CdkFlatTreeControl', () => {
-  let treeControl: FlatTreeControl<TestData>;
   let getLevel = (node: TestData) => node.level;
   let isExpandable = (node: TestData) => node.children && node.children.length > 0;
 
-  beforeEach(() => {
-    treeControl = new FlatTreeControl<TestData>(getLevel, isExpandable);
-  });
-
   describe('base tree control actions', () => {
+    let treeControl: FlatTreeControl<TestData>;
+
+    beforeEach(() => {
+      treeControl = new FlatTreeControl<TestData>(getLevel, isExpandable);
+    });
+
     it('should be able to expand and collapse dataNodes', () => {
       const nodes = generateData(10, 4);
       const secondNode = nodes[1];
@@ -139,6 +140,23 @@ describe('CdkFlatTreeControl', () => {
         .toBe(totalNumber, `Expect ${totalNumber} expanded nodes`);
     });
   });
+
+  it('maintains node expansion state based on trackBy function, if provided', () => {
+    const treeControl = new FlatTreeControl<TestData, string>(getLevel, isExpandable);
+
+    const nodes = generateData(2, 2);
+    const secondNode = nodes[1];
+    treeControl.dataNodes = nodes;
+    treeControl.trackBy = (node: TestData) => `${node.a} ${node.b} ${node.c}`;
+
+    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/flat-tree-control.ts

@@ -8,13 +8,23 @@
 
 import {BaseTreeControl} from './base-tree-control';
 
+/** Optional set of configuration that can be provided to the FlatTreeControl. */
+export interface FlatTreeControlOptions<T, K> {
+  trackBy?: (dataNode: T) => K;
+}
+
 /** Flat tree control. Able to expand/collapse a subtree recursively for flattened tree. */
-export class FlatTreeControl<T> extends BaseTreeControl<T> {
+export class FlatTreeControl<T, K = T> extends BaseTreeControl<T, K> {
 
   /** Construct with flat tree data node functions getLevel and isExpandable. */
-  constructor(public getLevel: (dataNode: T) => number,
-              public isExpandable: (dataNode: T) => boolean) {
+  constructor(
+      public getLevel: (dataNode: T) => number, public isExpandable: (dataNode: T) => boolean,
+      public options?: FlatTreeControlOptions<T, K>) {
     super();
+
+    if (this.options) {
+      this.trackBy = this.options.trackBy;
+    }
   }
 
   /**
@@ -48,6 +58,6 @@ export class FlatTreeControl<T> extends BaseTreeControl<T> {
    * data nodes of the tree.
    */
   expandAll(): void {
-    this.expansionModel.select(...this.dataNodes);
+    this.expansionModel.select(...this.dataNodes.map(node => this._trackByValue(node)));
   }
 }

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

@@ -13,12 +13,12 @@ import {Observable} from 'rxjs';
  * The CDKTree will use this TreeControl to expand/collapse a node.
  * User can also use it outside the `<cdk-tree>` to control the expansion status of the tree.
  */
-export interface TreeControl<T> {
+export interface TreeControl<T, K = T> {
   /** The saved tree nodes data for `expandAll` action. */
   dataNodes: T[];
 
   /** The expansion model */
-  expansionModel: SelectionModel<T>;
+  expansionModel: SelectionModel<K>;
 
   /** Whether the data node is expanded or collapsed. Return true if it's expanded. */
   isExpanded(dataNode: T): boolean;

tools/public_api_guard/cdk/tree.d.ts

@@ -1,9 +1,11 @@
-export declare abstract class BaseTreeControl<T> implements TreeControl<T> {
+export declare abstract class BaseTreeControl<T, K = T> implements TreeControl<T, K> {
     dataNodes: T[];
-    expansionModel: SelectionModel<T>;
+    expansionModel: SelectionModel<K>;
     getChildren: (dataNode: T) => (Observable<T[]> | T[] | undefined | null);
     getLevel: (dataNode: T) => number;
     isExpandable: (dataNode: T) => boolean;
+    trackBy?: (dataNode: T) => K;
+    protected _trackByValue(value: T | K): K;
     collapse(dataNode: T): void;
     collapseAll(): void;
     collapseDescendants(dataNode: T): void;
@@ -136,14 +138,19 @@ export declare class CdkTreeNodeToggle<T> {
     static ɵfac: i0.ɵɵFactoryDef<CdkTreeNodeToggle<any>, never>;
 }
 
-export declare class FlatTreeControl<T> extends BaseTreeControl<T> {
+export declare class FlatTreeControl<T, K = T> extends BaseTreeControl<T, K> {
     getLevel: (dataNode: T) => number;
     isExpandable: (dataNode: T) => boolean;
-    constructor(getLevel: (dataNode: T) => number, isExpandable: (dataNode: T) => boolean);
+    options?: FlatTreeControlOptions<T, K> | undefined;
+    constructor(getLevel: (dataNode: T) => number, isExpandable: (dataNode: T) => boolean, options?: FlatTreeControlOptions<T, K> | undefined);
     expandAll(): void;
     getDescendants(dataNode: T): T[];
 }
 
+export interface FlatTreeControlOptions<T, K> {
+    trackBy?: (dataNode: T) => K;
+}
+
 export declare function getTreeControlFunctionsMissingError(): Error;
 
 export declare function getTreeControlMissingError(): Error;
@@ -162,9 +169,9 @@ export declare class NestedTreeControl<T> extends BaseTreeControl<T> {
     getDescendants(dataNode: T): T[];
 }
 
-export interface TreeControl<T> {
+export interface TreeControl<T, K = T> {
     dataNodes: T[];
-    expansionModel: SelectionModel<T>;
+    expansionModel: SelectionModel<K>;
     readonly getChildren: (dataNode: T) => Observable<T[]> | T[] | undefined | null;
     readonly getLevel: (dataNode: T) => number;
     readonly isExpandable: (dataNode: T) => boolean;