feat(tree): Add support for trackBy (#11267)

Add trackBy support improve performance

Author: stasmaxymov

src/cdk/tree/nested-node.ts

@@ -69,7 +69,7 @@ export class CdkNestedTreeNode<T> extends CdkTreeNode<T> implements AfterContent
   }
 
   ngAfterContentInit() {
-    this._dataDiffer = this._differs.find([]).create();
+    this._dataDiffer = this._differs.find([]).create(this._tree.trackBy);
     if (!this._tree.treeControl.getChildren) {
       throw getTreeControlFunctionsMissingError();
     }

src/cdk/tree/tree.md

@@ -166,3 +166,13 @@ nodes should be sent to tree component based on current expansion status.
 
 The data source for nested tree has an option to leave the node expansion/collapsing event for each
 tree node component to handle.
+
+##### `trackBy`
+
+To improve performance, a `trackBy` function can be provided to the tree similar to Angular’s
+[`ngFor` `trackBy`](https://angular.io/api/common/NgForOf#change-propagation). This informs the
+tree how to uniquely identify nodes to track how the data changes with each update.
+
+```html
+<cdk-tree [dataSource]="dataSource" [treeControl]="treeControl" [trackBy]="trackByFn">
+```

src/cdk/tree/tree.spec.ts

@@ -6,7 +6,7 @@
  * found in the LICENSE file at https://angular.io/license
  */
 import {ComponentFixture, TestBed, fakeAsync, flush} from '@angular/core/testing';
-import {Component, ViewChild} from '@angular/core';
+import {Component, ViewChild, TrackByFunction} from '@angular/core';
 
 import {CollectionViewer, DataSource} from '@angular/cdk/collections';
 import {combineLatest, BehaviorSubject, Observable} from 'rxjs';
@@ -296,6 +296,97 @@ describe('CdkTree', () => {
           [`[topping_3] - [cheese_3] + [base_3]`]);
       });
     });
+
+    describe('with trackBy', () => {
+      let fixture: ComponentFixture<CdkTreeAppWithTrackBy>;
+      let component: CdkTreeAppWithTrackBy;
+
+      function createTrackByTestComponent(trackByStrategy: 'reference' | 'property' | 'index') {
+        configureCdkTreeTestingModule([CdkTreeAppWithTrackBy]);
+        fixture = TestBed.createComponent(CdkTreeAppWithTrackBy);
+        component = fixture.componentInstance;
+        component.trackByStrategy = trackByStrategy;
+        dataSource = component.dataSource as FakeDataSource;
+        tree = component.tree;
+        treeElement = fixture.nativeElement.querySelector('cdk-tree');
+
+        fixture.detectChanges();
+
+        // Each node receives an attribute 'initialIndex' the element's original place
+        getNodes(treeElement).forEach((node: Element, index: number) => {
+          node.setAttribute('initialIndex', index.toString());
+        });
+
+        // Prove that the attributes match their indicies
+        const initialNodes = getNodes(treeElement);
+        expect(initialNodes[0].getAttribute('initialIndex')).toBe('0');
+        expect(initialNodes[1].getAttribute('initialIndex')).toBe('1');
+        expect(initialNodes[2].getAttribute('initialIndex')).toBe('2');
+      }
+
+      function mutateData() {
+        // Swap first and second data in data array
+        const copiedData = component.dataSource.data.slice();
+        const temp = copiedData[0];
+        copiedData[0] = copiedData[1];
+        copiedData[1] = temp;
+
+        // Remove the third element
+        copiedData.splice(2, 1);
+
+        // Add new data
+        component.dataSource.data = copiedData;
+        component.dataSource.addData();
+      }
+
+      it('should add/remove/move nodes with reference-based trackBy', () => {
+        createTrackByTestComponent('reference');
+        mutateData();
+
+        // Expect that the first and second nodes were swapped and that the last node is new
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(3);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('0');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe(null);
+      });
+
+      it('should add/remove/move nodes with property-based trackBy', () => {
+        createTrackByTestComponent('property');
+        mutateData();
+
+        // Change each item reference to show that the trackby is checking the item properties.
+        // Otherwise this would cause them all to be removed/added.
+        component.dataSource.data = component.dataSource.data
+            .map(item => new TestData(item.pizzaTopping, item.pizzaCheese, item.pizzaBase, ));
+
+        // Expect that the first and second nodes were swapped and that the last node is new
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(3);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('0');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe(null);
+      });
+
+      it('should add/remove/move nodes with index-based trackBy', () => {
+        createTrackByTestComponent('index');
+        mutateData();
+
+        // Change each item reference to show that the trackby is checking the index.
+        // Otherwise this would cause them all to be removed/added.
+        component.dataSource.data = component.dataSource.data
+            .map(item => new TestData(item.pizzaTopping, item.pizzaCheese, item.pizzaBase, ));
+
+        // Expect first two to be the same since they were swapped but indicies are consistent.
+        // The third element was removed and caught by the tree so it was removed before another
+        // item was added, so it is without an initial index.
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(3);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('0');
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe(null);
+      });
+    });
   });
 
   describe('nested tree', () => {
@@ -599,6 +690,128 @@ describe('CdkTree', () => {
       });
     });
 
