fix(tree): improve nested tree node & fix nested tree control

Author: tinayuangao

src/cdk/tree/control/nested-tree-control.ts

@@ -26,15 +26,17 @@ export class NestedTreeControl<T> extends BaseTreeControl<T> {
   expandAll(): void {
     this.expansionModel.clear();
     let toBeExpanded = <any>[];
-    this.dataNodes.forEach(dataNode => toBeExpanded.push(...this.getDescendants(dataNode)));
+    this.dataNodes.forEach(dataNode =>
+        toBeExpanded.push(...this.getDescendants(dataNode), dataNode));
     this.expansionModel.select(...toBeExpanded);
   }
 
   /** Gets a list of descendant dataNodes of a subtree rooted at given data node recursively. */
   getDescendants(dataNode: T): T[] {
     const descendants = [];
     this._getDescendants(descendants, dataNode);
-    return descendants;
+    // Remove the node itself
+    return descendants.slice(1);
   }
 
   /** A helper function to get descendants recursively. */

src/cdk/tree/nested-node.ts

@@ -10,6 +10,9 @@ import {
   ContentChildren,
   Directive,
   ElementRef,
+  IterableDiffers,
+  IterableDiffer,
+  IterableChangeRecord,
   OnDestroy,
   QueryList,
 } from '@angular/core';
@@ -51,32 +54,33 @@ import {getTreeControlFunctionsMissingError} from './tree-errors';
   providers: [{provide: CdkTreeNode, useExisting: CdkNestedTreeNode}]
 })
 export class CdkNestedTreeNode<T> extends CdkTreeNode<T> implements AfterContentInit, OnDestroy {
+  /** Differ used to find the changes in the data provided by the data source. */
+  private _dataDiffer: IterableDiffer<T>;
+
   /** The children data dataNodes of current node. They will be placed in `CdkTreeNodeOutlet`. */
   protected _children: T[];
 
   /** The children node placeholder. */
   @ContentChildren(CdkTreeNodeOutlet) nodeOutlet: QueryList<CdkTreeNodeOutlet>;
 
   constructor(protected _elementRef: ElementRef,
-              protected _tree: CdkTree<T>) {
+              protected _tree: CdkTree<T>,
+              protected _differs: IterableDiffers) {
     super(_elementRef, _tree);
   }
 
   ngAfterContentInit() {
+    this._dataDiffer = this._differs.find([]).create();
     if (!this._tree.treeControl.getChildren) {
       throw getTreeControlFunctionsMissingError();
     }
     this._tree.treeControl.getChildren(this.data).pipe(takeUntil(this._destroyed))
         .subscribe(result => {
-          if (result && result.length) {
-            // In case when nodeOutlet is not in the DOM when children changes, save it in the node
-            // and add to nodeOutlet when it's available.
-            this._children = result as T[];
-            this._addChildrenNodes();
-          }
+          this._children = result;
+          this._updateChildrenNodes();
         });
     this.nodeOutlet.changes.pipe(takeUntil(this._destroyed))
-        .subscribe((_) => this._addChildrenNodes());
+        .subscribe((_) => this._updateChildrenNodes());
   }
 
   ngOnDestroy() {
@@ -86,12 +90,10 @@ export class CdkNestedTreeNode<T> extends CdkTreeNode<T> implements AfterContent
   }
 
   /** Add children dataNodes to the NodeOutlet */
-  protected _addChildrenNodes(): void {
-    this._clear();
-    if (this.nodeOutlet.length && this._children && this._children.length) {
-      this._children.forEach((child, index) => {
-        this._tree.insertNode(child, index, this.nodeOutlet.first.viewContainer);
-      });
+  protected _updateChildrenNodes(): void {
+    if (this.nodeOutlet.length && this._children) {
+      const viewContainer = this.nodeOutlet.first.viewContainer;
+      this._tree.renderNodeChanges(this._children, this._dataDiffer, viewContainer);
     }
   }
 

src/cdk/tree/tree.ts

@@ -263,27 +263,27 @@ export class CdkTree<T> implements CollectionViewer, OnInit, OnDestroy {
 
     if (dataStream) {
       this._dataSubscription = dataStream.pipe(takeUntil(this._onDestroy))
-        .subscribe(data => this._renderNodeChanges(data));
+        .subscribe(data => this.renderNodeChanges(data));
     } else {
       throw getTreeNoValidDataSourceError();
     }
   }
 
   /** Check for changes made in the data and render each change (node added/removed/moved). */
-  private _renderNodeChanges(dataNodes: T[]) {
-    const changes = this._dataDiffer.diff(dataNodes);
+  renderNodeChanges(data: T[], dataDiffer?: IterableDiffer<T>, viewContainer?: ViewContainerRef) {
+    const changes = (dataDiffer ? dataDiffer : this._dataDiffer).diff(data);
     if (!changes) { return; }
 
-    const viewContainer = this._nodeOutlet.viewContainer;
+    const container = viewContainer ? viewContainer : this._nodeOutlet.viewContainer;
     changes.forEachOperation(
       (item: IterableChangeRecord<T>, adjustedPreviousIndex: number, currentIndex: number) => {
         if (item.previousIndex == null) {
-          this.insertNode(dataNodes[currentIndex], currentIndex);
+          this.insertNode(data[currentIndex], currentIndex, container);
         } else if (currentIndex == null) {
-          viewContainer.remove(adjustedPreviousIndex);
+          container.remove(adjustedPreviousIndex);
         } else {
-          const view = viewContainer.get(adjustedPreviousIndex);
-          viewContainer.move(view!, currentIndex);
+          const view = container.get(adjustedPreviousIndex);
+          container.move(view!, currentIndex);
         }
       });
   }

src/lib/tree/node.ts

@@ -12,6 +12,7 @@ import {
   Directive,
   ElementRef,
   Input,
+  IterableDiffers,
   QueryList
 } from '@angular/core';
 import {
@@ -94,8 +95,9 @@ export class MatNestedTreeNode<T> extends _MatNestedTreeNodeMixinBase<T>
 
   constructor(protected _elementRef: ElementRef,
               protected _tree: CdkTree<T>,
+              protected _differs: IterableDiffers,
               @Attribute('tabindex') tabIndex: string) {
-    super(_elementRef, _tree);
+    super(_elementRef, _tree, _differs);
 
     this.tabIndex = Number(tabIndex) || 0;
   }