feat(data-table): add trackby (#5097)

src/demo-app/data-table/data-table-demo.html

@@ -4,6 +4,20 @@
     <button md-raised-button (click)="disconnect()" [disabled]="!dataSource">Disconnect Data Source</button>
   </div>
 
+  <div>
+    <button md-raised-button (click)="_peopleDatabase.shuffle(changeReferences)">Randomize Data</button>
+    <md-checkbox [(ngModel)]="changeReferences">Change references</md-checkbox>
+  </div>
+
+  <div class="demo-track-by-select">
+    <div class="demo-track-by-label">Track By</div>
+    <md-radio-group [(ngModel)]="trackByStrategy">
+      <md-radio-button [value]="'reference'">Reference</md-radio-button>
+      <md-radio-button [value]="'id'">ID</md-radio-button>
+      <md-radio-button [value]="'index'">Index</md-radio-button>
+    </md-radio-group>
+  </div>
+
   <div class="demo-highlighter">
     Highlight:
     <md-checkbox (change)="toggleHighlight('first', $event.checked)">First Row</md-checkbox>
@@ -19,7 +33,7 @@
                      (toggleColorColumn)="toggleColorColumn()">
   </table-header-demo>
 
-  <cdk-table #table [dataSource]="dataSource">
+  <cdk-table #table [dataSource]="dataSource" [trackBy]="userTrackBy">
 
     <!-- Column Definition: ID -->
     <ng-container cdkColumnDef="userId">

src/demo-app/data-table/data-table-demo.scss

@@ -17,6 +17,15 @@
     margin: 16px 0;
   }
 
+  .demo-track-by-select {
+    display: flex;
+    align-items: center;
+
+    .demo-track-by-label {
+      margin-right: 8px;
+    }
+  }
+
   .demo-highlighter .mat-checkbox {
     margin: 0 8px;
   }
@@ -102,4 +111,3 @@
     font-family: cursive;
   }
 }
-

src/demo-app/data-table/data-table-demo.ts

@@ -1,9 +1,11 @@
 import {Component} from '@angular/core';
-import {PeopleDatabase} from './people-database';
+import {PeopleDatabase, UserData} from './people-database';
 import {PersonDataSource} from './person-data-source';
 
 export type UserProperties = 'userId' | 'userName' | 'progress' | 'color' | undefined;
 
