Fix profiling screenshots data structure to map rootID to commitIndex

Propagate `rootID` throughout the code for `captureScreenshot`.

Rename private profiling maps of `store` to make relations more clear.

Fix missing cleanup for screenshots data in `set importedProfilingData` of `store`.

shells/browser/shared/src/background.js

@@ -122,7 +122,7 @@ chrome.runtime.onMessage.addListener((request, sender) => {
     }
 
     if (request.captureScreenshot) {
-      const { commitIndex } = request;
+      const { commitIndex, rootID } = request;
       try {
         chrome.tabs.captureVisibleTab(undefined, undefined, dataURL => {
           // TODO For some reason, sending a response using the third param (sendResponse) doesn't work,
@@ -134,6 +134,7 @@ chrome.runtime.onMessage.addListener((request, sender) => {
                 payload: {
                   commitIndex,
                   dataURL,
+                  rootID,
                 },
               });
             }

shells/browser/shared/src/main.js

@@ -77,11 +77,12 @@ function createPanelIfReactLoaded() {
             filename,
           });
         });
-        bridge.addListener('captureScreenshot', ({ commitIndex }) => {
+        bridge.addListener('captureScreenshot', ({ commitIndex, rootID }) => {
           chrome.runtime.sendMessage(
             {
               captureScreenshot: true,
               commitIndex,
+              rootID,
             },
             response => bridge.send('screenshotCaptured', response)
           );

shells/dev/src/backend.js

@@ -20,11 +20,12 @@ const bridge = new Bridge({
   },
 });
 
-bridge.addListener('captureScreenshot', ({ commitIndex }) => {
+bridge.addListener('captureScreenshot', ({ commitIndex, rootID }) => {
   html2canvas(document.body, { logging: false }).then(canvas => {
     bridge.send('screenshotCaptured', {
       commitIndex,
       dataURL: canvas.toDataURL(),
+      rootID,
     });
   });
 });

src/backend/agent.js

@@ -125,8 +125,14 @@ export default class Agent extends EventEmitter {
     }
   }
 
-  captureScreenshot = ({ commitIndex }: { commitIndex: number }) => {
-    this._bridge.send('captureScreenshot', { commitIndex });
+  captureScreenshot = ({
+    commitIndex,
+    rootID,
+  }: {
+    commitIndex: number,
+    rootID: number,
+  }) => {
+    this._bridge.send('captureScreenshot', { commitIndex, rootID });
   };
 
   getIDForNode(node: Object): number | null {

src/devtools/store.js

@@ -91,16 +91,20 @@ export default class Store extends EventEmitter {
   // Map of root (id) to a list of tree mutation that occur during profiling.
   // Once profiling is finished, these mutations can be used, along with the initial tree snapshots,
   // to reconstruct the state of each root for each commit.
-  _profilingOperations: Map<number, Array<Uint32Array>> = new Map();
+  _profilingOperationsByRootID: Map<number, Array<Uint32Array>> = new Map();
 
+  // Map of root (id) to a Map of screenshots by commit ID.
   // Stores screenshots for each commit (when profiling).
-  _profilingScreenshots: Map<number, string> = new Map();
+  _profilingScreenshotsByRootID: Map<number, Map<number, string>> = new Map();
 
   // Snapshot of the state of the main Store (including all roots) when profiling started.
   // Once profiling is finished, this snapshot can be used along with "operations" messages emitted during profiling,
   // to reconstruct the state of each root for each commit.
   // It's okay to use a single root to store this information because node IDs are unique across all roots.
-  _profilingSnapshot: Map<number, ProfilingSnapshotNode> = new Map();
+  _profilingSnapshotsByElementID: Map<
+    number,
+    ProfilingSnapshotNode
+  > = new Map();
 
   // Incremented each time the store is mutated.
   // This enables a passive effect to detect a mutation between render and commit phase.
@@ -187,14 +191,14 @@ export default class Store extends EventEmitter {
     if (this._ownersMap.size !== 0) {
       throw new Error('Expected _ownersMap to be empty.');
     }
-    if (this._profilingOperations.size !== 0) {
-      throw new Error('Expected _profilingOperations to be empty.');
+    if (this._profilingOperationsByRootID.size !== 0) {
+      throw new Error('Expected _profilingOperationsByRootID to be empty.');
     }
-    if (this._profilingScreenshots.size !== 0) {
-      throw new Error('Expected _profilingScreenshots to be empty.');
+    if (this._profilingScreenshotsByRootID.size !== 0) {
+      throw new Error('Expected _profilingScreenshotsByRootID to be empty.');
     }
-    if (this._profilingSnapshot.size !== 0) {
-      throw new Error('Expected _profilingSnapshot to be empty.');
+    if (this._profilingSnapshotsByElementID.size !== 0) {
+      throw new Error('Expected _profilingSnapshotsByElementID to be empty.');
     }
     if (this._rootIDToCapabilities.size !== 0) {
       throw new Error('Expected _rootIDToCapabilities to be empty.');
@@ -239,7 +243,8 @@ export default class Store extends EventEmitter {
   // Profiling data has been recorded for at least one root.
   get hasProfilingData(): boolean {
     return (
-      this._importedProfilingData !== null || this._profilingOperations.size > 0
+      this._importedProfilingData !== null ||
+      this._profilingOperationsByRootID.size > 0
     );
   }
 
@@ -248,8 +253,9 @@ export default class Store extends EventEmitter {
   }
   set importedProfilingData(value: ImportedProfilingData | null): void {
     this._importedProfilingData = value;
-    this._profilingOperations = new Map();
-    this._profilingSnapshot = new Map();
+    this._profilingOperationsByRootID = new Map();
+    this._profilingScreenshotsByRootID = new Map();
+    this._profilingSnapshotsByElementID = new Map();
     this._profilingCache.invalidate();
 
     this.emit('importedProfilingData');
@@ -268,15 +274,15 @@ export default class Store extends EventEmitter {
   }
 
   get profilingOperations(): Map<number, Array<Uint32Array>> {
-    return this._profilingOperations;
+    return this._profilingOperationsByRootID;
   }
 
-  get profilingScreenshots(): Map<number, string> {
-    return this._profilingScreenshots;
+  get profilingScreenshots(): Map<number, Map<number, string>> {
+    return this._profilingScreenshotsByRootID;
   }
 
   get profilingSnapshot(): Map<number, ProfilingSnapshotNode> {
-    return this._profilingSnapshot;
+    return this._profilingSnapshotsByElementID;
   }
 
   get revision(): number {
@@ -305,9 +311,9 @@ export default class Store extends EventEmitter {
 
   clearProfilingData(): void {
     this._importedProfilingData = null;
-    this._profilingOperations = new Map();
-    this._profilingScreenshots = new Map();
-    this._profilingSnapshot = new Map();
+    this._profilingOperationsByRootID = new Map();
+    this._profilingScreenshotsByRootID = new Map();
+    this._profilingSnapshotsByElementID = new Map();
 
     // Invalidate suspense cache if profiling data is being (re-)recorded.
     // Note that we clear now because any existing data is "stale".
@@ -639,17 +645,17 @@ export default class Store extends EventEmitter {
   }
 
   _captureScreenshot = throttle(
-    memoize((commitIndex: number) => {
-      this._bridge.send('captureScreenshot', { commitIndex });
+    memoize((rootID: number, commitIndex: number) => {
+      this._bridge.send('captureScreenshot', { commitIndex, rootID });
     }),
     THROTTLE_CAPTURE_SCREENSHOT_DURATION
   );
 
-  _takeProfilingSnapshotRecursive = (id: number) => {
-    const element = this.getElementByID(id);
+  _takeProfilingSnapshotRecursive = (elementID: number) => {
+    const element = this.getElementByID(elementID);
     if (element !== null) {
-      this._profilingSnapshot.set(id, {
-        id,
+      this._profilingSnapshotsByElementID.set(elementID, {
+        id: elementID,
         children: element.children.slice(0),
         displayName: element.displayName,
         key: element.key,
@@ -703,17 +709,17 @@ export default class Store extends EventEmitter {
     const rootID = operations[1];
 
     if (this._isProfiling) {
-      let profilingOperations = this._profilingOperations.get(rootID);
+      let profilingOperations = this._profilingOperationsByRootID.get(rootID);
       if (profilingOperations == null) {
         profilingOperations = [operations];
-        this._profilingOperations.set(rootID, profilingOperations);
+        this._profilingOperationsByRootID.set(rootID, profilingOperations);
       } else {
         profilingOperations.push(operations);
       }
 
       if (this._captureScreenshots) {
         const commitIndex = profilingOperations.length - 1;
-        this._captureScreenshot(commitIndex);
+        this._captureScreenshot(rootID, commitIndex);
       }
     }
 
@@ -990,9 +996,9 @@ export default class Store extends EventEmitter {
   onProfilingStatus = (isProfiling: boolean) => {
     if (isProfiling) {
       this._importedProfilingData = null;
-      this._profilingOperations = new Map();
-      this._profilingScreenshots = new Map();
-      this._profilingSnapshot = new Map();
+      this._profilingOperationsByRootID = new Map();
+      this._profilingScreenshotsByRootID = new Map();
+      this._profilingSnapshotsByElementID = new Map();
       this.roots.forEach(this._takeProfilingSnapshotRecursive);
     }
 
@@ -1011,11 +1017,23 @@ export default class Store extends EventEmitter {
   onScreenshotCaptured = ({
     commitIndex,
     dataURL,
+    rootID,
   }: {|
     commitIndex: number,
     dataURL: string,
+    rootID: number,
   |}) => {
-    this._profilingScreenshots.set(commitIndex, dataURL);
+    let profilingScreenshotsForRootByCommitIndex = this._profilingScreenshotsByRootID.get(
+      rootID
+    );
+    if (!profilingScreenshotsForRootByCommitIndex) {
+      profilingScreenshotsForRootByCommitIndex = new Map();
+      this._profilingScreenshotsByRootID.set(
+        rootID,
+        profilingScreenshotsForRootByCommitIndex
+      );
+    }
+    profilingScreenshotsForRootByCommitIndex.set(commitIndex, dataURL);
   };
 
   onBridgeShutdown = () => {

src/devtools/views/Profiler/SidebarCommitInfo.js

@@ -24,10 +24,13 @@ export default function SidebarCommitInfo(_: Props) {
     profilingScreenshots,
   } = useContext(StoreContext);
 
+  const screenshotsByCommitIndex =
+    rootID !== null ? profilingScreenshots.get(rootID) : null;
   const screenshot =
-    selectedCommitIndex !== null
-      ? profilingScreenshots.get(selectedCommitIndex)
+    screenshotsByCommitIndex != null && selectedCommitIndex !== null
+      ? screenshotsByCommitIndex.get(selectedCommitIndex)
       : null;
+
   const [
     isScreenshotModalVisible,
     setIsScreenshotModalVisible,