From dd84e8ff2bbb8193f585ec23d870716a0e2151c0 Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Wed, 10 Apr 2019 08:57:42 -0700
Subject: [PATCH] Newly selected components always auto-expand their ancestors

---
 src/devtools/store.js                         |  7 ++++-
 src/devtools/views/Components/TreeContext.js  | 31 ++++++++++---------
 .../views/Profiler/ProfilerContext.js         | 11 ++++---
 3 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/src/devtools/store.js b/src/devtools/store.js
index a764394184..66e84b578c 100644
--- a/src/devtools/store.js
+++ b/src/devtools/store.js
@@ -353,7 +353,12 @@ export default class Store extends EventEmitter {
           break;
         }
         const child = ((this._idToElement.get(childID): any): Element);
-        index += child.isCollapsed ? 1 : child.weight;
+
+        // We intentionally ignore collapsed state when determining an item's index.
+        // We only do this for the direct path to an element.
+        // That's because the index of the element is meaningless if it's inside of a collapsed tree.
+        // If this index is used to display the element, the caller should also un-collapse its ancestors.
+        index += child.weight;
       }
 
       previousID = current.id;
diff --git a/src/devtools/views/Components/TreeContext.js b/src/devtools/views/Components/TreeContext.js
index 26a925a717..915312d521 100644
--- a/src/devtools/views/Components/TreeContext.js
+++ b/src/devtools/views/Components/TreeContext.js
@@ -22,6 +22,7 @@ import React, {
   useCallback,
   useContext,
   useEffect,
+  useLayoutEffect,
   useMemo,
   useReducer,
   useRef,
@@ -110,6 +111,8 @@ function reduceTreeState(store: Store, state: State, action: Action): State {
     selectedElementID,
   } = state;
 
+  let lookupIDForIndex = true;
+
   // Base tree should ignore selected element changes when the owner's tree is active.
   if (ownerStack.length === 0) {
     switch (type) {
@@ -128,6 +131,11 @@ function reduceTreeState(store: Store, state: State, action: Action): State {
         selectedElementIndex = ((payload: any): number | null);
         break;
       case 'SELECT_ELEMENT_BY_ID':
+        // Skip lookup in this case; it would be redundant.
+        // It might also cause problems if the specified element was inside of a (not yet expanded) subtree.
+        lookupIDForIndex = false;
+
+        selectedElementID = payload;
         selectedElementIndex =
           payload === null
             ? null
@@ -168,7 +176,7 @@ function reduceTreeState(store: Store, state: State, action: Action): State {
   }
 
   // Keep selected item ID and index in sync.
-  if (selectedElementIndex !== state.selectedElementIndex) {
+  if (lookupIDForIndex && selectedElementIndex !== state.selectedElementIndex) {
     if (selectedElementIndex === null) {
       selectedElementID = null;
     } else {
@@ -687,19 +695,14 @@ function TreeContextController({ children, viewElementSource }: Props) {
     return () => bridge.removeListener('selectFiber', handleSelectFiber);
   }, [bridge, dispatch]);
 
-  // If a newly-selected search result is inside of a collapsed subtree, auto expand it.
-  // We also need to handle when the search text changed (selecting a new element) without changing the index.
-  const prevSearchIndex = useRef<number | null>(null);
-  const prevSearchText = useRef<string>('');
-  useEffect(() => {
-    if (
-      state.searchIndex !== prevSearchIndex.current ||
-      state.searchText !== prevSearchText.current
-    ) {
-      prevSearchIndex.current = state.searchIndex;
-      prevSearchText.current = state.searchText;
+  // If a newly-selected search result or inspection selection is inside of a collapsed subtree, auto expand it.
+  // This needs to be a layout effect to avoid temporarily flashing an incorrect selection.
+  const prevSelectedElementID = useRef<number | null>(null);
+  useLayoutEffect(() => {
+    if (state.selectedElementID !== prevSelectedElementID.current) {
+      prevSelectedElementID.current = state.selectedElementID;
 
-      if (state.searchIndex !== null && state.selectedElementID !== null) {
+      if (state.selectedElementID !== null) {
         let element = store.getElementByID(state.selectedElementID);
         while (element !== null && element.parentID > 0) {
           element = ((store.getElementByID(element.parentID): any): Element);
@@ -709,7 +712,7 @@ function TreeContextController({ children, viewElementSource }: Props) {
         }
       }
     }
-  }, [state.searchIndex, state.searchText, state.selectedElementID, store]);
+  }, [state.selectedElementID, store]);
 
   // Mutations to the underlying tree may impact this context (e.g. search results, selection state).
   useEffect(() => {
diff --git a/src/devtools/views/Profiler/ProfilerContext.js b/src/devtools/views/Profiler/ProfilerContext.js
index d06b5dab73..c8a89c0d9e 100644
--- a/src/devtools/views/Profiler/ProfilerContext.js
+++ b/src/devtools/views/Profiler/ProfilerContext.js
@@ -79,7 +79,7 @@ type Props = {|
 
 function ProfilerContextController({ children }: Props) {
   const store = useContext(StoreContext);
-  const { selectElementAtIndex, selectedElementID } = useContext(TreeContext);
+  const { selectElementByID, selectedElementID } = useContext(TreeContext);
 
   const subscription = useMemo(
     () => ({
@@ -152,13 +152,14 @@ function ProfilerContextController({ children }: Props) {
       selectFiberID(id);
       selectFiberName(name);
       if (id !== null) {
-        const index = store.getIndexOfElementID(id);
-        if (index !== null) {
-          selectElementAtIndex(index);
+        // If this element is still in the store, then select it in the Components tab as well.
+        const element = store.getElementByID(id);
+        if (element !== null) {
+          selectElementByID(id);
         }
       }
     },
-    [selectElementAtIndex, selectFiberID, selectFiberName, store]
+    [selectElementByID, selectFiberID, selectFiberName, store]
   );
 
   if (isProfiling) {