feat: 🎸 added getNumberOfVisibleDescendants to TreeState

index.d.ts

@@ -144,6 +144,7 @@ export interface TreeState {
   getNodeAt: (state: State, index: number) => Node;
   getTree: (state: State) => Node[];
   createFromTree: (tree: Node[]) => State;
+  getNumberOfVisibleDescendants: (state: State, index: number) => number;
 }
 
 export const selectors: Selectors;

src/state/TreeState.js

@@ -1,7 +1,7 @@
 import {getFlattenedTreePaths} from '../selectors/getFlattenedTree';
 import {getNodeFromPath} from '../selectors/nodes';
 
-class State {
+export class State {
   flattenedTree = null;
   tree = null;
 
@@ -41,6 +41,29 @@ export default class TreeState {
     return getNodeFromPath(rowPath, state.tree);
   };
 
+  /**
+   * Given a state and an index, finds the number of visible descendants
+   * @param {State} state - The current state
+   * @param {number} index - The visible row index
+   * @return {number} The number of visible descendants
+   */
+  static getNumberOfVisibleDescendants = (state, index) => {
+    const {id} = TreeState.getNodeAt(state, index);
+
+    const {flattenedTree} = state;
+    let i;
+
+    for (i = index; i < flattenedTree.length; i++) {
+      const path = flattenedTree[i];
+
+      if (!path.some(p => p === id)) {
+        break;
+      }
+    }
+
+    return Math.max(i - 1 - index, 0);
+  };
+
   /**
    * Given a state, gets the tree
    * @param {State} state - The current state

src/state/__tests__/TreeState.test.js

@@ -61,4 +61,47 @@ describe('TreeState', () => {
       expect(() => TreeState.getTree(() => {})).toThrowError('Expected a State instance but got function');
     });
   });
+
+  describe('getNumberOfVisibleDescendants', () => {
+    test('should fail for when invalid state is supplied', () => {
+      expect(() => TreeState.getNumberOfVisibleDescendants('state', 0)).toThrowError(
+        'Expected a State instance but got string',
+      );
+      expect(() => TreeState.getNumberOfVisibleDescendants(1225, 0)).toThrowError(
+        'Expected a State instance but got number',
+      );
+      expect(() => TreeState.getNumberOfVisibleDescendants([], 0)).toThrowError(
+        'Expected a State instance but got object',
+      );
+      expect(() => TreeState.getNumberOfVisibleDescendants({}, 0)).toThrowError(
+        'Expected a State instance but got object',
+      );
+      expect(() => TreeState.getNumberOfVisibleDescendants(true, 0)).toThrowError(
+        'Expected a State instance but got boolean',
+      );
+      expect(() => TreeState.getNumberOfVisibleDescendants(() => {}, 0)).toThrowError(
+        'Expected a State instance but got function',
+      );
+    });
+
+    test('should get a correct number of descendants for a node with deep descendants', () => {
+      expect(TreeState.getNumberOfVisibleDescendants(TreeState.createFromTree(Nodes), 0)).toEqual(4);
+    });
+
+    test('should get a correct number of descendants for a node without grand children', () => {
+      expect(TreeState.getNumberOfVisibleDescendants(TreeState.createFromTree(Nodes), 1)).toEqual(2);
+    });
+
+    test('should get a correct number of descendants for a node with deep descendants', () => {
+      expect(TreeState.getNumberOfVisibleDescendants(TreeState.createFromTree(Nodes), 0)).toEqual(4);
+    });
+
+    test('should get 0 descendants for a node that does not have any descendants in the root node', () => {
+      expect(TreeState.getNumberOfVisibleDescendants(TreeState.createFromTree(Nodes), 6)).toEqual(0);
+    });
+
+    test('should get 0 descendants for a node that does not have any descendants', () => {
+      expect(TreeState.getNumberOfVisibleDescendants(TreeState.createFromTree(Nodes), 3)).toEqual(0);
+    });
+  });
 });