refactor(cdk/tree): address more PR comments; factor out helper methods

Author: BobobUnicorn

src/cdk/a11y/key-manager/tree-key-manager.ts

@@ -24,9 +24,16 @@ import {Typeahead} from './typeahead';
  * keyboard events occur.
  */
 export class TreeKeyManager<T extends TreeKeyManagerItem> implements TreeKeyManagerStrategy<T> {
+  /** The index of the currently active (focused) item. */
   private _activeItemIndex = -1;
+  /** The currently active (focused) item. */
   private _activeItem: T | null = null;
+  /** Whether or not we activate the item when it's focused. */
   private _shouldActivationFollowFocus = false;
+  /**
+   * The orientation that the tree is laid out in. In `rtl` mode, the behavior of Left and
+   * Right arrow are switched.
+   */
   private _horizontalOrientation: 'ltr' | 'rtl' = 'ltr';
 
   // Keep tree items focusable when disabled. Align with
@@ -40,6 +47,7 @@ export class TreeKeyManager<T extends TreeKeyManagerItem> implements TreeKeyMana
   /** Function to determine equivalent items. */
   private _trackByFn: (item: T) => unknown = (item: T) => item;
 
+  /** Synchronous cache of the items to manage. */
   private _items: T[] = [];
 
   private _typeahead?: Typeahead<T>;

src/cdk/tree/tree.ts

@@ -75,12 +75,12 @@ type RenderingData<T> =
   | {
       flattenedNodes: null;
       nodeType: null;
-      renderNodes: T[];
+      renderNodes: readonly T[];
     }
   | {
-      flattenedNodes: T[];
+      flattenedNodes: readonly T[];
       nodeType: 'nested' | 'flat';
-      renderNodes: [];
+      renderNodes: readonly T[];
     };
 
 /**
@@ -342,6 +342,14 @@ export class CdkTree<T, K = T>
     }
   }
 
+  private _getExpansionModel() {
+    if (!this.treeControl) {
+      this._expansionModel ??= new SelectionModel<K>(true);
+      return this._expansionModel;
+    }
+    return this.treeControl.expansionModel;
+  }
+
   /** Set up a subscription for the data provided by the data source. */
   private _subscribeToDataChanges() {
     if (this._dataSubscription) {
@@ -365,15 +373,17 @@ export class CdkTree<T, K = T>
       return;
     }
 
-    let expansionModel;
-    if (!this.treeControl) {
-      this._expansionModel = new SelectionModel<K>(true);
-      expansionModel = this._expansionModel;
-    } else {
-      expansionModel = this.treeControl.expansionModel;
-    }
+    this._dataSubscription = this._getRenderData(dataStream)
+      .pipe(takeUntil(this._onDestroy))
+      .subscribe(renderingData => {
+        this._renderDataChanges(renderingData);
+      });
+  }
 
-    this._dataSubscription = combineLatest([
+  /** Given an Observable containing a stream of the raw data, returns an Observable containing the RenderingData */
+  private _getRenderData(dataStream: Observable<readonly T[]>): Observable<RenderingData<T>> {
+    const expansionModel = this._getExpansionModel();
+    return combineLatest([
       dataStream,
       this._nodeType,
       // We don't use the expansion data directly, however we add it here to essentially
@@ -384,24 +394,19 @@ export class CdkTree<T, K = T>
           this._emitExpansionChanges(expansionChanges);
         }),
       ),
-    ])
-      .pipe(
-        switchMap(([data, nodeType]) => {
-          if (nodeType === null) {
-            return observableOf([{renderNodes: data}, nodeType] as const);
-          }
+    ]).pipe(
+      switchMap(([data, nodeType]) => {
+        if (nodeType === null) {
+          return observableOf({renderNodes: data, flattenedNodes: null, nodeType} as const);
+        }
 
-          // If we're here, then we know what our node type is, and therefore can
-          // perform our usual rendering pipeline, which necessitates converting the data
-          return this._convertData(data, nodeType).pipe(
-            map(convertedData => [convertedData, nodeType] as const),
-          );
-        }),
-        takeUntil(this._onDestroy),
-      )
-      .subscribe(([data, nodeType]) => {
-        this._renderDataChanges({nodeType, ...data} as RenderingData<T>);
-      });
+        // If we're here, then we know what our node type is, and therefore can
+        // perform our usual rendering pipeline, which necessitates converting the data
+        return this._computeRenderingData(data, nodeType).pipe(
+          map(convertedData => ({...convertedData, nodeType}) as const),
+        );
+      }),
+    );
   }
 
   private _renderDataChanges(data: RenderingData<T>) {
@@ -586,10 +591,9 @@ export class CdkTree<T, K = T>
 
   /** Whether the data node is expanded or collapsed. Returns true if it's expanded. */
   isExpanded(dataNode: T): boolean {
-    return (
-      this.treeControl?.isExpanded(dataNode) ??
-      this._expansionModel?.isSelected(this._getExpansionKey(dataNode)) ??
-      false
+    return !!(
+      this.treeControl?.isExpanded(dataNode) ||
+      this._expansionModel?.isSelected(this._getExpansionKey(dataNode))
     );
   }
 
@@ -733,26 +737,13 @@ export class CdkTree<T, K = T>
           if (!expanded) {
             return [];
           }
-          const startIndex = flattenedNodes.findIndex(node => this._getExpansionKey(node) === key);
-          const level = levelAccessor(dataNode) + 1;
-          const results: T[] = [];
-
-          // Goes through flattened tree nodes in the `flattenedNodes` array, and get all direct
-          // descendants. The level of descendants of a tree node must be equal to the level of the
-          // given tree node + 1.
-          // If we reach a node whose level is equal to the level of the tree node, we hit a sibling.
-          // If we reach a node whose level is greater than the level of the tree node, we hit a
-          // sibling of an ancestor.
-          for (let i = startIndex + 1; i < flattenedNodes.length; i++) {
-            const currentLevel = levelAccessor(flattenedNodes[i]);
-            if (level > currentLevel) {
-              break;
-            }
-            if (level === currentLevel) {
-              results.push(flattenedNodes[i]);
-            }
-          }
-          return results;
+          return this._findChildrenByLevel(
+            levelAccessor,
+            flattenedNodes,
+
+            dataNode,
+            1,
+          );
         }),
       );
     }
