license import { CollectionViewer , DataSource } from '@angular/cdk/collections' ; import { FlatTreeControl , TreeControl } from '@angular/cdk/tree' ; import { BehaviorSubject , merge , Observable } from 'rxjs' ; import { map , take } from 'rxjs/operators' ; /**
 * Tree flattener to convert a normal type of node to node with children & level information.
 * Transform nested nodes of type `T` to flattened nodes of type `F`.
 *
 * For example, the input data of type `T` is nested, and contains its children data:
 *   SomeNode: {
 *     key: 'Fruits',
 *     children: [
 *       NodeOne: {
 *         key: 'Apple',
 *       },
 *       NodeTwo: {
 *        key: 'Pear',
 *      }
 *    ]
 *  }
 *  After flattener flatten the tree, the structure will become
 *  SomeNode: {
 *    key: 'Fruits',
 *    expandable: true,
 *    level: 1
 *  },
 *  NodeOne: {
 *    key: 'Apple',
 *    expandable: false,
 *    level: 2
 *  },
 *  NodeTwo: {
 *   key: 'Pear',
 *   expandable: false,
 *   level: 2
 * }
 * and the output flattened type is `F` with additional information.
 */ export class MatTreeFlattener < T , F > { constructor ( public transformFunction , public getLevel , public isExpandable , public getChildren ) { } _flattenNode ( node , level , resultNodes , parentMap ) { const flatNode = this . transformFunction ( node , level ) ; resultNodes . push ( flatNode ) ; if ( this . isExpandable ( flatNode ) ) { const childrenNodes = this . getChildren ( node ) ; if ( childrenNodes ) { if ( Array . isArray ( childrenNodes ) ) { this . _flattenChildren ( childrenNodes , level , resultNodes , parentMap ) ; } else { childrenNodes . pipe ( take ( 1 ) ) . subscribe ( children => { this . _flattenChildren ( children , level , resultNodes , parentMap ) ; } ) ; } } } return resultNodes ; } _flattenChildren ( children , level , resultNodes , parentMap ) { children . forEach ( ( child , index ) => { let childParentMap = parentMap . slice ( ) ; childParentMap . push ( index != children . length - 1 ) ; this . _flattenNode ( child , level + 1 , resultNodes , childParentMap ) ; } ) ; } /**
     * Flatten a list of node type T to flattened version of node F.
     * Please note that type T may be nested, and the length of `structuredData` may be different
     * from that of returned list `F[]`.
     */ flattenNodes ( structuredData ) { let resultNodes = [ ] ; structuredData . forEach ( node => this . _flattenNode ( node , 0 , resultNodes , [ ] ) ) ; return resultNodes ; } /**
     * Expand flattened node with current expansion status.
     * The returned list may have different length.
     */ expandFlattenedNodes ( nodes , treeControl ) { let results = [ ] ; let currentExpand = [ ] ; currentExpand [ 0 ] = true ; nodes . forEach ( node => { let expand = true ; for ( let i = 0 ; i <= this . getLevel ( node ) ; i ++ ) { expand = expand && currentExpand [ i ] ; } if ( expand ) { results . push ( node ) ; } if ( this . isExpandable ( node ) ) { currentExpand [ this . getLevel ( node ) + 1 ] = treeControl . isExpanded ( node ) ; } } ) ; return results ; } } /**
 * Data source for flat tree.
 * The data source need to handle expansion/collapsion of the tree node and change the data feed
 * to `MatTree`.
 * 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 > { _flattenedData = new BehaviorSubject < F [ ] > ( [ ] ) ; _expandedData = new BehaviorSubject < F [ ] > ( [ ] ) ; _data : BehaviorSubject < T [ ] > ; get data ( ) { return this . _data . value ; } set data ( value ) { this . _data . next ( value ) ; this . _flattenedData . next ( this . treeFlattener . flattenNodes ( this . data ) ) ; this . treeControl . dataNodes = this . _flattenedData . value ; } constructor ( private treeControl , private treeFlattener , initialData = [ ] ) { super ( ) ; this . _data = new BehaviorSubject < T [ ] > ( initialData ) ; } connect ( collectionViewer ) { const changes = [ collectionViewer . viewChange , this . treeControl . expansionModel . onChange , this . _flattenedData ] ; return merge ( ... changes ) . pipe ( map ( ( ) => { this . _expandedData . next ( this . treeFlattener . expandFlattenedNodes ( this . _flattenedData . value , this . treeControl ) ) ; return this . _expandedData . value ; } ) ) ; } disconnect ( ) { } }