refactor: improve trace updates highlighting (#41749)

Summary:
Pull Request resolved: #41749

Changelog: [Internal]

- Each trace update frame will have its own unique id, which will be the react tag of the corresponding host fiber.
- Based on these ids, previous frames will be rerendered, not
just removed once new frames are sent from React DevTools.
  - Imagine a case when there are 2 components on the screen: the first one rerenders once in a second and the second component rerenders much more frequently, each 5 milliseconds. With our previous implementation, update frames for first component will be removed once the second component has been rerendered.
- Each frame will have a lifetime for 2 seconds, it resets if frame with the same id was sent again from JS (basically component rerendered again, while we were highlighting it).

Android demo:
https://pxl.cl/3Vllz

Reviewed By: sammy-SC

Differential Revision: D51708054

fbshipit-source-id: 7abff9c1a334dccb3a1c08a46487d4bb99cdc448

packages/react-native/Libraries/Debugging/DebuggingOverlay.js

@@ -10,7 +10,7 @@
 
 import type {
   ElementRectangle,
-  Overlay,
+  TraceUpdate,
 } from './DebuggingOverlayNativeComponent';
 
 import View from '../Components/View/View';
@@ -26,7 +26,7 @@ const isNativeComponentReady =
   UIManager.hasViewManagerConfig('DebuggingOverlay');
 
 type DebuggingOverlayHandle = {
-  highlightTraceUpdates(updates: Overlay[]): void,
+  highlightTraceUpdates(updates: TraceUpdate[]): void,
   highlightElements(elements: ElementRectangle[]): void,
   clearElementsHighlight(): void,
 };
@@ -44,11 +44,11 @@ function DebuggingOverlay(
         }
 
         const nonEmptyRectangles = updates.filter(
-          ({rect, color}) => rect.width >= 0 && rect.height >= 0,
+          ({rectangle, color}) => rectangle.width >= 0 && rectangle.height >= 0,
         );
 
         if (nativeComponentRef.current != null) {
-          Commands.draw(
+          Commands.highlightTraceUpdates(
             nativeComponentRef.current,
             JSON.stringify(nonEmptyRectangles),
           );

packages/react-native/Libraries/Debugging/DebuggingOverlayNativeComponent.js

@@ -20,8 +20,10 @@ type NativeProps = $ReadOnly<{|
   ...ViewProps,
 |}>;
 export type DebuggingOverlayNativeComponentType = HostComponent<NativeProps>;
-export type Overlay = {
-  rect: ElementRectangle,
+
+export type TraceUpdate = {
+  id: number,
+  rectangle: ElementRectangle,
   color: ?ProcessedColorValue,
 };
 
@@ -33,11 +35,11 @@ export type ElementRectangle = {
 };
 
 interface NativeCommands {
-  +draw: (
+  +highlightTraceUpdates: (
     viewRef: React.ElementRef<DebuggingOverlayNativeComponentType>,
     // TODO(T144046177): Ideally we can pass array of Overlay, but currently
     // Array type is not supported in RN codegen for building native commands.
-    overlays: string,
+    updates: string,
   ) => void;
   +highlightElements: (
     viewRef: React.ElementRef<DebuggingOverlayNativeComponentType>,
@@ -50,7 +52,11 @@ interface NativeCommands {
 }
 
 export const Commands: NativeCommands = codegenNativeCommands<NativeCommands>({
-  supportedCommands: ['draw', 'highlightElements', 'clearElementsHighlights'],
+  supportedCommands: [
+    'highlightTraceUpdates',
+    'highlightElements',
+    'clearElementsHighlights',
+  ],
 });
 
 export default (codegenNativeComponent<NativeProps>(

packages/react-native/Libraries/Debugging/DebuggingOverlayRegistry.js

@@ -20,8 +20,9 @@ import type {
   ReactDevToolsAgentEvents,
   ReactDevToolsGlobalHook,
 } from '../Types/ReactDevToolsTypes';
-import type {Overlay} from './DebuggingOverlayNativeComponent';
+import type {TraceUpdate} from './DebuggingOverlayNativeComponent';
 
+import {findNodeHandle} from '../ReactNative/RendererProxy';
 import processColor from '../StyleSheet/processColor';
 
 // TODO(T171193075): __REACT_DEVTOOLS_GLOBAL_HOOK__ is always injected in dev-bundles,
@@ -98,7 +99,7 @@ class DebuggingOverlayRegistry {
   #onDrawTraceUpdates: (
     ...ReactDevToolsAgentEvents['drawTraceUpdates']
   ) => void = traceUpdates => {
-    const promisesToResolve: Array<Promise<Overlay>> = [];
+    const promisesToResolve: Array<Promise<TraceUpdate>> = [];
 
     traceUpdates.forEach(({node, color}) => {
       const publicInstance = this.#getPublicInstanceFromInstance(node);
@@ -107,11 +108,18 @@ class DebuggingOverlayRegistry {
         return;
       }
 
-      const frameToDrawPromise = new Promise<Overlay>(resolve => {
+      const frameToDrawPromise = new Promise<TraceUpdate>((resolve, reject) => {
         // TODO(T171095283): We should refactor this to use `getBoundingClientRect` when Paper is no longer supported.
         publicInstance.measure((x, y, width, height, left, top) => {
+          const id = findNodeHandle(node);
+          if (id == null) {
+            reject();
+            return;
+          }
+
           resolve({
-            rect: {x: left, y: top, width, height},
+            id,
+            rectangle: {x: left, y: top, width, height},
             color: processColor(color),
           });
         });

packages/react-native/Libraries/__tests__/__snapshots__/public-api-test.js.snap

@@ -4751,7 +4751,7 @@ declare export function createNodeList<T>(
 
 exports[`public API should not change unintentionally Libraries/Debugging/DebuggingOverlay.js 1`] = `
 "type DebuggingOverlayHandle = {
-  highlightTraceUpdates(updates: Overlay[]): void,
+  highlightTraceUpdates(updates: TraceUpdate[]): void,
   highlightElements(elements: ElementRectangle[]): void,
   clearElementsHighlight(): void,
 };
@@ -4769,8 +4769,9 @@ exports[`public API should not change unintentionally Libraries/Debugging/Debugg
   ...ViewProps,
 |}>;
 export type DebuggingOverlayNativeComponentType = HostComponent<NativeProps>;
-export type Overlay = {
-  rect: ElementRectangle,
+export type TraceUpdate = {
+  id: number,
+  rectangle: ElementRectangle,
   color: ?ProcessedColorValue,
 };
 export type ElementRectangle = {
@@ -4780,9 +4781,9 @@ export type ElementRectangle = {
   height: number,
 };
 interface NativeCommands {
-  +draw: (
+  +highlightTraceUpdates: (
     viewRef: React.ElementRef<DebuggingOverlayNativeComponentType>,
-    overlays: string
+    updates: string
   ) => void;
   +highlightElements: (
     viewRef: React.ElementRef<DebuggingOverlayNativeComponentType>,

packages/react-native/React/Fabric/Mounting/ComponentViews/DebuggingOverlay/RCTDebuggingOverlayComponentView.mm

@@ -50,9 +50,9 @@ - (void)handleCommand:(const NSString *)commandName args:(const NSArray *)args
   RCTDebuggingOverlayHandleCommand(self, commandName, args);
 }
 
-- (void)draw:(NSString *)overlays
+- (void)highlightTraceUpdates:(NSString *)updates
 {
-  [_overlay draw:overlays];
+  [_overlay highlightTraceUpdates:updates];
 }
 
 - (void)highlightElements:(NSString *)elements

packages/react-native/React/Views/RCTDebuggingOverlay.h

@@ -9,10 +9,18 @@
 
 #import <React/RCTView.h>
 
+@interface TraceUpdateTuple : NSObject
+
+@property (nonatomic, strong, readonly) UIView *view;
+@property (nonatomic, copy, readonly) dispatch_block_t cleanupBlock;
+
+- (instancetype)initWithView:(UIView *)view cleanupBlock:(dispatch_block_t)cleanupBlock;
+
+@end
+
 @interface RCTDebuggingOverlay : RCTView
 
-- (void)draw:(NSString *)serializedNodes;
-- (void)clearTraceUpdatesViews;
+- (void)highlightTraceUpdates:(NSString *)serializedUpdates;
 - (void)highlightElements:(NSString *)serializedElements;
 - (void)clearElementsHighlights;
 

packages/react-native/React/Views/RCTDebuggingOverlay.m

@@ -11,55 +11,99 @@
 #import <React/RCTLog.h>
 #import <React/RCTUtils.h>
 
+@implementation TraceUpdateTuple
+
+- (instancetype)initWithView:(UIView *)view cleanupBlock:(dispatch_block_t)cleanupBlock
+{
+  if (self = [super init]) {
+    _view = view;
+    _cleanupBlock = cleanupBlock;
+  }
+
+  return self;
+}
+
+@end
+
 @implementation RCTDebuggingOverlay {
   NSMutableArray<UIView *> *_highlightedElements;
-  NSMutableArray<UIView *> *_highlightedTraceUpdates;
+  NSMutableDictionary<NSNumber *, TraceUpdateTuple *> *_idToTraceUpdateMap;
 }
 
-- (void)draw:(NSString *)serializedNodes
+- (instancetype)initWithFrame:(CGRect)frame
 {
-  [self clearTraceUpdatesViews];
+  self = [super initWithFrame:frame];
+  if (self) {
+    _idToTraceUpdateMap = [NSMutableDictionary new];
+  }
+  return self;
+}
 
+- (void)highlightTraceUpdates:(NSString *)serializedUpdates
+{
   NSError *error = nil;
-  id deserializedNodes = RCTJSONParse(serializedNodes, &error);
+  id deserializedUpdates = RCTJSONParse(serializedUpdates, &error);
 
   if (error) {
-    RCTLogError(@"Failed to parse serialized nodes passed to RCTDebuggingOverlay");
+    RCTLogError(@"Failed to parse serialized updates passed to RCTDebuggingOverlay");
     return;
   }
 
-  if (![deserializedNodes isKindOfClass:[NSArray class]]) {
-    RCTLogError(@"Expected to receive nodes as an array, got %@", NSStringFromClass([deserializedNodes class]));
+  if (![deserializedUpdates isKindOfClass:[NSArray class]]) {
+    RCTLogError(@"Expected to receive updates as an array, got %@", NSStringFromClass([deserializedUpdates class]));
     return;
   }
 
-  _highlightedTraceUpdates = [NSMutableArray new];
-  for (NSDictionary *node in deserializedNodes) {
-    NSDictionary *nodeRectangle = node[@"rect"];
-    NSNumber *nodeColor = node[@"color"];
+  for (NSDictionary *update in deserializedUpdates) {
+    NSNumber *identifier = [RCTConvert NSNumber:update[@"id"]];
+    NSDictionary *nodeRectangle = update[@"rectangle"];
+    UIColor *nodeColor = [RCTConvert UIColor:update[@"color"]];
 
     CGRect rect = [RCTConvert CGRect:nodeRectangle];
 
+    TraceUpdateTuple *possiblyRegisteredTraceUpdateTuple = [_idToTraceUpdateMap objectForKey:identifier];
+    if (possiblyRegisteredTraceUpdateTuple != nil) {
+      dispatch_block_t cleanupBlock = [possiblyRegisteredTraceUpdateTuple cleanupBlock];
+      UIView *view = [possiblyRegisteredTraceUpdateTuple view];
+
+      dispatch_block_cancel(cleanupBlock);
+
+      view.frame = rect;
+      view.layer.borderColor = nodeColor.CGColor;
+
+      dispatch_block_t newCleanupBlock = dispatch_block_create(0, ^{
+        [self->_idToTraceUpdateMap removeObjectForKey:identifier];
+        [view removeFromSuperview];
+      });
+
+      [_idToTraceUpdateMap setObject:[[TraceUpdateTuple alloc] initWithView:view cleanupBlock:newCleanupBlock]
+                              forKey:identifier];
+
+      dispatch_after(dispatch_time(DISPATCH_TIME_NOW, 2 * NSEC_PER_SEC), dispatch_get_main_queue(), newCleanupBlock);
+
+      continue;
+    }
+
     UIView *box = [[UIView alloc] initWithFrame:rect];
     box.backgroundColor = [UIColor clearColor];
 
     box.layer.borderWidth = 2.0f;
-    box.layer.borderColor = [RCTConvert UIColor:nodeColor].CGColor;
+    box.layer.borderColor = nodeColor.CGColor;
+
+    dispatch_block_t unmountViewAndPerformCleanup = dispatch_block_create(0, ^{
+      [self->_idToTraceUpdateMap removeObjectForKey:identifier];
+      [box removeFromSuperview];
+    });
+
+    TraceUpdateTuple *traceUpdateTuple = [[TraceUpdateTuple alloc] initWithView:box
+                                                                   cleanupBlock:unmountViewAndPerformCleanup];
 
+    [_idToTraceUpdateMap setObject:traceUpdateTuple forKey:identifier];
     [self addSubview:box];
-    [_highlightedTraceUpdates addObject:box];
-  }
-}
 
-- (void)clearTraceUpdatesViews
-{
-  if (_highlightedTraceUpdates != nil) {
-    for (UIView *v in _highlightedTraceUpdates) {
-      [v removeFromSuperview];
-    }
+    dispatch_after(
+        dispatch_time(DISPATCH_TIME_NOW, 2 * NSEC_PER_SEC), dispatch_get_main_queue(), unmountViewAndPerformCleanup);
   }
-
-  _highlightedTraceUpdates = nil;
 }
 
 - (void)highlightElements:(NSString *)serializedElements

packages/react-native/React/Views/RCTDebuggingOverlayManager.m

@@ -22,13 +22,13 @@ - (UIView *)view
   return [RCTDebuggingOverlay new];
 }
 
-RCT_EXPORT_METHOD(draw : (nonnull NSNumber *)viewTag nodes : (NSString *)serializedNodes)
+RCT_EXPORT_METHOD(highlightTraceUpdates : (nonnull NSNumber *)viewTag nodes : (NSString *)serializedUpdates)
 {
   [self.bridge.uiManager addUIBlock:^(RCTUIManager *uiManager, NSDictionary<NSNumber *, UIView *> *viewRegistry) {
     UIView *view = viewRegistry[viewTag];
 
     if ([view isKindOfClass:[RCTDebuggingOverlay class]]) {
-      [(RCTDebuggingOverlay *)view draw:serializedNodes];
+      [(RCTDebuggingOverlay *)view highlightTraceUpdates:serializedUpdates];
     } else {
       RCTLogError(@"Expected view to be RCTDebuggingOverlay, got %@", NSStringFromClass([view class]));
     }