support an optional trackBy function in FlatTreeControl

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;
@@ -135,14 +137,19 @@ export declare class CdkTreeNodeToggle<T> {
     static ɵfac: i0.ɵɵFactoryDef<CdkTreeNodeToggle<any>>;
 }
 
-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;
@@ -161,9 +168,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;