+export type TrackByStrategy = 'id' | 'reference' | 'index';
+
 @Component({
   moduleId: module.id,
   selector: 'data-table-demo',
@@ -13,9 +15,11 @@ export type UserProperties = 'userId' | 'userName' | 'progress' | 'color' | unde
 export class DataTableDemo {
   dataSource: PersonDataSource | null;
   propertiesToDisplay: UserProperties[] = [];
+  trackByStrategy: TrackByStrategy = 'reference';
+  changeReferences = false;
   highlights = new Set<string>();
 
-  constructor(private _peopleDatabase: PeopleDatabase) {
+  constructor(public _peopleDatabase: PeopleDatabase) {
     this.connect();
   }
 
@@ -35,6 +39,14 @@ export class DataTableDemo {
     return distanceFromMiddle / 50 + .3;
   }
 
+  userTrackBy = (index: number, item: UserData) => {
+    switch (this.trackByStrategy) {
+      case 'id': return item.id;
+      case 'reference': return item;
+      case 'index': return index;
+    }
+  }
+
   toggleColorColumn() {
     let colorColumnIndex = this.propertiesToDisplay.indexOf('color');
     if (colorColumnIndex == -1) {

src/demo-app/data-table/people-database.ts

@@ -1,6 +1,7 @@
 import {Injectable} from '@angular/core';
 import {NAMES} from '../dataset/names';
 import {COLORS} from '../dataset/colors';
+import {BehaviorSubject} from 'rxjs/BehaviorSubject';
 
 export let LATEST_ID: number = 0;
 
@@ -13,7 +14,9 @@ export interface UserData {
 
 @Injectable()
 export class PeopleDatabase {
-  data: UserData[] = [];
+  dataChange: BehaviorSubject<UserData[]> = new BehaviorSubject<UserData[]>([]);
+
+  get data(): UserData[] { return this.dataChange.value; }
 
   constructor() {
     this.initialize();
@@ -21,20 +24,44 @@ export class PeopleDatabase {
 
   initialize() {
     LATEST_ID = 0;
-    this.data = [];
+    this.dataChange.next([]);
     for (let i = 0; i < 100; i++) { this.addPerson(); }
   }
 
+  shuffle(changeReferences: boolean) {
+    let copiedData = this.data.slice();
+    for (let i = copiedData.length; i; i--) {
+      let j = Math.floor(Math.random() * i);
+      [copiedData[i - 1], copiedData[j]] = [copiedData[j], copiedData[i - 1]];
+    }
+
+    if (changeReferences) {
+      copiedData = copiedData.map(userData => {
+        return {
+          id: userData.id,
+          name: userData.name,
+          progress: userData.progress,
+          color: userData.color
+        };
+      });
+    }
+
+    this.dataChange.next(copiedData);
+  }
+
   addPerson() {
     const name =
         NAMES[Math.round(Math.random() * (NAMES.length - 1))] + ' ' +
         NAMES[Math.round(Math.random() * (NAMES.length - 1))].charAt(0) + '.';
 
-    this.data.push({
+    const copiedData = this.data.slice();
+    copiedData.push({
       id: (++LATEST_ID).toString(),
       name: name,
       progress: Math.round(Math.random() * 100).toString(),
       color: COLORS[Math.round(Math.random() * (COLORS.length - 1))]
     });
+
+    this.dataChange.next(copiedData);
   }
 }

src/demo-app/data-table/person-data-source.ts

@@ -10,7 +10,12 @@ export class PersonDataSource extends DataSource<any> {
   }
 
   connect(collectionViewer: CollectionViewer): Observable<UserData[]> {
-    return collectionViewer.viewChange.map((view: {start: number, end: number}) => {
+    const changeStreams = Observable.combineLatest(
+        collectionViewer.viewChange,
+        this._peopleDatabase.dataChange);
+    return changeStreams.map((result: any[]) => {
+      const view: {start: number, end: number} = result[0];
+
       // Set the rendered rows length to the virtual page size. Fill in the data provided
       // from the index start until the end index or pagination size, whichever is smaller.
       this._renderedData.length = this._peopleDatabase.data.length;

src/lib/core/data-table/data-table.spec.ts

@@ -25,7 +25,8 @@ describe('CdkTable', () => {
         SimpleCdkTableApp,
         DynamicDataSourceCdkTableApp,
         CustomRoleCdkTableApp,
-        RowContextCdkTableApp
+        TrackByCdkTableApp,
+        RowContextCdkTableApp,
       ],
     }).compileComponents();
   }));
@@ -145,6 +146,117 @@ describe('CdkTable', () => {
     expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
   });
 
+  describe('with trackBy', () => {
+
+    let trackByComponent: TrackByCdkTableApp;
+    let trackByFixture: ComponentFixture<TrackByCdkTableApp>;
+
+    function createTestComponentWithTrackyByTable(trackByStrategy) {
+      trackByFixture = TestBed.createComponent(TrackByCdkTableApp);
+
+      trackByComponent = trackByFixture.componentInstance;
+      trackByComponent.trackByStrategy = trackByStrategy;
+
+      dataSource = trackByComponent.dataSource as FakeDataSource;
+      table = trackByComponent.table;
+      tableElement = trackByFixture.nativeElement.querySelector('cdk-table');
+
+      trackByFixture.detectChanges();  // Let the component and table create embedded views
+      trackByFixture.detectChanges();  // Let the cells render
+
+      // Each row receives an attribute 'initialIndex' the element's original place
+      getRows(tableElement).forEach((row: Element, index: number) => {
+        row.setAttribute('initialIndex', index.toString());
+      });
+
+      // Prove that the attributes match their indicies
+      const initialRows = getRows(tableElement);
+      expect(initialRows[0].getAttribute('initialIndex')).toBe('0');
+      expect(initialRows[1].getAttribute('initialIndex')).toBe('1');
+      expect(initialRows[2].getAttribute('initialIndex')).toBe('2');
+    }
+
+    // Swap first two elements, remove the third, add new data
+    function mutateData() {
+      // Swap first and second data in data array
+      const copiedData = trackByComponent.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
+      trackByComponent.dataSource.data = copiedData;
+      trackByComponent.dataSource.addData();
+    }
+
+    it('should add/remove/move rows with reference-based trackBy', () => {
+      createTestComponentWithTrackyByTable('reference');
+      mutateData();
+
+      // Expect that the first and second rows were swapped and that the last row is new
+      const changedRows = getRows(tableElement);
+      expect(changedRows.length).toBe(3);
+      expect(changedRows[0].getAttribute('initialIndex')).toBe('1');
+      expect(changedRows[1].getAttribute('initialIndex')).toBe('0');
+      expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
+    });
+
+    it('should add/remove/move rows with changed references without property-based trackBy', () => {
+      createTestComponentWithTrackyByTable('reference');
+      mutateData();
+
+      // Change each item reference to show that the trackby is not checking the item properties.
+      trackByComponent.dataSource.data = trackByComponent.dataSource.data
+          .map(item => { return {a: item.a, b: item.b, c: item.c}; });
+
+      // Expect that all the rows are considered new since their references are all different
+      const changedRows = getRows(tableElement);
+      expect(changedRows.length).toBe(3);
+      expect(changedRows[0].getAttribute('initialIndex')).toBe(null);
+      expect(changedRows[1].getAttribute('initialIndex')).toBe(null);
+      expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
+    });
+
+    it('should add/remove/move rows with changed references with property-based trackBy', () => {
+      createTestComponentWithTrackyByTable('propertyA');
+      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.
+      trackByComponent.dataSource.data = trackByComponent.dataSource.data
+          .map(item => { return {a: item.a, b: item.b, c: item.c}; });
+
+      // Expect that the first and second rows were swapped and that the last row is new
+      const changedRows = getRows(tableElement);
+      expect(changedRows.length).toBe(3);
+      expect(changedRows[0].getAttribute('initialIndex')).toBe('1');
+      expect(changedRows[1].getAttribute('initialIndex')).toBe('0');
+      expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
+    });
+
+    it('should add/remove/move rows with changed references with index-based trackBy', () => {
+      createTestComponentWithTrackyByTable('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.
+      trackByComponent.dataSource.data = trackByComponent.dataSource.data
+          .map(item => { return {a: item.a, b: item.b, c: item.c}; });
+
+      // Expect first two to be the same since they were swapped but indicies are consistent.
+      // The third element was removed and caught by the table so it was removed before another
+      // item was added, so it is without an initial index.
+      const changedRows = getRows(tableElement);
+      expect(changedRows.length).toBe(3);
+      expect(changedRows[0].getAttribute('initialIndex')).toBe('0');
+      expect(changedRows[1].getAttribute('initialIndex')).toBe('1');
+      expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
+    });
+  });
+
   it('should match the right table content with dynamic data', () => {
     const initialDataLength = dataSource.data.length;
     expect(dataSource.data.length).toBe(3);
@@ -406,6 +518,41 @@ class DynamicDataSourceCdkTableApp {
   @ViewChild(CdkTable) table: CdkTable<TestData>;
 }
 
+@Component({
+  template: `
+    <cdk-table [dataSource]="dataSource" [trackBy]="trackBy">
+      <ng-container cdkColumnDef="column_a">
+        <cdk-header-cell *cdkHeaderCellDef> Column A</cdk-header-cell>
+        <cdk-cell *cdkCellDef="let row"> {{row.a}}</cdk-cell>
+      </ng-container>
+
+      <ng-container cdkColumnDef="column_b">
+        <cdk-header-cell *cdkHeaderCellDef> Column B</cdk-header-cell>
+        <cdk-cell *cdkCellDef="let row"> {{row.b}}</cdk-cell>
+      </ng-container>
+
+      <cdk-header-row *cdkHeaderRowDef="columnsToRender"></cdk-header-row>
+      <cdk-row *cdkRowDef="let row; columns: columnsToRender"></cdk-row>
+    </cdk-table>
+  `
+})
+class TrackByCdkTableApp {
+  trackByStrategy: 'reference' | 'propertyA' | 'index' = 'reference';
+
+  dataSource: FakeDataSource = new FakeDataSource();
+  columnsToRender = ['column_a', 'column_b'];
+
+  @ViewChild(CdkTable) table: CdkTable<TestData>;
+
+  trackBy = (index: number, item: TestData) => {
+    switch (this.trackByStrategy) {
+      case 'reference': return item;
+      case 'propertyA': return item.a;
+      case 'index': return index;
+    }
+  }
+}
+
 @Component({
   template: `
     <cdk-table [dataSource]="dataSource" role="treegrid">

src/lib/core/data-table/data-table.ts

@@ -17,12 +17,14 @@ import {
   ElementRef,
   EmbeddedViewRef,
   Input,
+  isDevMode,
   IterableChangeRecord,
   IterableDiffer,
   IterableDiffers,
   NgIterable,
   QueryList,
   Renderer2,
+  TrackByFunction,
   ViewChild,
   ViewContainerRef,
   ViewEncapsulation
@@ -103,6 +105,24 @@ export class CdkTable<T> implements CollectionViewer {
   /** Differ used to find the changes in the data provided by the data source. */
   private _dataDiffer: IterableDiffer<T>;
 
+  /**
+   * Tracking function that will be used to check the differences in data changes. Used similarly
+   * to ngFor trackBy function. Optimize row operations by identifying a row based on its data
+   * relative to the function to know if a row should be added/removed/moved.
+   * Accepts a function that takes two parameters, `index` and `item`.
+   */
+  @Input()
+  set trackBy(fn: TrackByFunction<T>) {
+    if (isDevMode() &&
+        fn != null && typeof fn !== 'function' &&
+        <any>console && <any>console.warn) {
+        console.warn(`trackBy must be a function, but received ${JSON.stringify(fn)}.`);
+    }
+    this._trackByFn = fn;
+  }
+  get trackBy(): TrackByFunction<T> { return this._trackByFn; }
+  private _trackByFn: TrackByFunction<T>;
+
   // TODO(andrewseguin): Remove max value as the end index
   //   and instead calculate the view on init and scroll.
   /**
@@ -156,10 +176,6 @@ export class CdkTable<T> implements CollectionViewer {
     if (!role) {
       renderer.setAttribute(elementRef.nativeElement, 'role', 'grid');
     }
-
-    // TODO(andrewseguin): Add trackby function input.
-    // Find and construct an iterable differ that can be used to find the diff in an array.
-    this._dataDiffer = this._differs.find(this._data).create();
   }
 
   ngOnDestroy() {
@@ -199,8 +215,11 @@ export class CdkTable<T> implements CollectionViewer {
   }
 
   ngAfterViewInit() {
-    this._isViewInitialized = true;
+    // Find and construct an iterable differ that can be used to find the diff in an array.
+    this._dataDiffer = this._differs.find([]).create(this._trackByFn);
+
     this._renderHeaderRow();
+    this._isViewInitialized = true;
   }
 
   ngDoCheck() {