fix(tree): Fix flat mat tree data source (#9556)

Author: tinayuangao

src/cdk-experimental/tree/tree.ts

@@ -25,6 +25,7 @@ import {
 import {BehaviorSubject} from 'rxjs/BehaviorSubject';
 import {takeUntil} from 'rxjs/operators/takeUntil';
 import {Subject} from 'rxjs/Subject';
+import {Subscription} from 'rxjs/Subscription';
 import {CdkTreeNodeDef, CdkTreeNode, CdkTreeNodeOutletContext} from './node';
 import {CdkTreeNodeOutlet} from './outlet';
 import {TreeControl} from './control/tree-control';
@@ -65,6 +66,9 @@ export class CdkTree<T> implements CollectionViewer, OnInit, OnDestroy {
   /** Stores the node definition that does not have a when predicate. */
   private _defaultNodeDef: CdkTreeNodeDef<T> | null;
 
+  /** Data subscription */
+  private _dataSubscription : Subscription | null;
+
   /**
    * Provides a stream containing the latest data array to render. Influenced by the tree's
    * stream of view window (what dataNodes are currently on screen).
@@ -118,14 +122,19 @@ export class CdkTree<T> implements CollectionViewer, OnInit, OnDestroy {
     if (this.dataSource) {
       this.dataSource.disconnect(this);
     }
+
+    if (this._dataSubscription) {
+      this._dataSubscription.unsubscribe();
+      this._dataSubscription = null;
+    }
   }
 
   ngAfterContentChecked() {
     const defaultNodeDefs = this._nodeDefs.filter(def => !def.when);
     if (defaultNodeDefs.length > 1) { throw getTreeMultipleDefaultNodeDefsError(); }
     this._defaultNodeDef = defaultNodeDefs[0];
 
-    if (this.dataSource) {
+    if (this.dataSource && this._nodeDefs && !this._dataSubscription) {
       this._observeRenderChanges();
     }
   }
@@ -145,6 +154,11 @@ export class CdkTree<T> implements CollectionViewer, OnInit, OnDestroy {
       this.dataSource.disconnect(this);
     }
 
+    if (this._dataSubscription) {
+      this._dataSubscription.unsubscribe();
+      this._dataSubscription = null;
+    }
+
     // Remove the all dataNodes if there is now no data source
     if (!dataSource) {
       this._nodeOutlet.viewContainer.clear();
@@ -155,7 +169,7 @@ export class CdkTree<T> implements CollectionViewer, OnInit, OnDestroy {
 
   /** Set up a subscription for the data provided by the data source. */
   private _observeRenderChanges() {
-    this.dataSource.connect(this).pipe(takeUntil(this._onDestroy))
+    this._dataSubscription = this.dataSource.connect(this).pipe(takeUntil(this._onDestroy))
       .subscribe(data => {
         this._data = data;
         this._renderNodeChanges(data);

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

@@ -136,18 +136,20 @@ export class MatTreeFlatDataSource<T, F> implements DataSource<F> {
   }
 
   connect(collectionViewer: CollectionViewer): Observable<F[]> {
-    return merge([
+    const changes = [
       collectionViewer.viewChange,
-      this.treeControl.expansionModel.onChange,
-      this._flattenedData])
-      .pipe(map(() => {
-        this._expandedData.next(
-          this.treeFlattener.expandFlattenedNodes(this._flattenedData.value, this.treeControl));
-        return this._expandedData.value;
-      }));
+      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() {
+    // no op
   }
 }