fix(tree): allow use of FlatTreeControl's trackBy function in MatTreeFlatDataSource and MatTreeFlattener (#19600)

* Allow use of FlatTreeControl's trackBy function in MatTreeFlatDataSource and MatTreeFlattener

* Update public API goldens

Author: tjaneczko

src/material/tree/data-source/flat-data-source.ts

@@ -45,7 +45,7 @@ import {map, take} from 'rxjs/operators';
  * }
  * and the output flattened type is `F` with additional information.
  */
-export class MatTreeFlattener<T, F> {
+export class MatTreeFlattener<T, F, K = F> {
 
   constructor(public transformFunction: (node: T, level: number) => F,
               public getLevel: (node: F) => number,
@@ -97,7 +97,7 @@ export class MatTreeFlattener<T, F> {
    * Expand flattened node with current expansion status.
    * The returned list may have different length.
    */
-  expandFlattenedNodes(nodes: F[], treeControl: TreeControl<F>): F[] {
+  expandFlattenedNodes(nodes: F[], treeControl: TreeControl<F, K>): F[] {
     let results: F[] = [];
     let currentExpand: boolean[] = [];
     currentExpand[0] = true;
@@ -126,7 +126,7 @@ export class MatTreeFlattener<T, F> {
  * The nested tree nodes of type `T` are flattened through `MatTreeFlattener`, and converted
  * to type `F` for `MatTree` to consume.
  */
-export class MatTreeFlatDataSource<T, F> extends DataSource<F> {
+export class MatTreeFlatDataSource<T, F, K = F> extends DataSource<F> {
   _flattenedData = new BehaviorSubject<F[]>([]);
 
   _expandedData = new BehaviorSubject<F[]>([]);
@@ -139,8 +139,8 @@ export class MatTreeFlatDataSource<T, F> extends DataSource<F> {
     this._treeControl.dataNodes = this._flattenedData.value;
   }
 
-  constructor(private _treeControl: FlatTreeControl<F>,
-              private _treeFlattener: MatTreeFlattener<T, F>,
+  constructor(private _treeControl: FlatTreeControl<F, K>,
+              private _treeFlattener: MatTreeFlattener<T, F, K>,
               initialData: T[] = []) {
     super();
     this._data = new BehaviorSubject<T[]>(initialData);

tools/public_api_guard/material/tree.d.ts

@@ -21,26 +21,26 @@ export declare class MatTree<T> extends CdkTree<T> {
     static ɵfac: i0.ɵɵFactoryDef<MatTree<any>, never>;
 }
 
-export declare class MatTreeFlatDataSource<T, F> extends DataSource<F> {
+export declare class MatTreeFlatDataSource<T, F, K = F> extends DataSource<F> {
     _data: BehaviorSubject<T[]>;
     _expandedData: BehaviorSubject<F[]>;
     _flattenedData: BehaviorSubject<F[]>;
     get data(): T[];
     set data(value: T[]);
-    constructor(_treeControl: FlatTreeControl<F>, _treeFlattener: MatTreeFlattener<T, F>, initialData?: T[]);
+    constructor(_treeControl: FlatTreeControl<F, K>, _treeFlattener: MatTreeFlattener<T, F, K>, initialData?: T[]);
     connect(collectionViewer: CollectionViewer): Observable<F[]>;
     disconnect(): void;
 }
 
-export declare class MatTreeFlattener<T, F> {
+export declare class MatTreeFlattener<T, F, K = F> {
     getChildren: (node: T) => Observable<T[]> | T[] | undefined | null;
     getLevel: (node: F) => number;
     isExpandable: (node: F) => boolean;
     transformFunction: (node: T, level: number) => F;
     constructor(transformFunction: (node: T, level: number) => F, getLevel: (node: F) => number, isExpandable: (node: F) => boolean, getChildren: (node: T) => Observable<T[]> | T[] | undefined | null);
     _flattenChildren(children: T[], level: number, resultNodes: F[], parentMap: boolean[]): void;
     _flattenNode(node: T, level: number, resultNodes: F[], parentMap: boolean[]): F[];
-    expandFlattenedNodes(nodes: F[], treeControl: TreeControl<F>): F[];
+    expandFlattenedNodes(nodes: F[], treeControl: TreeControl<F, K>): F[];
     flattenNodes(structuredData: T[]): F[];
 }