feat(data-table): add trackby input

Author: andrewseguin

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

@@ -1,6 +1,19 @@
-<div class="demo-data-source-actions">
+<div class="demo-database-actions">
   <button md-raised-button (click)="connect()">Connect New Data Source</button>
   <button md-raised-button (click)="disconnect()" [disabled]="!dataSource">Disconnect Data Source</button>
+
+  <button md-raised-button (click)="_peopleDatabase.shuffle(changeReferences)">Randomize Data</button>
+  <md-checkbox [(ngModel)]="changeReferences">Change references</md-checkbox>
+</div>
+
+<span>Track By</span>
+<md-radio-group [(ngModel)]="trackBy">
+  <md-radio-button [value]="'id'"> ID </md-radio-button>
+  <md-radio-button [value]="'reference'"> Reference </md-radio-button>
+  <md-radio-button [value]="'index'"> Index </md-radio-button>
+</md-radio-group>
+
+<div class="demo-data-source-actions">
 </div>
 
 <div class="demo-table-container mat-elevation-z4">
@@ -9,7 +22,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

@@ -91,4 +91,3 @@
     font-family: cursive;
   }
 }
-

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

@@ -1,20 +1,28 @@
 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';
 
+export type TrackByStrategy = 'id' | 'reference' | 'index';
+
 @Component({
   moduleId: module.id,
   selector: 'data-table-demo',
   templateUrl: 'data-table-demo.html',
   styleUrls: ['data-table-demo.css'],
 })
 export class DataTableDemo {
+  static TRACK_BY: 'id' | 'reference' | 'index' = 'reference';
+
+  get trackBy(): TrackByStrategy { return DataTableDemo.TRACK_BY; }
+  set trackBy(trackBy: TrackByStrategy) { DataTableDemo.TRACK_BY = trackBy; }
+
   dataSource: PersonDataSource;
   propertiesToDisplay: UserProperties[] = [];
+  changeReferences = false;
 
-  constructor(private _peopleDatabase: PeopleDatabase) {
+  constructor(public _peopleDatabase: PeopleDatabase) {
     this.connect();
   }
 
@@ -34,6 +42,14 @@ export class DataTableDemo {
     return distanceFromMiddle / 50 + .3;
   }
 
+  userTrackBy(index: number, item: UserData) {
+    switch (DataTableDemo.TRACK_BY) {
+      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();
@@ -25,16 +28,40 @@ export class PeopleDatabase {
     for (let i = 0; i < 100; i++) { this.addPerson(); }
   }
 
+  randomize(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

@@ -21,7 +21,12 @@ describe('CdkTable', () => {
 
     TestBed.configureTestingModule({
       imports: [CdkDataTableModule],
-      declarations: [SimpleCdkTableApp, DynamicDataSourceCdkTableApp, CustomRoleCdkTableApp],
+      declarations: [
+          SimpleCdkTableApp,
+        DynamicDataSourceCdkTableApp,
+        CustomRoleCdkTableApp,
+        TrackByCdkTableApp,
+      ],
     }).compileComponents();
 
     fixture = TestBed.createComponent(SimpleCdkTableApp);
@@ -138,6 +143,119 @@ describe('CdkTable', () => {
     expect(changedRows[2].getAttribute('initialIndex')).toBe(null);
   });
 
+  describe('should properly use trackBy when diffing to add/remove/move rows', () => {
+
+    afterEach(() => {
+      // Resetting the static value of the trackby function for TrackByCdkTableApp
+      TrackByCdkTableApp.TRACK_BY = 'reference';
+    });
+
+    function createComponent() {
+      fixture = TestBed.createComponent(TrackByCdkTableApp);
+
+      component = fixture.componentInstance;
+      dataSource = component.dataSource as FakeDataSource;
+      table = component.table;
+      tableElement = fixture.nativeElement.querySelector('cdk-table');
+
+      fixture.detectChanges();  // Let the component and table create embedded views
+      fixture.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 changeData() {
+      // 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('with reference-based trackBy', () => {
+      createComponent();
+      changeData();
+
+      // 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('with changed references without property-based trackBy', () => {
+      createComponent();
+      changeData();
+
+      // Change each item reference to show that the trackby is not checking the item properties.
+      component.dataSource.data = component.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('with changed references with property-based trackBy', () => {
+      TrackByCdkTableApp.TRACK_BY = 'propertyA';
+      createComponent();
+      changeData();
+
+      // 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 => { 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('with changed references with index-based trackBy', () => {
+      TrackByCdkTableApp.TRACK_BY = 'index';
+      createComponent();
+      changeData();
+
+      // 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 => { 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);
@@ -316,6 +434,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 {
+  static TRACK_BY: 'reference' | 'propertyA' | 'index' = 'reference';
+
+  dataSource: FakeDataSource = new FakeDataSource();
+  columnsToRender = ['column_a', 'column_b'];
+
+  @ViewChild(CdkTable) table: CdkTable<TestData>;
+
+  trackBy(index: number, item: TestData) {
+    switch (TrackByCdkTableApp.TRACK_BY) {
+      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

@@ -15,13 +15,13 @@ import {
   ContentChildren,
   Directive,
   ElementRef,
-  Input,
+  Input, isDevMode,
   IterableChangeRecord,
   IterableDiffer,
   IterableDiffers,
   NgIterable,
   QueryList,
-  Renderer2,
+  Renderer2, TrackByFunction,
   ViewChild,
   ViewContainerRef,
   ViewEncapsulation
@@ -93,6 +93,22 @@ export class CdkTable<T> implements CollectionViewer {
   /** Differ used to find the changes in the data provided by the data source. */
   private _dataDiffer: IterableDiffer<T> = null;
 
+  /**
+   * Tracking function that will be used to check the differences in data changes.
+   * 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.
   /**
@@ -146,10 +162,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() {
@@ -189,6 +201,9 @@ export class CdkTable<T> implements CollectionViewer {
   }
 
   ngAfterViewInit() {
+    // 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();
 
     if (this.dataSource) {