Fixed some edge cases with collapsed by default. Still some bugs existing.

src/__tests__/__snapshots__/store-test.js.snap

@@ -179,19 +179,62 @@ exports[`Store collapseNodesByDefault:true should filter DOM nodes from the stor
     ▸ <Parent>
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 1: mount 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 1: mount 1`] = `
+[root]
+  ▸ <Wrapper>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 2: expand deepest node 1`] = `
+[root]
+  ▾ <Wrapper>
+    ▾ <Nested>
+      ▾ <Nested>
+        ▾ <Nested>
+            <Nested>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 3: collapse root 1`] = `
+[root]
+  ▸ <Wrapper>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 4: expand root 1`] = `
+[root]
+  ▾ <Wrapper>
+    ▾ <Nested>
+      ▾ <Nested>
+        ▾ <Nested>
+            <Nested>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 5: collapse middle node 1`] = `
+[root]
+  ▾ <Wrapper>
+    ▸ <Nested>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding deep parts of the tree: 6: expand middle node 1`] = `
+[root]
+  ▾ <Wrapper>
+    ▾ <Nested>
+      ▾ <Nested>
+        ▾ <Nested>
+            <Nested>
+`;
+
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 1: mount 1`] = `
 [root]
   ▸ <Grandparent>
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 2: expand Grandparent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 2: expand Grandparent 1`] = `
 [root]
   ▾ <Grandparent>
     ▸ <Parent>
     ▸ <Parent>
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 3: expand first Parent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 3: expand first Parent 1`] = `
 [root]
   ▾ <Grandparent>
     ▾ <Parent>
@@ -200,7 +243,7 @@ exports[`Store collapseNodesByDefault:true should support collapsing parts of th
     ▸ <Parent>
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 4: expand second Parent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 4: expand second Parent 1`] = `
 [root]
   ▾ <Grandparent>
     ▾ <Parent>
@@ -211,7 +254,7 @@ exports[`Store collapseNodesByDefault:true should support collapsing parts of th
         <Child key="1">
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 5: collapse first Parent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 5: collapse first Parent 1`] = `
 [root]
   ▾ <Grandparent>
     ▸ <Parent>
@@ -220,14 +263,14 @@ exports[`Store collapseNodesByDefault:true should support collapsing parts of th
         <Child key="1">
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 6: collapse second Parent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 6: collapse second Parent 1`] = `
 [root]
   ▾ <Grandparent>
     ▸ <Parent>
     ▸ <Parent>
 `;
 
-exports[`Store collapseNodesByDefault:true should support collapsing parts of the tree: 7: collapse Grandparent 1`] = `
+exports[`Store collapseNodesByDefault:true should support expanding parts of the tree: 7: collapse Grandparent 1`] = `
 [root]
   ▸ <Grandparent>
 `;
@@ -255,12 +298,14 @@ exports[`Store collapseNodesByDefault:true should support mount and update opera
 
 exports[`Store collapseNodesByDefault:true should support mount and update operations: 1: mount 1`] = `
 [root]
-  ▸ <Grandparent>
+  ▸ <Parent>
+  ▸ <Parent>
 `;
 
 exports[`Store collapseNodesByDefault:true should support mount and update operations: 2: update 1`] = `
 [root]
-  ▸ <Grandparent>
+  ▸ <Parent>
+  ▸ <Parent>
 `;
 
 exports[`Store collapseNodesByDefault:true should support mount and update operations: 3: unmount 1`] = ``;

src/__tests__/setupTests.js

@@ -31,8 +31,11 @@ env.beforeEach(() => {
   const agent = new Agent();
   agent.addBridge(bridge);
 
-  initBackend(global.__REACT_DEVTOOLS_GLOBAL_HOOK__, agent, global);
+  const hook = global.__REACT_DEVTOOLS_GLOBAL_HOOK__;
 
+  initBackend(hook, agent, global);
+
+  global.agent = agent;
   global.bridge = bridge;
   global.store = new Store(bridge);
 });

src/__tests__/store-test.js

