Make sure blocks are added before the batch is flushed

Author: janicduplessis

Libraries/NativeAnimation/RCTNativeAnimatedModule.h

@@ -10,9 +10,10 @@
 #import <React/RCTBridgeModule.h>
 #import <React/RCTEventDispatcher.h>
 #import <React/RCTEventEmitter.h>
+#import <React/RCTUIManager.h>
 
 #import "RCTValueAnimatedNode.h"
 
-@interface RCTNativeAnimatedModule : RCTEventEmitter <RCTBridgeModule, RCTValueAnimatedNodeObserver, RCTEventDispatcherObserver>
+@interface RCTNativeAnimatedModule : RCTEventEmitter <RCTBridgeModule, RCTValueAnimatedNodeObserver, RCTEventDispatcherObserver, RCTUIManagerObserver>
 
 @end

Libraries/NativeAnimation/RCTNativeAnimatedModule.m

@@ -28,6 +28,7 @@ - (void)invalidate
 {
   [_nodesManager stopAnimationLoop];
   [self.bridge.eventDispatcher removeDispatchObserver:self];
+  [self.bridge.uiManager removeUIManagerObserver:self];
 }
 
 - (dispatch_queue_t)methodQueue
@@ -44,6 +45,7 @@ - (void)setBridge:(RCTBridge *)bridge
   _preOperations = [NSMutableArray new];
 
   [bridge.eventDispatcher addDispatchObserver:self];
+  [bridge.uiManager addUIManagerObserver:self];
 }
 
 #pragma mark -- API
@@ -191,7 +193,7 @@ - (void)addPreOperationBlock:(AnimatedOperation)operation
   [_preOperations addObject:operation];
 }
 
-- (void)batchDidComplete
+- (void)uiManagerWillFlushUIBlocks:(RCTUIManager *)uiManager
 {
   if (_preOperations.count == 0 && _operations.count == 0) {
     return;
@@ -202,13 +204,13 @@ - (void)batchDidComplete
   _preOperations = [NSMutableArray new];
   _operations = [NSMutableArray new];
 
-  [self.bridge.uiManager prependUIBlock:^(__unused RCTUIManager *uiManager, __unused NSDictionary<NSNumber *, UIView *> *viewRegistry) {
+  [uiManager prependUIBlock:^(__unused RCTUIManager *manager, __unused NSDictionary<NSNumber *, UIView *> *viewRegistry) {
     for (AnimatedOperation operation in preOperations) {
       operation(_nodesManager);
     }
   }];
 
-  [self.bridge.uiManager addUIBlock:^(__unused RCTUIManager *uiManager, __unused NSDictionary<NSNumber *, UIView *> *viewRegistry) {
+  [uiManager addUIBlock:^(__unused RCTUIManager *manager, __unused NSDictionary<NSNumber *, UIView *> *viewRegistry) {
     for (AnimatedOperation operation in operations) {
       operation(_nodesManager);
     }

React/Modules/RCTUIManager.h

@@ -50,6 +50,20 @@ RCT_EXTERN NSString *const RCTUIManagerRootViewKey;
 
 @class RCTUIManager;
 
+/**
+ * Allows to hook into UIManager internals. This can be used to execute code at
+ * specific points during the view updating process.
+ */
+@protocol RCTUIManagerObserver <NSObject>
+
+/**
+ * Called before flushing UI blocks at the end of a batch. This is called from
+ * the UIManager queue. Can be used to add UI operations in that batch.
+ */
+- (void)uiManagerWillFlushUIBlocks:(RCTUIManager *)manager;
+
+@end
+
 @protocol RCTScrollableProtocol;
 
 /**
@@ -112,6 +126,17 @@ RCT_EXTERN NSString *const RCTUIManagerRootViewKey;
  */
 - (void)prependUIBlock:(RCTViewManagerUIBlock)block;
 
+/**
+ * Add a UIManagerObserver. See the RCTUIManagerObserver protocol for more info. This
+ * method can be called safely from any queue.
+ */
+- (void)addUIManagerObserver:(id<RCTUIManagerObserver>)observer;
+
+/**
+ * Remove a UIManagerObserver. This method can be called safely from any queue.
+ */
+- (void)removeUIManagerObserver:(id<RCTUIManagerObserver>)observer;
+
 /**
  * Used by native animated module to bypass the process of updating the values through the shadow
  * view hierarchy. This method will directly update native views, which means that updates for

React/Modules/RCTUIManager.m

@@ -224,6 +224,7 @@ @implementation RCTUIManager
   NSDictionary *_componentDataByName;
 
   NSMutableSet<id<RCTComponent>> *_bridgeTransactionListeners;
+  NSMutableSet<id<RCTUIManagerObserver>> *_uiManagerObservers;
 #if !TARGET_OS_TV
   UIInterfaceOrientation _currentInterfaceOrientation;
 #endif
@@ -326,6 +327,7 @@ - (void)setBridge:(RCTBridge *)bridge
   _rootViewTags = [NSMutableSet new];
 
   _bridgeTransactionListeners = [NSMutableSet new];
+  _uiManagerObservers = [NSMutableSet new];
 
   _viewsToBeDeleted = [NSMutableSet new];
 
@@ -736,6 +738,20 @@ - (void)_amendPendingUIBlocksWithStylePropagationUpdateForShadowView:(RCTShadowV
   }
 }
 
+- (void)addUIManagerObserver:(id<RCTUIManagerObserver>)observer
+{
+  dispatch_async(RCTGetUIManagerQueue(), ^{
+    [_uiManagerObservers addObject:observer];
+  });
+}
+
+- (void)removeUIManagerObserver:(id<RCTUIManagerObserver>)observer
+{
+  dispatch_async(RCTGetUIManagerQueue(), ^{
+    [_uiManagerObservers removeObject:observer];
+  });
+}
+
 /**
  * A method to be called from JS, which takes a container ID and then releases
  * all subviews for that container upon receipt.
@@ -1177,6 +1193,10 @@ - (void)_layoutAndMount
     }
   }];
 
+  for (id<RCTUIManagerObserver> observer in _uiManagerObservers) {
+    [observer uiManagerWillFlushUIBlocks:self];
+  }
+
   [self flushUIBlocks];
 }