From d7447a8458c8196425e94ef746ba6c9d1be6f421 Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Thu, 2 May 2019 15:12:57 -0700
Subject: [PATCH] Fixed bug in Store.getIndexOfElementID() that caused roots
 with multiple top-level children to return an incorrect item index

---
 src/__tests__/store-test.js | 87 +++++++++++++++++++++++++++++++++++++
 src/devtools/store.js       | 13 +++---
 2 files changed, 94 insertions(+), 6 deletions(-)

diff --git a/src/__tests__/store-test.js b/src/__tests__/store-test.js
index 067448426c..5942574a77 100644
--- a/src/__tests__/store-test.js
+++ b/src/__tests__/store-test.js
@@ -696,4 +696,91 @@ describe('Store', () => {
       expect(store).toMatchSnapshot('4: toggle fallback on');
     });
   });
+
+  describe('getIndexOfElementID', () => {
+    beforeEach(() => {
+      store.collapseNodesByDefault = false;
+    });
+
+    it('should support a single root with a single child', () => {
+      const Grandparent = () => (
+        <React.Fragment>
+          <Parent />
+          <Parent />
+        </React.Fragment>
+      );
+      const Parent = () => <Child />;
+      const Child = () => null;
+
+      act(() =>
+        ReactDOM.render(<Grandparent />, document.createElement('div'))
+      );
+
+      for (let i = 0; i < store.numElements; i++) {
+        expect(store.getIndexOfElementID(store.getElementIDAtIndex(i))).toBe(i);
+      }
+    });
+
+    it('should support multiple roots with one children each', () => {
+      const Grandparent = () => <Parent />;
+      const Parent = () => <Child />;
+      const Child = () => null;
+
+      act(() => {
+        ReactDOM.render(<Grandparent />, document.createElement('div'));
+        ReactDOM.render(<Grandparent />, document.createElement('div'));
+      });
+
+      for (let i = 0; i < store.numElements; i++) {
+        expect(store.getIndexOfElementID(store.getElementIDAtIndex(i))).toBe(i);
+      }
+    });
+
+    it('should support a single root with multiple top level children', () => {
+      const Grandparent = () => <Parent />;
+      const Parent = () => <Child />;
+      const Child = () => null;
+
+      act(() =>
+        ReactDOM.render(
+          <React.Fragment>
+            <Grandparent />
+            <Grandparent />
+          </React.Fragment>,
+          document.createElement('div')
+        )
+      );
+
+      for (let i = 0; i < store.numElements; i++) {
+        expect(store.getIndexOfElementID(store.getElementIDAtIndex(i))).toBe(i);
+      }
+    });
+
+    it('should support multiple roots with multiple top level children', () => {
+      const Grandparent = () => <Parent />;
+      const Parent = () => <Child />;
+      const Child = () => null;
+
+      act(() => {
+        ReactDOM.render(
+          <React.Fragment>
+            <Grandparent />
+            <Grandparent />
+          </React.Fragment>,
+          document.createElement('div')
+        );
+        ReactDOM.render(
+          <React.Fragment>
+            <Grandparent />
+            <Grandparent />
+          </React.Fragment>,
+          document.createElement('div')
+        );
+      });
+
+      for (let i = 0; i < store.numElements; i++) {
+        expect(store.getIndexOfElementID(store.getElementIDAtIndex(i))).toBe(i);
+      }
+    });
+  });
 });
diff --git a/src/devtools/store.js b/src/devtools/store.js
index 6da87d3032..9c129d7162 100644
--- a/src/devtools/store.js
+++ b/src/devtools/store.js
@@ -429,12 +429,6 @@ export default class Store extends EventEmitter {
     let index = 0;
     while (true) {
       const current = ((this._idToElement.get(currentID): any): Element);
-      if (current.parentID === 0) {
-        // We found the root; stop crawling.
-        break;
-      }
-
-      index++;
 
       const { children } = current;
       for (let i = 0; i < children.length; i++) {
@@ -446,6 +440,13 @@ export default class Store extends EventEmitter {
         index += child.isCollapsed ? 1 : child.weight;
       }
 
+      if (current.parentID === 0) {
+        // We found the root; stop crawling.
+        break;
+      }
+
+      index++;
+
       previousID = current.id;
       currentID = current.parentID;
     }