+    describe('with trackBy', () => {
+      let fixture: ComponentFixture<NestedCdkTreeAppWithTrackBy>;
+      let component: NestedCdkTreeAppWithTrackBy;
+
+      function createTrackByTestComponent(trackByStrategy: 'reference' | 'property' | 'index') {
+        configureCdkTreeTestingModule([NestedCdkTreeAppWithTrackBy]);
+        fixture = TestBed.createComponent(NestedCdkTreeAppWithTrackBy);
+        component = fixture.componentInstance;
+        component.trackByStrategy = trackByStrategy;
+        dataSource = component.dataSource as FakeDataSource;
+
+        tree = component.tree;
+        treeElement = fixture.nativeElement.querySelector('cdk-tree');
+
+        fixture.detectChanges();
+
+        // Each node receives an attribute 'initialIndex' the element's original place
+        getNodes(treeElement).forEach((node: Element, index: number) => {
+          node.setAttribute('initialIndex', index.toString());
+        });
+
+        // Prove that the attributes match their indicies
+        const initialNodes = getNodes(treeElement);
+        expect(initialNodes.length).toBe(3);
+        initialNodes.forEach((node, index) => {
+          expect(node.getAttribute('initialIndex')).toBe(`${index}`);
+        });
+
+        const parent = dataSource.data[0];
+        dataSource.addChild(parent, false);
+        dataSource.addChild(parent, false);
+        dataSource.addChild(parent, false);
+        getNodes(initialNodes[0]).forEach((node: Element, index: number) => {
+          node.setAttribute('initialIndex', `c${index}`);
+        });
+        getNodes(initialNodes[0]).forEach((node, index) => {
+          expect(node.getAttribute('initialIndex')).toBe(`c${index}`);
+        });
+      }
+
+      function mutateChildren(parent: TestData) {
+        // Swap first and second data in data array
+        const copiedData = parent.children.slice();
+        const temp = copiedData[0];
+        copiedData[0] = copiedData[1];
+        copiedData[1] = temp;
+
+        // Remove the third element
+        copiedData.splice(2, 1);
+
+        // Add new data
+        parent.children = copiedData;
+        parent.observableChildren.next(copiedData);
+        component.dataSource.addChild(parent, false);
+      }
+
+      it('should add/remove/move children nodes with reference-based trackBy', () => {
+        createTrackByTestComponent('reference');
+        mutateChildren(dataSource.data[0]);
+
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(6);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('0');
+
+        // Expect that the first and second child nodes were swapped and that the last node is new
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('c1');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe('c0');
+        expect(changedNodes[3].getAttribute('initialIndex')).toBe(null);
+
+        expect(changedNodes[4].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[5].getAttribute('initialIndex')).toBe('2');
+      });
+
+      it('should add/remove/move children nodes with property-based trackBy', () => {
+        createTrackByTestComponent('property');
+        mutateChildren(dataSource.data[0]);
+
+        // Change each item reference to show that the trackby is checking the item properties.
+        // Otherwise this would cause them all to be removed/added.
+        dataSource.data[0].observableChildren.next(dataSource.data[0].children
+          .map(item => new TestData(item.pizzaTopping, item.pizzaCheese, item.pizzaBase)));
+
+        // Expect that the first and second nodes were swapped and that the last node is new
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(6);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('0');
+
+        // Expect that the first and second child nodes were swapped and that the last node is new
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('c1');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe('c0');
+        expect(changedNodes[3].getAttribute('initialIndex')).toBe(null);
+
+        expect(changedNodes[4].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[5].getAttribute('initialIndex')).toBe('2');
+      });
+
+      it('should add/remove/move children nodes with index-based trackBy', () => {
+        createTrackByTestComponent('index');
+        mutateChildren(dataSource.data[0]);
+
+        // Change each item reference to show that the trackby is checking the index.
+        // Otherwise this would cause them all to be removed/added.
+        dataSource.data[0].observableChildren.next(dataSource.data[0].children
+          .map(item => new TestData(item.pizzaTopping, item.pizzaCheese, item.pizzaBase)));
+
+        const changedNodes = getNodes(treeElement);
+        expect(changedNodes.length).toBe(6);
+        expect(changedNodes[0].getAttribute('initialIndex')).toBe('0');
+
+        // Expect first two children to be the same since they were swapped
+        // but indicies are consistent.
+        // The third element was removed and caught by the tree so it was removed before another
+        // item was added, so it is without an initial index.
+        expect(changedNodes[1].getAttribute('initialIndex')).toBe('c0');
+        expect(changedNodes[2].getAttribute('initialIndex')).toBe('c1');
+        expect(changedNodes[3].getAttribute('initialIndex')).toBe(null);
+
+        expect(changedNodes[4].getAttribute('initialIndex')).toBe('1');
+        expect(changedNodes[5].getAttribute('initialIndex')).toBe('2');
+      });
+    });
+
     it('should throw an error when missing function in nested tree', fakeAsync(() => {
       configureCdkTreeTestingModule([NestedCdkErrorTreeApp]);
       expect(() => {
@@ -1144,3 +1357,66 @@ class DepthNestedCdkTreeApp {
 
   @ViewChild(CdkTree) tree: CdkTree<TestData>;
 }
+
+@Component({
+  template: `
+    <cdk-tree [dataSource]="dataSource" [treeControl]="treeControl" [trackBy]="trackByFn">
+      <cdk-tree-node *cdkTreeNodeDef="let node" class="customNodeClass">
+                     {{node.pizzaTopping}} - {{node.pizzaCheese}} + {{node.pizzaBase}}
+      </cdk-tree-node>
+    </cdk-tree>
+  `
+})
+class CdkTreeAppWithTrackBy {
+  trackByStrategy: 'reference' | 'property' | 'index' = 'reference';
+
+  trackByFn: TrackByFunction<TestData> = (index, item) => {
+    switch (this.trackByStrategy) {
+      case 'reference': return item;
+      case 'property': return item.pizzaBase;
+      case 'index': return index;
+    }
+  }
+
+  getLevel = (node: TestData) => node.level;
+  isExpandable = (node: TestData) => node.children.length > 0;
+
+  treeControl: TreeControl<TestData> = new FlatTreeControl(this.getLevel, this.isExpandable);
+  dataSource: FakeDataSource = new FakeDataSource(this.treeControl);
+
+  @ViewChild(CdkTree) tree: CdkTree<TestData>;
+}
+
+@Component({
+  template: `
+    <cdk-tree [dataSource]="dataArray" [treeControl]="treeControl" [trackBy]="trackByFn">
+      <cdk-nested-tree-node *cdkTreeNodeDef="let node">
+           [{{node.pizzaTopping}}] - [{{node.pizzaCheese}}] + [{{node.pizzaBase}}]
+         <ng-template cdkTreeNodeOutlet></ng-template>
+      </cdk-nested-tree-node>
+    </cdk-tree>
+  `
+})
+class NestedCdkTreeAppWithTrackBy {
+  trackByStrategy: 'reference' | 'property' | 'index' = 'reference';
+
+  trackByFn: TrackByFunction<TestData> = (index, item) => {
+    switch (this.trackByStrategy) {
+      case 'reference': return item;
+      case 'property': return item.pizzaBase;
+      case 'index': return index;
+    }
+  }
+
+  getChildren = (node: TestData) => node.observableChildren;
+
+  treeControl: TreeControl<TestData> = new NestedTreeControl(this.getChildren);
+
+  dataSource: FakeDataSource = new FakeDataSource(this.treeControl);
+
+  get dataArray() {
+    return this.dataSource.data;
+  }
+
+  @ViewChild(CdkTree) tree: CdkTree<TestData>;
+}

src/cdk/tree/tree.ts

@@ -24,7 +24,8 @@ import {
   QueryList,
   ViewChild,
   ViewContainerRef,
-  ViewEncapsulation
+  ViewEncapsulation,
+  TrackByFunction
 } from '@angular/core';
 import {BehaviorSubject, Observable, of as observableOf, Subject, Subscription} from 'rxjs';
 import {takeUntil} from 'rxjs/operators';
@@ -165,6 +166,14 @@ export class CdkTree<T>
   /** The tree controller */
   @Input() treeControl: TreeControl<T>;
 
+  /**
+   * Tracking function that will be used to check the differences in data changes. Used similarly
+   * to `ngFor` `trackBy` function. Optimize node operations by identifying a node based on its data
+   * relative to the function to know if a node should be added/removed/moved.
+   * Accepts a function that takes two parameters, `index` and `item`.
+   */
+  @Input() trackBy: TrackByFunction<T>;
+
   // Outlets within the tree's template where the dataNodes will be inserted.
   @ViewChild(CdkTreeNodeOutlet) _nodeOutlet: CdkTreeNodeOutlet;
 
@@ -184,7 +193,7 @@ export class CdkTree<T>
               private _changeDetectorRef: ChangeDetectorRef) {}
 
   ngOnInit() {
-    this._dataDiffer = this._differs.find([]).create();
+    this._dataDiffer = this._differs.find([]).create(this.trackBy);
     if (!this.treeControl) {
       throw getTreeControlMissingError();
     }