Be able to handle different snapshots for different roots

Summary:
^

This change allows to take different snapshots for different nodes and render them each on the visualiser.

At the moment, more than likely, this is not really used. At the same time, it fixes an issue whereas any subtree update can override and set the only visible snapshot.

Reviewed By: LukeDefeo, antonk52

Differential Revision: D39821920

fbshipit-source-id: ab8f6a4a2a5e96801c951a4e3009cc571a617f22

desktop/plugins/public/ui-debugger/components/Visualization2D.tsx

@@ -8,7 +8,7 @@
  */
 
 import React from 'react';
-import {Id, Tag, UINode} from '../types';
+import {Id, Snapshot, Tag, UINode} from '../types';
 import {styled, Layout, theme} from 'flipper-plugin';
 import {Typography} from 'antd';
 
@@ -16,7 +16,7 @@ export const Visualization2D: React.FC<
   {
     root: Id;
     nodes: Map<Id, UINode>;
-    snapshot?: String;
+    snapshots: Map<Id, Snapshot>;
     hoveredNode?: Id;
     selectedNode?: Id;
     onSelectNode: (id: Id) => void;
@@ -26,7 +26,7 @@ export const Visualization2D: React.FC<
 > = ({
   root,
   nodes,
-  snapshot,
+  snapshots,
   hoveredNode,
   selectedNode,
   onSelectNode,
@@ -35,6 +35,7 @@ export const Visualization2D: React.FC<
 }) => {
   //todo, do a bfs search for the first bounds found
   const rootBounds = nodes.get(root)?.bounds;
+  const rootSnapshot = snapshots.get(root);
 
   if (!rootBounds) {
     return null;
@@ -61,15 +62,16 @@ export const Visualization2D: React.FC<
           height: toPx(rootBounds.height),
         }}>
         <OuterBorder />
-        {snapshot ? (
+        {rootSnapshot ? (
           <img
-            src={'data:image/jpeg;base64,' + snapshot}
+            src={'data:image/jpeg;base64,' + rootSnapshot}
             style={{maxWidth: '100%'}}
           />
         ) : null}
         <Visualization2DNode
           nodeId={root}
           nodes={nodes}
+          snapshots={snapshots}
           hoveredNode={hoveredNode}
           selectedNode={selectedNode}
           onSelectNode={onSelectNode}
@@ -85,6 +87,7 @@ function Visualization2DNode({
   parentId,
   nodeId,
   nodes,
+  snapshots,
   hoveredNode,
   selectedNode,
   onSelectNode,
@@ -94,6 +97,7 @@ function Visualization2DNode({
   nodeId: Id;
   parentId?: Id;
   nodes: Map<Id, UINode>;
+  snapshots: Map<Id, Snapshot>;
   modifierPressed: boolean;
   hoveredNode?: Id;
   selectedNode?: Id;
@@ -101,6 +105,7 @@ function Visualization2DNode({
   onHoverNode: (id?: Id) => void;
 }) {
   const node = nodes.get(nodeId);
+  const snapshot = snapshots.get(nodeId);
 
   if (!node) {
     return null;
@@ -130,6 +135,7 @@ function Visualization2DNode({
       key={childId}
       nodeId={childId}
       nodes={nodes}
+      snapshots={snapshots}
       hoveredNode={hoveredNode}
       onSelectNode={onSelectNode}
       onHoverNode={onHoverNode}
@@ -177,11 +183,19 @@ function Visualization2DNode({
         onSelectNode(nodeId);
       }}>
       <NodeBorder tags={node.tags}></NodeBorder>
-
+      {snapshot ? (
+        <img
+          src={'data:image/jpeg;base64,' + snapshot}
+          style={{maxWidth: '100%'}}
+        />
+      ) : (
+        <>
           {/* Dirty hack to avoid showing highly overlapping text */}
           {!hasOverlappingChild && !isZeroWidthOrHeight && node.bounds
             ? node.name
             : null}
+        </>
+      )}
       {children}
     </div>
   );

desktop/plugins/public/ui-debugger/components/main.tsx

@@ -19,7 +19,7 @@ import {
 import {Typography} from 'antd';
 
 import {useHotkeys} from 'react-hotkeys-hook';
-import {Id, UINode} from '../types';
+import {Id, Snapshot, UINode} from '../types';
 import {PerfStats} from './PerfStats';
 import {Tree} from './Tree';
 import {Visualization2D} from './Visualization2D';
@@ -29,7 +29,7 @@ export function Component() {
   const instance = usePlugin(plugin);
   const rootId = useValue(instance.rootId);
   const nodes: Map<Id, UINode> = useValue(instance.nodes);
-  const snapshot: String | undefined = useValue(instance.snapshot);
+  const snapshots: Map<Id, Snapshot> = useValue(instance.snapshots);
 
   const [showPerfStats, setShowPerfStats] = useState(false);
   const [selectedNode, setSelectedNode] = useState<Id | undefined>(undefined);
@@ -72,7 +72,7 @@ export function Component() {
             <Visualization2D
               root={rootId}
               nodes={nodes}
-              snapshot={snapshot}
+              snapshots={snapshots}
               hoveredNode={hoveredNode}
               onHoverNode={setHoveredNode}
               selectedNode={selectedNode}

desktop/plugins/public/ui-debugger/index.tsx

@@ -8,7 +8,7 @@
  */
 
 import {PluginClient, createState, createDataSource} from 'flipper-plugin';
-import {Events, Id, PerfStatsEvent, UINode} from './types';
+import {Events, Id, PerfStatsEvent, Snapshot, UINode} from './types';
 
 export function plugin(client: PluginClient<Events>) {
   const rootId = createState<Id | undefined>(undefined);
@@ -23,11 +23,13 @@ export function plugin(client: PluginClient<Events>) {
   });
 
   const nodesAtom = createState<Map<Id, UINode>>(new Map());
-  const snapshotAtom = createState<String | undefined>(undefined);
-  client.onMessage('subtreeUpdate', ({nodes, snapshot}) => {
-    snapshotAtom.set(snapshot);
+  const snapshotsAtom = createState<Map<Id, Snapshot>>(new Map());
+  client.onMessage('subtreeUpdate', (event) => {
+    snapshotsAtom.update((draft) => {
+      draft.set(event.rootId, event.snapshot);
+    });
     nodesAtom.update((draft) => {
-      for (const node of nodes) {
+      for (const node of event.nodes) {
         draft.set(node.id, node);
       }
     });
@@ -38,7 +40,7 @@ export function plugin(client: PluginClient<Events>) {
     nodesAtom.set(new Map(nodes.map((node) => [node.id, node])));
   });
 
-  return {rootId, snapshot: snapshotAtom, nodes: nodesAtom, perfEvents};
+  return {rootId, snapshots: snapshotsAtom, nodes: nodesAtom, perfEvents};
 }
 
 export {Component} from './components/main';

desktop/plugins/public/ui-debugger/types.tsx

@@ -10,15 +10,17 @@
 export type Events = {
   init: {rootId: string};
   nativeScan: {txId: number; nodes: UINode[]};
-  subtreeUpdate: {
-    txId: number;
-    rootId: String;
-    nodes: UINode[];
-    snapshot: String;
-  };
+  subtreeUpdate: SubtreeUpdateEvent;
   perfStats: PerfStatsEvent;
 };
 
+export type SubtreeUpdateEvent = {
+  txId: number;
+  rootId: Id;
+  nodes: UINode[];
+  snapshot: Snapshot;
+};
+
 export type PerfStatsEvent = {
   txId: number;
   observerType: string;
@@ -48,6 +50,7 @@ export type Bounds = {
 };
 
 export type Id = string;
+export type Snapshot = string;
 
 export type Tag = 'Native' | 'Declarative' | 'Android' | 'Litho ';