refactor: improve trace updates highlighting

Summary: Changelog: [Internal]

Differential Revision: https://internalfb.com/D51708054

fbshipit-source-id: e66c80d0253e1c5343777423a9525068b202e2e4

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/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]));
     }

packages/react-native/ReactAndroid/src/main/java/com/facebook/react/views/debuggingoverlay/DebuggingOverlay.java

@@ -13,54 +13,42 @@
 import android.graphics.RectF;
 import android.view.View;
 import androidx.annotation.UiThread;
-import com.facebook.react.uimanager.PixelUtil;
+import com.facebook.react.bridge.UiThreadUtil;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 
 public class DebuggingOverlay extends View {
 
-  private final Paint mOverlayPaint = new Paint();
-  private List<Overlay> mOverlays = new ArrayList<Overlay>();
+  private final Paint mTraceUpdatePaint = new Paint();
+  private HashMap<Integer, TraceUpdate> mTraceUpdatesToDisplayMap = new HashMap();
+  private HashMap<Integer, Runnable> mTraceUpdateIdToCleanupRunnableMap = new HashMap();
 
   private final Paint mHighlightedElementsPaint = new Paint();
   private List<RectF> mHighlightedElementsRectangles = new ArrayList<>();
 
-  public static class Overlay {
-
-    private final int mColor;
-    private final RectF mRect;
-
-    public Overlay(int color, RectF rect) {
-      mColor = color;
-      mRect = rect;
-    }
-
-    public int getColor() {
-      return mColor;
-    }
-
-    public RectF getPixelRect() {
-      return new RectF(
-          PixelUtil.toPixelFromDIP(mRect.left),
-          PixelUtil.toPixelFromDIP(mRect.top),
-          PixelUtil.toPixelFromDIP(mRect.right),
-          PixelUtil.toPixelFromDIP(mRect.bottom));
-    }
-  }
-
   public DebuggingOverlay(Context context) {
     super(context);
 
-    mOverlayPaint.setStyle(Paint.Style.STROKE);
-    mOverlayPaint.setStrokeWidth(6);
+    mTraceUpdatePaint.setStyle(Paint.Style.STROKE);
+    mTraceUpdatePaint.setStrokeWidth(6);
 
     mHighlightedElementsPaint.setStyle(Paint.Style.FILL);
     mHighlightedElementsPaint.setColor(0xCCC8E6FF);
   }
 
   @UiThread
-  public void setOverlays(List<Overlay> overlays) {
-    mOverlays = overlays;
+  public void setTraceUpdates(List<TraceUpdate> traceUpdates) {
+    for (TraceUpdate traceUpdate : traceUpdates) {
+      int traceUpdateId = traceUpdate.getId();
+      if (mTraceUpdateIdToCleanupRunnableMap.containsKey(traceUpdateId)) {
+        UiThreadUtil.removeOnUiThread(mTraceUpdateIdToCleanupRunnableMap.get(traceUpdateId));
+        mTraceUpdateIdToCleanupRunnableMap.remove(traceUpdateId);
+      }
+
+      mTraceUpdatesToDisplayMap.put(traceUpdateId, traceUpdate);
+    }
+
     invalidate();
   }
 
@@ -81,9 +69,23 @@ public void onDraw(Canvas canvas) {
     super.onDraw(canvas);
 
     // Draw border outside of the given overlays to be aligned with web trace highlights
-    for (Overlay overlay : mOverlays) {
-      mOverlayPaint.setColor(overlay.getColor());
-      canvas.drawRect(overlay.getPixelRect(), mOverlayPaint);
+    for (TraceUpdate traceUpdate : mTraceUpdatesToDisplayMap.values()) {
+      mTraceUpdatePaint.setColor(traceUpdate.getColor());
+      canvas.drawRect(traceUpdate.getRectangle(), mTraceUpdatePaint);
+
+      int traceUpdateId = traceUpdate.getId();
+      Runnable block =
+          () -> {
+            mTraceUpdatesToDisplayMap.remove(traceUpdateId);
+            mTraceUpdateIdToCleanupRunnableMap.remove(traceUpdateId);
+
+            invalidate();
+          };
+
+      if (!mTraceUpdateIdToCleanupRunnableMap.containsKey(traceUpdateId)) {
+        mTraceUpdateIdToCleanupRunnableMap.put(traceUpdateId, block);
+        UiThreadUtil.runOnUiThread(block, 2000);
+      }
     }
 
     for (RectF elementRectangle : mHighlightedElementsRectangles) {