@@ -195,22 +195,32 @@ describe('Store', () => {
     });
 
     it('should support mount and update operations', () => {
-      const Grandparent = ({ count }) => (
-        <React.Fragment>
-          <Parent count={count} />
-          <Parent count={count} />
-        </React.Fragment>
-      );
       const Parent = ({ count }) =>
         new Array(count).fill(true).map((_, index) => <Child key={index} />);
       const Child = () => <div>Hi!</div>;
 
       const container = document.createElement('div');
 
-      act(() => ReactDOM.render(<Grandparent count={4} />, container));
+      act(() =>
+        ReactDOM.render(
+          <React.Fragment>
+            <Parent count={1} />
+            <Parent count={3} />
+          </React.Fragment>,
+          container
+        )
+      );
       expect(store).toMatchSnapshot('1: mount');
 
-      act(() => ReactDOM.render(<Grandparent count={2} />, container));
+      act(() =>
+        ReactDOM.render(
+          <React.Fragment>
+            <Parent count={2} />
+            <Parent count={1} />
+          </React.Fragment>,
+          container
+        )
+      );
       expect(store).toMatchSnapshot('2: update');
 
       act(() => ReactDOM.unmountComponentAtNode(container));
@@ -311,7 +321,7 @@ describe('Store', () => {
       expect(store).toMatchSnapshot('2: resolved');
     });
 
-    it('should support collapsing parts of the tree', () => {
+    it('should support expanding parts of the tree', () => {
       const Grandparent = ({ count }) => (
         <React.Fragment>
           <Parent count={count} />
@@ -353,5 +363,48 @@ describe('Store', () => {
       act(() => store.toggleIsCollapsed(grandparentID, true));
       expect(store).toMatchSnapshot('7: collapse Grandparent');
     });
+
+    it('should support expanding deep parts of the tree', () => {
+      const Wrapper = ({ forwardedRef }) => (
+        <Nested depth={3} forwardedRef={forwardedRef} />
+      );
+      const Nested = ({ depth, forwardedRef }) =>
+        depth > 0 ? (
+          <Nested depth={depth - 1} forwardedRef={forwardedRef} />
+        ) : (
+          <div ref={forwardedRef} />
+        );
+
+      const ref = React.createRef();
+
+      act(() =>
+        ReactDOM.render(
+          <Wrapper forwardedRef={ref} />,
+          document.createElement('div')
+        )
+      );
+      expect(store).toMatchSnapshot('1: mount');
+
+      const deepestedNodeID = global.agent.getIDForNode(ref.current);
+
+      act(() => store.toggleIsCollapsed(deepestedNodeID, false));
+      expect(store).toMatchSnapshot('2: expand deepest node');
+
+      const rootID = store.getElementIDAtIndex(0);
+
+      act(() => store.toggleIsCollapsed(rootID, true));
+      expect(store).toMatchSnapshot('3: collapse root');
+
+      act(() => store.toggleIsCollapsed(rootID, false));
+      expect(store).toMatchSnapshot('4: expand root');
+
+      const id = store.getElementIDAtIndex(1);
+
+      act(() => store.toggleIsCollapsed(id, true));
+      expect(store).toMatchSnapshot('5: collapse middle node');
+
+      act(() => store.toggleIsCollapsed(id, false));
+      expect(store).toMatchSnapshot('6: expand middle node');
+    });
   });
 });

src/devtools/store.js

@@ -315,9 +315,8 @@ export default class Store extends EventEmitter {
 
     // Find the element in the tree using the weight of each node...
     // Skip over the root itself, because roots aren't visible in the Elements tree.
-    const firstChildID = ((root: any): Element).children[0];
-    let currentElement = ((this._idToElement.get(firstChildID): any): Element);
-    let currentWeight = rootWeight;
+    let currentElement = ((root: any): Element);
+    let currentWeight = rootWeight - 1;
     while (index !== currentWeight) {
       const numChildren = currentElement.children.length;
       for (let i = 0; i < numChildren; i++) {
@@ -449,25 +448,66 @@ export default class Store extends EventEmitter {
     // We do this to avoid mismatches on e.g. CommitTreeBuilder that would cause errors.
   }
 
+  // TODO Maybe split this into two methods: expand() and collapse()
   toggleIsCollapsed(id: number, isCollapsed: boolean): void {
     const element = this.getElementByID(id);
     if (element !== null) {
-      if (element.type === ElementTypeRoot) {
-        throw Error('Root nodes cannot be collapsed');
+      if (isCollapsed) {
+        if (element.type === ElementTypeRoot) {
+          throw Error('Root nodes cannot be collapsed');
+        }
+
+        if (element.isCollapsed) {
+          return;
+        }
+
+        element.isCollapsed = true;
+
+        const weightDelta = 1 - element.weight;
+
+        let parentElement = ((this._idToElement.get(
+          element.parentID
+        ): any): Element);
+        while (parentElement != null) {
+          parentElement.weight += weightDelta;
+          parentElement = this._idToElement.get(parentElement.parentID);
+        }
+      } else {
+        let currentElement = element;
+        while (currentElement != null) {
+          const oldWeight = currentElement.isCollapsed
+            ? 1
+            : currentElement.weight;
+          currentElement.isCollapsed = false;
+          const newWeight = currentElement.isCollapsed
+            ? 1
+            : currentElement.weight;
+          const weightDelta = newWeight - oldWeight;
+
+          let parentElement = ((this._idToElement.get(
+            currentElement.parentID
+          ): any): Element);
+          while (parentElement != null) {
+            parentElement.weight += weightDelta;
+            if (parentElement.isCollapsed) {
+              break;
+            }
+            parentElement = this._idToElement.get(parentElement.parentID);
+          }
+
+          currentElement =
+            currentElement.parentID !== 0
+              ? this.getElementByID(currentElement.parentID)
+              : null;
+        }
       }
 
-      const oldWeight = element.isCollapsed ? 1 : element.weight;
-      element.isCollapsed = isCollapsed;
-      const newWeight = element.isCollapsed ? 1 : element.weight;
-      const weightDelta = newWeight - oldWeight;
-
-      this._weightAcrossRoots += weightDelta;
-
-      let parentElement = this._idToElement.get(element.parentID);
-      while (parentElement != null) {
-        parentElement.weight += weightDelta;
-        parentElement = this._idToElement.get(parentElement.parentID);
-      }
+      let weightAcrossRoots = 0;
+      this._roots.forEach(rootID => {
+        const { weight } = ((this.getElementByID(rootID): any): Element);
+        weightAcrossRoots += weight;
+      });
+      this._weightAcrossRoots = weightAcrossRoots;
 
       // The Tree context's search reducer expects an explicit list of ids for nodes that were added or removed.
       // In this  case, we can pass it empty arrays since nodes in a collapsed tree are still there (just hidden).
@@ -867,19 +907,23 @@ export default class Store extends EventEmitter {
 
   // Used for Jest snapshot testing.
   // May also be useful for visually debugging the tree, so it lives on the Store.
-  __toSnapshot = () => {
+  __toSnapshot = (includeWeight: boolean = false) => {
     const snapshotLines = [];
 
     let rootWeight = 0;
 
     this._roots.forEach(rootID => {
-      snapshotLines.push('[root]');
-
       const { weight } = ((this.getElementByID(rootID): any): Element);
 
+      snapshotLines.push('[root]' + (includeWeight ? ` (${weight})` : ''));
+
       for (let i = rootWeight; i < rootWeight + weight; i++) {
         const element = ((this.getElementAtIndex(i): any): Element);
 
+        if (element == null) {
+          throw Error(`No element for index ${i}`);
+        }
+
         let prefix = ' ';
         if (element.children.length > 0) {
           prefix = element.isCollapsed ? '▸' : '▾';
@@ -890,15 +934,29 @@ export default class Store extends EventEmitter {
           key = ` key="${element.key}"`;
         }
 
+        let suffix = '';
+        if (includeWeight) {
+          suffix = ` (${element.isCollapsed ? 1 : element.weight})`;
+        }
+
         snapshotLines.push(
           `${'  '.repeat(element.depth + 1)}${prefix} <${element.displayName ||
-            'null'}${key}>`
+            'null'}${key}>${suffix}`
         );
       }
 
       rootWeight += weight;
     });
 
+    // Make sure the pretty-printed test align with the Store's reported number of total rows.
+    if (rootWeight !== this._weightAcrossRoots) {
+      throw Error(
+        `Inconsistent store state. Individual root weights (${rootWeight}) do not match total weight (${
+          this._weightAcrossRoots
+        })`
+      );
+    }
+
     return snapshotLines.join('\n');
   };
 }

src/devtools/views/Components/TreeContext.js

@@ -748,11 +748,8 @@ function TreeContextController({ children, viewElementSource }: Props) {
 
       if (state.selectedElementID !== null) {
         let element = store.getElementByID(state.selectedElementID);
-        while (element !== null && element.parentID > 0) {
-          element = ((store.getElementByID(element.parentID): any): Element);
-          if (element.isCollapsed) {
-            store.toggleIsCollapsed(element.id, false);
-          }
+        if (element !== null && element.parentID > 0) {
+          store.toggleIsCollapsed(element.parentID, false);
         }
       }
     }