@@ -763,6 +754,42 @@ export class CdkTree<T, K = T>
     throw getTreeControlMissingError();
   }
 
+  /**
+   * Given the list of flattened nodes, the level accessor, and the level range within
+   * which to consider children, finds the children for a given node.
+   *
+   * For example, for direct children, `levelDelta` would be 1. For all descendants,
+   * `levelDelta` would be Infinity.
+   */
+  private _findChildrenByLevel(
+    levelAccessor: (node: T) => number,
+    flattenedNodes: readonly T[],
+    dataNode: T,
+    levelDelta: number,
+  ): T[] {
+    const key = this._getExpansionKey(dataNode);
+    const startIndex = flattenedNodes.findIndex(node => this._getExpansionKey(node) === key);
+    const dataNodeLevel = levelAccessor(dataNode);
+    const expectedLevel = dataNodeLevel + levelDelta;
+    const results: T[] = [];
+
+    // Goes through flattened tree nodes in the `flattenedNodes` array, and get all
+    // descendants within a certain level range.
+    //
+    // If we reach a node whose level is equal to or less than the level of the tree node,
+    // we hit a sibling or parent's sibling, and should stop.
+    for (let i = startIndex + 1; i < flattenedNodes.length; i++) {
+      const currentLevel = levelAccessor(flattenedNodes[i]);
+      if (currentLevel <= dataNodeLevel) {
+        break;
+      }
+      if (currentLevel <= expectedLevel) {
+        results.push(flattenedNodes[i]);
+      }
+    }
+    return results;
+  }
+
   /**
    * Adds the specified node component to the tree's internal registry.
    *
@@ -842,27 +869,12 @@ export class CdkTree<T, K = T>
       return observableOf(this.treeControl.getDescendants(dataNode));
     }
     if (this.levelAccessor) {
-      const key = this._getExpansionKey(dataNode);
-      const startIndex = this._flattenedNodes.value.findIndex(
-        node => this._getExpansionKey(node) === key,
+      const results = this._findChildrenByLevel(
+        this.levelAccessor,
+        this._flattenedNodes.value,
+        dataNode,
+        Infinity,
       );
-      const results: T[] = [];
-
-      // Goes through flattened tree nodes in the `dataNodes` array, and get all descendants.
-      // The level of descendants of a tree node must be greater than the level of the given
-      // tree node.
-      // If we reach a node whose level is equal to the level of the tree node, we hit a sibling.
-      // If we reach a node whose level is greater than the level of the tree node, we hit a
-      // sibling of an ancestor.
-      const currentLevel = this.levelAccessor(dataNode);
-      for (
-        let i = startIndex + 1;
-        i < this._flattenedNodes.value.length &&
-        currentLevel < this.levelAccessor(this._flattenedNodes.value[i]);
-        i++
-      ) {
-        results.push(this._flattenedNodes.value[i]);
-      }
       return observableOf(results);
     }
     if (this.childrenAccessor) {
@@ -1003,7 +1015,7 @@ export class CdkTree<T, K = T>
    *
    * This also computes parent, level, and group data.
    */
-  private _convertData(
+  private _computeRenderingData(
     nodes: readonly T[],
     nodeType: 'flat' | 'nested',
   ): Observable<{