Merge pull request #215 from bvaughn/persist-selection

Try to restore selection between reloads

src/backend/agent.js

@@ -3,10 +3,19 @@
 import EventEmitter from 'events';
 import memoize from 'memoize-one';
 import throttle from 'lodash.throttle';
-import { LOCAL_STORAGE_RELOAD_AND_PROFILE_KEY, __DEBUG__ } from '../constants';
+import {
+  LOCAL_STORAGE_RELOAD_AND_PROFILE_KEY,
+  SESSION_STORAGE_LAST_SELECTION_KEY,
+  __DEBUG__,
+} from '../constants';
 import { hideOverlay, showOverlay } from './views/Highlighter';
 
-import type { RendererID, RendererInterface } from './types';
+import type {
+  PathFrame,
+  PathMatch,
+  RendererID,
+  RendererInterface,
+} from './types';
 import type { Bridge } from '../types';
 
 const debug = (methodName, ...args) => {
@@ -46,10 +55,17 @@ type OverrideSuspenseParams = {|
   forceFallback: boolean,
 |};
 
+type PersistedSelection = {|
+  rendererID: number,
+  path: Array<PathFrame>,
+|};
+
 export default class Agent extends EventEmitter {
   _bridge: Bridge = ((null: any): Bridge);
   _isProfiling: boolean = false;
   _rendererInterfaces: { [key: RendererID]: RendererInterface } = {};
+  _persistedSelection: PersistedSelection | null = null;
+  _persistedSelectionMatch: PathMatch | null = null;
 
   constructor() {
     super();
@@ -59,6 +75,15 @@ export default class Agent extends EventEmitter {
 
       localStorage.removeItem(LOCAL_STORAGE_RELOAD_AND_PROFILE_KEY);
     }
+
+    if (typeof sessionStorage !== 'undefined') {
+      const persistedSelectionString = sessionStorage.getItem(
+        SESSION_STORAGE_LAST_SELECTION_KEY
+      );
+      if (persistedSelectionString != null) {
+        this._persistedSelection = JSON.parse(persistedSelectionString);
+      }
+    }
   }
 
   addBridge(bridge: Bridge) {
@@ -316,6 +341,19 @@ export default class Agent extends EventEmitter {
     } else {
       renderer.selectElement(id);
       this._bridge.send('selectElement');
+
+      // When user selects an element, stop trying to restore the selection,
+      // and instead remember the current selection for the next reload.
+      if (
+        this._persistedSelectionMatch === null ||
+        this._persistedSelectionMatch.id !== id
+      ) {
+        this._persistedSelection = null;
+        this._persistedSelectionMatch = null;
+        renderer.setTrackedPath(null);
+        this._throttledPersistSelection(rendererID, id);
+      }
+
       // TODO: If there was a way to change the selected DOM element
       // in native Elements tab without forcing a switch to it, we'd do it here.
       // For now, it doesn't seem like there is a way to do that:
@@ -388,6 +426,14 @@ export default class Agent extends EventEmitter {
     if (this._isProfiling) {
       rendererInterface.startProfiling();
     }
+
+    // When the renderer is attached, we need to tell it whether
+    // we remember the previous selection that we'd like to restore.
+    // It'll start tracking mounts for matches to the last selection path.
+    const selection = this._persistedSelection;
+    if (selection !== null && selection.rendererID === rendererID) {
+      rendererInterface.setTrackedPath(selection.path);
+    }
   }
 
   syncSelectionFromNativeElementsPanel = () => {
@@ -480,6 +526,32 @@ export default class Agent extends EventEmitter {
     //
     // this._bridge.send('operations', operations, [operations.buffer]);
     this._bridge.send('operations', operations);
+
+    if (this._persistedSelection !== null) {
+      const rendererID = operations[0];
+      if (this._persistedSelection.rendererID === rendererID) {
+        // Check if we can select a deeper match for the persisted selection.
+        const renderer = this._rendererInterfaces[rendererID];
+        const prevMatch = this._persistedSelectionMatch;
+        const nextMatch = renderer.getBestMatchForTrackedPath();
+        this._persistedSelectionMatch = nextMatch;
+        const prevMatchID = prevMatch !== null ? prevMatch.id : null;
+        const nextMatchID = nextMatch !== null ? nextMatch.id : null;
+        if (prevMatchID !== nextMatchID) {
+          if (nextMatchID !== null) {
+            // We moved forward, unlocking a deeper node.
+            this._bridge.send('selectFiber', nextMatchID);
+          }
+        }
+        if (nextMatch !== null && nextMatch.isFullMatch) {
+          // We've just unlocked the innermost selected node.
+          // There's no point tracking it further.
+          this._persistedSelection = null;
+          this._persistedSelectionMatch = null;
+          renderer.setTrackedPath(null);
+        }
+      }
+    }
   };
 
   _onClick = (event: MouseEvent) => {
@@ -530,4 +602,22 @@ export default class Agent extends EventEmitter {
     // because those are usually unintentional as you lift the cursor.
     { leading: false }
   );
+
+  _throttledPersistSelection = throttle((rendererID: number, id: number) => {
+    // This is throttled, so both renderer and selected ID
+    // might not be available by the time we read them.
+    // This is why we need the defensive checks here.
+    const renderer = this._rendererInterfaces[rendererID];
+    const path = renderer != null ? renderer.getPathForElement(id) : null;
+    if (typeof sessionStorage !== 'undefined') {
+      if (path !== null) {
+        sessionStorage.setItem(
+          SESSION_STORAGE_LAST_SELECTION_KEY,
+          JSON.stringify(({ rendererID, path }: PersistedSelection))
+        );
+      } else {
+        sessionStorage.removeItem(SESSION_STORAGE_LAST_SELECTION_KEY);
+      }
+    }
+  }, 1000);
 }

src/backend/renderer.js

@@ -36,6 +36,8 @@ import type {
   Interaction,
   Interactions,
   InteractionWithCommits,
+  PathFrame,
+  PathMatch,
   ProfilingSummary,
   ReactRenderer,
   RendererInterface,
@@ -537,6 +539,11 @@ export function attach(
   // When a mount or update is in progress, this value tracks the root that is being operated on.
   let currentRootID: number = -1;
 
+  // Track the order in which roots were added.
+  // We will use it to disambiguate roots when restoring selection between reloads.
+  let nextRootIndex = 0;
+  const rootInsertionOrder: Map<number, number> = new Map();
+
   function getFiberID(primaryFiber: Fiber): number {
     if (!fiberToIDMap.has(primaryFiber)) {
       const id = getUID();
@@ -758,6 +765,18 @@ export function attach(
   }
 
   function recordUnmount(fiber: Fiber, isSimulated: boolean) {
+    if (trackedPathMatchFiber !== null) {
+      // We're in the process of trying to restore previous selection.
+      // If this fiber matched but is being unmounted, there's no use trying.
+      // Reset the state so we don't keep holding onto it.
+      if (
+        fiber === trackedPathMatchFiber ||
+        fiber === trackedPathMatchFiber.alternate
+      ) {
+        setTrackedPath(null);
+      }
+    }
+
     const isRoot = fiber.tag === HostRoot;
     const primaryFiber = getPrimaryFiber(fiber);
     if (!fiberToIDMap.has(primaryFiber)) {
@@ -807,6 +826,12 @@ export function attach(
       debug('mountFiberRecursively()', fiber, parentFiber);
     }
 
+    // If we have the tree selection from previous reload, try to match this Fiber.
+    // Also remember whether to do the same for siblings.
+    const mightSiblingsBeOnTrackedPath = updateTrackedPathStateBeforeMount(
+      fiber
+    );
+
     const shouldIncludeInTree = !shouldFilterFiber(fiber);
     if (shouldIncludeInTree) {
       recordMount(fiber, parentFiber);
@@ -840,6 +865,10 @@ export function attach(
       }
     }
 
+    // We're exiting this Fiber now, and entering its siblings.
+    // If we have selection to restore, we might need to re-activate tracking.
+    updateTrackedPathStateAfterMount(mightSiblingsBeOnTrackedPath);
+
     if (traverseSiblings && fiber.sibling !== null) {
       mountFiberRecursively(fiber.sibling, parentFiber, true);
     }
@@ -1111,9 +1140,15 @@ export function attach(
         hook.emit('operations', ops);
       });
     } else {
+      // Before the traversals, remember to start tracking
+      // our path in case we have selection to restore.
+      if (trackedPath !== null) {
+        mightBeOnTrackedPath = true;
+      }
       // If we have not been profiling, then we can just walk the tree and build up its current state as-is.
       hook.getFiberRoots(rendererID).forEach(root => {
         currentRootID = getFiberID(getPrimaryFiber(root.current));
+        rootInsertionOrder.set(currentRootID, nextRootIndex++);
 
         if (isProfiling) {
           // If profiling is active, store commit time and duration, and the current interactions.
@@ -1151,6 +1186,12 @@ export function attach(
 
     currentRootID = getFiberID(getPrimaryFiber(current));
 
+    // Before the traversals, remember to start tracking
+    // our path in case we have selection to restore.
+    if (trackedPath !== null) {
+      mightBeOnTrackedPath = true;
+    }
+
     if (isProfiling) {
       // If profiling is active, store commit time and duration, and the current interactions.
       // The frontend may request this information after profiling has stopped.
@@ -1176,16 +1217,19 @@ export function attach(
         current.memoizedState != null && current.memoizedState.element != null;
       if (!wasMounted && isMounted) {
         // Mount a new root.
+        rootInsertionOrder.set(currentRootID, nextRootIndex++);
         mountFiberRecursively(current, null);
       } else if (wasMounted && isMounted) {
         // Update an existing root.
         updateFiberRecursively(current, alternate, null);
       } else if (wasMounted && !isMounted) {
         // Unmount an existing root.
+        rootInsertionOrder.delete(currentRootID);
         recordUnmount(current, false);
       }
     } else {
       // Mount a new root.
+      rootInsertionOrder.set(currentRootID, nextRootIndex++);
       mountFiberRecursively(current, null);
     }
 
@@ -1979,14 +2023,153 @@ export function attach(
     scheduleUpdate(fiber);
   }
 
+  // Remember if we're trying to restore the selection after reload.
+  // In that case, we'll do some extra checks for matching mounts.
+  let trackedPath: Array<PathFrame> | null = null;
+  let trackedPathMatchFiber: Fiber | null = null;
+  let trackedPathMatchDepth = -1;
+  let mightBeOnTrackedPath = false;
+
+  function setTrackedPath(path: Array<PathFrame> | null) {
+    if (path === null) {
+      trackedPathMatchFiber = null;
+      trackedPathMatchDepth = -1;
+      mightBeOnTrackedPath = false;
+    } else if (trackedPath !== null) {
+      throw new Error('Tracked path can only be set once.');
+    }
+    trackedPath = path;
+  }
+
+  // We call this before traversing a new mount.
+  // It remembers whether this Fiber is the next best match for tracked path.
+  // The return value signals whether we should keep matching siblings or not.
+  function updateTrackedPathStateBeforeMount(fiber: Fiber): boolean {
+    if (trackedPath === null || !mightBeOnTrackedPath) {
+      // Fast path: there's nothing to track so do nothing and ignore siblings.
+      return false;
+    }
+    const returnFiber = fiber.return;
+    const returnAlternate = returnFiber !== null ? returnFiber.alternate : null;
+    // By now we know there's some selection to restore, and this is a new Fiber.
+    // Is this newly mounted Fiber a direct child of the current best match?
+    // (This will also be true for new roots if we haven't matched anything yet.)
+    if (
+      trackedPathMatchFiber === returnFiber ||
+      (trackedPathMatchFiber === returnAlternate && returnAlternate !== null)
+    ) {
+      // Is this the next Fiber we should select? Let's compare the frames.
+      const actualFrame = getPathFrame(fiber);
+      const expectedFrame = trackedPath[trackedPathMatchDepth + 1];
+      if (expectedFrame === undefined) {
+        throw new Error('Expected to see a frame at the next depth.');
+      }
+      if (
+        actualFrame.index === expectedFrame.index &&
+        actualFrame.key === expectedFrame.key &&
+        actualFrame.displayName === expectedFrame.displayName
+      ) {
+        // We have our next match.
+        trackedPathMatchFiber = fiber;
+        trackedPathMatchDepth++;
+        // Are we out of frames to match?
+        if (trackedPathMatchDepth === trackedPath.length - 1) {
+          // There's nothing that can possibly match afterwards.
+          // Don't check the children.
+          mightBeOnTrackedPath = false;
+        } else {
+          // Check the children, as they might reveal the next match.
+          mightBeOnTrackedPath = true;
+        }
+        // In either case, since we have a match, we don't need
+        // to check the siblings. They'll never match.
+        return false;
+      }
+    }
+    // This Fiber's parent is on the path, but this Fiber itself isn't.
+    // There's no need to check its children--they won't be on the path either.
+    mightBeOnTrackedPath = false;
+    // However, one of its siblings may be on the path so keep searching.
+    return true;
+  }
+
+  function updateTrackedPathStateAfterMount(mightSiblingsBeOnTrackedPath) {
+    // updateTrackedPathStateBeforeMount() told us whether to match siblings.
+    // Now that we're entering siblings, let's use that information.
+    mightBeOnTrackedPath = mightSiblingsBeOnTrackedPath;
+  }
+
+  function getPathFrame(fiber: Fiber): PathFrame {
+    const { displayName, key } = getDataForFiber(fiber);
+    let index = fiber.index;
+    if (fiber.tag === HostRoot) {
+      // Roots don't have a real index.
+      // Instead, we'll use the order in which it mounted.
+      const id = getFiberID(getPrimaryFiber(fiber));
+      const order = rootInsertionOrder.get(id);
+      if (typeof order !== 'number') {
+        throw new Error('Expected mounted root to have known insertion order.');
+      }
+      index = order;
+    }
+    return {
+      displayName,
+      key,
+      index,
+    };
+  }
+
+  // Produces a serializable representation that does a best effort
+  // of identifying a particular Fiber between page reloads.
+  // The return path will contain Fibers that are "invisible" to the store
+  // because their keys and indexes are important to restoring the selection.
+  function getPathForElement(id: number): Array<PathFrame> | null {
+    let fiber = idToFiberMap.get(id);
+    if (fiber == null) {
+      return null;
+    }
+    const keyPath = [];
+    while (fiber !== null) {
+      keyPath.push(getPathFrame(fiber));
+      fiber = fiber.return;
+    }
+    keyPath.reverse();
+    return keyPath;
+  }
+
+  function getBestMatchForTrackedPath(): PathMatch | null {
+    if (trackedPath === null) {
+      // Nothing to match.
+      return null;
+    }
+    if (trackedPathMatchFiber === null) {
+      // We didn't find anything.
+      return null;
+    }
+    // Find the closest Fiber store is aware of.
+    let fiber = trackedPathMatchFiber;
+    while (fiber !== null && shouldFilterFiber(fiber)) {
+      fiber = fiber.return;
+    }
+    if (fiber === null) {
+      return null;
+    }
+    return {
+      id: getFiberID(getPrimaryFiber(fiber)),
+      isFullMatch: trackedPathMatchDepth === trackedPath.length - 1,
+    };
+  }
+
   return {
     cleanup,
     flushInitialOperations,
+    getBestMatchForTrackedPath,
     getCommitDetails,
     getFiberIDFromNative,
     getFiberCommits,
     getInteractions,
     findNativeByFiberID,
+    getPathForElement,
     getProfilingDataForDownload,
     getProfilingSummary,
     handleCommitFiberRoot,
@@ -2001,6 +2184,7 @@ export function attach(
     setInHook,
     setInProps,
     setInState,
+    setTrackedPath,
     startProfiling,
     stopProfiling,
   };

src/backend/types.js

@@ -14,7 +14,7 @@ export type Fiber = Object;
 // (e.g. props, state, context, hooks) then we could add a bitmask field for this
 // to keep the number of attributes small.
 export type FiberData = {|
-  key: React$Key | null,
+  key: string | null,
   displayName: string | null,
   type: ElementType,
 |};
@@ -90,10 +90,22 @@ export type ProfilingSummary = {|
   rootID: number,
 |};
 
+export type PathFrame = {|
+  key: string | null,
+  index: number,
+  displayName: string | null,
+|};
+
+export type PathMatch = {|
+  id: number,
+  isFullMatch: boolean,
+|};
+
 export type RendererInterface = {
   cleanup: () => void,
   findNativeByFiberID: (id: number) => ?NativeType,
   flushInitialOperations: () => void,
+  getBestMatchForTrackedPath: () => PathMatch | null,
   getCommitDetails: (rootID: number, commitIndex: number) => CommitDetails,
   getFiberIDFromNative: (
     component: NativeType,
@@ -103,6 +115,7 @@ export type RendererInterface = {
   getInteractions: (rootID: number) => Interactions,
   getProfilingDataForDownload: (rootID: number) => Object,
   getProfilingSummary: (rootID: number) => ProfilingSummary,
+  getPathForElement: (id: number) => Array<PathFrame> | null,
   handleCommitFiberRoot: (fiber: Object) => void,
   handleCommitFiberUnmount: (fiber: Object) => void,
   inspectElement: (id: number) => InspectedElement | null,
@@ -120,6 +133,7 @@ export type RendererInterface = {
   ) => void,
   setInProps: (id: number, path: Array<string | number>, value: any) => void,
   setInState: (id: number, path: Array<string | number>, value: any) => void,
+  setTrackedPath: (path: Array<PathFrame> | null) => void,
   startProfiling: () => void,
   stopProfiling: () => void,
 };

src/constants.js

@@ -8,4 +8,7 @@ export const TREE_OPERATION_UPDATE_TREE_BASE_DURATION = 4;
 export const LOCAL_STORAGE_RELOAD_AND_PROFILE_KEY =
   'React::DevTools::reloadAndProfile';
 
+export const SESSION_STORAGE_LAST_SELECTION_KEY =
+  'React::DevTools::lastSelection';
+
 export const __DEBUG__ = false;