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

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

@@ -25,16 +25,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.expansionModel.select(...toBeExpanded);
+    const allNodes = this.dataNodes.reduce((accumulator, dataNode) =>
+        [...accumulator, ...this.getDescendants(dataNode), dataNode], []);
+    this.expansionModel.select(...allNodes);
   }
 
   /** 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.splice(1);
   }
 
   /** A helper function to get descendants recursively. */

src/cdk/tree/nested-node.ts

@@ -10,6 +10,8 @@ import {
   ContentChildren,
   Directive,
   ElementRef,
+  IterableDiffers,
+  IterableDiffer,
   OnDestroy,
   QueryList,
 } from '@angular/core';
@@ -51,32 +53,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 +89,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

@@ -259,22 +259,22 @@ 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> = this._dataDiffer,
+                    viewContainer: ViewContainerRef = this._nodeOutlet.viewContainer) {
+    const changes = dataDiffer.diff(data);
     if (!changes) { return; }
 
-    const 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, viewContainer);
         } else if (currentIndex == null) {
           viewContainer.remove(adjustedPreviousIndex);
         } else {

src/demo-app/tree/tree-demo.html

@@ -23,11 +23,11 @@
   <mat-expansion-panel>
     <mat-expansion-panel-header>Nested tree</mat-expansion-panel-header>
     <mat-tree [dataSource]="nestedDataSource" [treeControl]="nestedTreeControl">
-      <mat-tree-node *matTreeNodeDef="let node" matTreeNodeToggle>
-        <li>
+      <mat-nested-tree-node *matTreeNodeDef="let node" matTreeNodeToggle>
+        <li class="mat-tree-node">
           <div>{{node.filename}}:  {{node.type}}</div>
         </li>
-      </mat-tree-node>
+      </mat-nested-tree-node>
       <mat-nested-tree-node *matTreeNodeDef="let node; when: hasNestedChild">
         <li>
           <div class="mat-tree-node">

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;
   }