Swizzle RCTUIManager

ios/LayoutReanimation/REAAnimationsManager.m

@@ -1,7 +1,7 @@
 #import <RNReanimated/REAAnimationsManager.h>
 #import <RNReanimated/REASharedElement.h>
 #import <RNReanimated/REASharedTransitionManager.h>
-#import <RNReanimated/REAUIManager.h>
+#import <RNReanimated/REASwizzledUIManager.h>
 #import <React/RCTComponentData.h>
 #import <React/RCTTextView.h>
 #import <React/UIView+Private.h>
@@ -30,7 +30,7 @@ BOOL REANodeFind(id<RCTComponent> view, int (^block)(id<RCTComponent>))
 
 @implementation REAAnimationsManager {
   RCTUIManager *_uiManager;
-  REAUIManager *_reaUiManager;
+  REASwizzledUIManager *_reaSwizzledUIManager;
   NSMutableSet<NSNumber *> *_enteringViews;
   NSMutableDictionary<NSNumber *, REASnapshot *> *_enteringViewTargetValues;
   NSMutableDictionary<NSNumber *, UIView *> *_exitingViews;
@@ -59,7 +59,6 @@ - (instancetype)initWithUIManager:(RCTUIManager *)uiManager
 {
   if (self = [super init]) {
     _uiManager = uiManager;
-    _reaUiManager = (REAUIManager *)uiManager;
     _exitingViews = [NSMutableDictionary new];
     _exitingSubviewsCountMap = [NSMutableDictionary new];
     _ancestorsToRemove = [NSMutableSet new];
@@ -74,6 +73,7 @@ - (instancetype)initWithUIManager:(RCTUIManager *)uiManager
       [_currentKeys addObject:[NSString stringWithFormat:@"current%@", [key capitalizedString]]];
     }
     _sharedTransitionManager = [[REASharedTransitionManager alloc] initWithAnimationsManager:self];
+    _reaSwizzledUIManager = [[REASwizzledUIManager alloc] initWithUIManager:uiManager withAnimationManager:self];
   }
   return self;
 }
@@ -106,7 +106,7 @@ - (void)setAnimationRemovingBlock:(REAAnimationRemovingBlock)clearAnimation
 - (UIView *)viewForTag:(NSNumber *)tag
 {
   UIView *view;
-  (view = [_reaUiManager viewForReactTag:tag]) || (view = [_exitingViews objectForKey:tag]) ||
+  (view = [_uiManager viewForReactTag:tag]) || (view = [_exitingViews objectForKey:tag]) ||
       (view = [_sharedTransitionManager getTransitioningView:tag]);
   return view;
 }

ios/LayoutReanimation/REASwizzledUIManager.h

@@ -0,0 +1,6 @@
+#import <RNReanimated/REAAnimationsManager.h>
+
+@interface REASwizzledUIManager : NSObject
+- (instancetype)initWithUIManager:(RCTUIManager *)uiManager
+             withAnimationManager:(REAAnimationsManager *)animationsManager;
+@end

ios/LayoutReanimation/REAUIManager.mm → ios/LayoutReanimation/REASwizzledUIManager.mm

@@ -1,91 +1,91 @@
-#import <Foundation/Foundation.h>
 #import <RNReanimated/FeaturesConfig.h>
-#import <RNReanimated/REAIOSScheduler.h>
-#import <RNReanimated/REAUIManager.h>
-#import <RNReanimated/Scheduler.h>
-#import <React/RCTComponentData.h>
+#import <RNReanimated/REASwizzledUIManager.h>
 #import <React/RCTLayoutAnimation.h>
 #import <React/RCTLayoutAnimationGroup.h>
-#import <React/RCTModalHostView.h>
 #import <React/RCTRootShadowView.h>
 #import <React/RCTRootViewInternal.h>
-#import <React/RCTUIManagerObserverCoordinator.h>
-
-#if __has_include(<RNScreens/RNSScreen.h>)
-#import <RNScreens/RNSScreen.h>
-#endif
-
-@interface RCTUIManager (REA)
-- (void)_manageChildren:(NSNumber *)containerTag
-        moveFromIndices:(NSArray<NSNumber *> *)moveFromIndices
-          moveToIndices:(NSArray<NSNumber *> *)moveToIndices
-      addChildReactTags:(NSArray<NSNumber *> *)addChildReactTags
-           addAtIndices:(NSArray<NSNumber *> *)addAtIndices
-        removeAtIndices:(NSArray<NSNumber *> *)removeAtIndices
-               registry:(NSMutableDictionary<NSNumber *, id<RCTComponent>> *)registry;
-
-- (RCTViewManagerUIBlock)uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *)rootShadowView;
+#import <React/RCTUIManager.h>
+#import <objc/runtime.h>
 
+@interface RCTUIManager (Reanimated)
+@property REAAnimationsManager *animationsManager;
 - (NSArray<id<RCTComponent>> *)_childrenToRemoveFromContainer:(id<RCTComponent>)container
                                                     atIndices:(NSArray<NSNumber *> *)atIndices;
 @end
 
-@implementation REAUIManager {
-  NSMutableDictionary<NSNumber *, NSMutableSet<id<RCTComponent>> *> *_toBeRemovedRegister;
-  NSMutableDictionary<NSNumber *, NSNumber *> *_parentMapper;
-  REAAnimationsManager *_animationsManager;
-  std::weak_ptr<reanimated::Scheduler> _scheduler;
+@implementation RCTUIManager (Reanimated)
+@dynamic animationsManager;
+- (void)setAnimationsManager:(REAAnimationsManager *)animationsManager
+{
+  objc_setAssociatedObject(self, @selector(animationsManager), animationsManager, OBJC_ASSOCIATION_RETAIN);
 }
+- (id)animationsManager
+{
+  return objc_getAssociatedObject(self, @selector(animationsManager));
+}
+@end
 
-+ (NSString *)moduleName
+@implementation REASwizzledUIManager
+
+- (instancetype)initWithUIManager:(RCTUIManager *)uiManager
+             withAnimationManager:(REAAnimationsManager *)animationsManager
 {
-  return NSStringFromClass([RCTUIManager class]);
+  if (self = [super init]) {
+    [uiManager setAnimationsManager:animationsManager];
+    [self swizzleMethods];
+  }
+  return self;
 }
 
-- (void)invalidate
+- (void)swizzleMethods
 {
-  [_animationsManager invalidate];
-  [super invalidate];
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    [self swizzleMethod:@selector(uiBlockWithLayoutUpdateForRootView:)
+               forClass:[RCTUIManager class]
+                   with:@selector(reanimated_uiBlockWithLayoutUpdateForRootView:)
+              fromClass:[self class]];
+    SEL manageChildrenOriginal = @selector
+        (_manageChildren:moveFromIndices:moveToIndices:addChildReactTags:addAtIndices:removeAtIndices:registry:);
+    SEL manageChildrenReanimated =
+        @selector(reanimated_manageChildren:
+                            moveFromIndices:moveToIndices:addChildReactTags:addAtIndices:removeAtIndices:registry:);
+    [self swizzleMethod:manageChildrenOriginal
+               forClass:[RCTUIManager class]
+                   with:manageChildrenReanimated
+              fromClass:[self class]];
+  });
 }
 
-- (void)setBridge:(RCTBridge *)bridge
+- (void)swizzleMethod:(SEL)originalSelector
+             forClass:(Class)originalClass
+                 with:(SEL)newSelector
+            fromClass:(Class)newClass
 {
-  if (!_blockSetter) {
-    _blockSetter = true;
-
-    self.bridge = bridge;
-    [super setValue:bridge forKey:@"_bridge"];
-    [self setValue:[bridge.uiManager valueForKey:@"_shadowViewRegistry"] forKey:@"_shadowViewRegistry"];
-    [self setValue:[bridge.uiManager valueForKey:@"_viewRegistry"] forKey:@"_viewRegistry"];
-    [self setValue:[bridge.uiManager valueForKey:@"_nativeIDRegistry"] forKey:@"_nativeIDRegistry"];
-    [self setValue:[bridge.uiManager valueForKey:@"_shadowViewsWithUpdatedProps"]
-            forKey:@"_shadowViewsWithUpdatedProps"];
-    [self setValue:[bridge.uiManager valueForKey:@"_shadowViewsWithUpdatedChildren"]
-            forKey:@"_shadowViewsWithUpdatedChildren"];
-    [self setValue:[bridge.uiManager valueForKey:@"_pendingUIBlocks"] forKey:@"_pendingUIBlocks"];
-    [self setValue:[bridge.uiManager valueForKey:@"_rootViewTags"] forKey:@"_rootViewTags"];
-    [self setValue:[bridge.uiManager valueForKey:@"_observerCoordinator"] forKey:@"_observerCoordinator"];
-    [self setValue:[bridge.uiManager valueForKey:@"_componentDataByName"] forKey:@"_componentDataByName"];
-
-    _blockSetter = false;
-  }
+  Method originalMethod = class_getInstanceMethod(originalClass, originalSelector);
+  Method newMethod = class_getInstanceMethod(newClass, newSelector);
+  IMP originalImplementation = method_getImplementation(originalMethod);
+  IMP newImplementation = method_getImplementation(newMethod);
+  class_replaceMethod(originalClass, newSelector, originalImplementation, method_getTypeEncoding(originalMethod));
+  class_replaceMethod(originalClass, originalSelector, newImplementation, method_getTypeEncoding(newMethod));
 }
 
-- (void)_manageChildren:(NSNumber *)containerTag
-        moveFromIndices:(NSArray<NSNumber *> *)moveFromIndices
-          moveToIndices:(NSArray<NSNumber *> *)moveToIndices
-      addChildReactTags:(NSArray<NSNumber *> *)addChildReactTags
-           addAtIndices:(NSArray<NSNumber *> *)addAtIndices
-        removeAtIndices:(NSArray<NSNumber *> *)removeAtIndices
-               registry:(NSMutableDictionary<NSNumber *, id<RCTComponent>> *)registry
+- (void)reanimated_manageChildren:(NSNumber *)containerTag
+                  moveFromIndices:(NSArray<NSNumber *> *)moveFromIndices
+                    moveToIndices:(NSArray<NSNumber *> *)moveToIndices
+                addChildReactTags:(NSArray<NSNumber *> *)addChildReactTags
+                     addAtIndices:(NSArray<NSNumber *> *)addAtIndices
+                  removeAtIndices:(NSArray<NSNumber *> *)removeAtIndices
+                         registry:(NSMutableDictionary<NSNumber *, id<RCTComponent>> *)registry
 {
   bool isLayoutAnimationEnabled = reanimated::FeaturesConfig::isLayoutAnimationEnabled();
   id<RCTComponent> container;
   NSArray<id<RCTComponent>> *permanentlyRemovedChildren;
   BOOL containerIsRootOfViewController = NO;
+  RCTUIManager *originalSelf = (RCTUIManager *)self;
   if (isLayoutAnimationEnabled) {
     container = registry[containerTag];
-    permanentlyRemovedChildren = [self _childrenToRemoveFromContainer:container atIndices:removeAtIndices];
+    permanentlyRemovedChildren = [originalSelf _childrenToRemoveFromContainer:container atIndices:removeAtIndices];
 
     if ([container isKindOfClass:[UIView class]]) {
       UIViewController *controller = ((UIView *)container).reactViewController;
@@ -97,22 +97,23 @@ - (void)_manageChildren:(NSNumber *)containerTag
     // of some view controller. In that case, we skip running exiting animations
     // in its children, to prevent issues with RN Screens.
     if (containerIsRootOfViewController) {
-      NSArray<id<RCTComponent>> *permanentlyRemovedChildren = [self _childrenToRemoveFromContainer:container
-                                                                                         atIndices:removeAtIndices];
+      NSArray<id<RCTComponent>> *permanentlyRemovedChildren =
+          [originalSelf _childrenToRemoveFromContainer:container atIndices:removeAtIndices];
       for (UIView *view in permanentlyRemovedChildren) {
-        [_animationsManager endAnimationsRecursive:view];
+        [originalSelf.animationsManager endAnimationsRecursive:view];
       }
-      [_animationsManager removeAnimationsFromSubtree:(UIView *)container];
+      [originalSelf.animationsManager removeAnimationsFromSubtree:(UIView *)container];
     }
   }
 
-  [super _manageChildren:containerTag
-         moveFromIndices:moveFromIndices
-           moveToIndices:moveToIndices
-       addChildReactTags:addChildReactTags
-            addAtIndices:addAtIndices
-         removeAtIndices:removeAtIndices
-                registry:registry];
+  // call original method
+  [self reanimated_manageChildren:containerTag
+                  moveFromIndices:moveFromIndices
+                    moveToIndices:moveToIndices
+                addChildReactTags:addChildReactTags
+                     addAtIndices:addAtIndices
+                  removeAtIndices:removeAtIndices
+                         registry:registry];
 
   if (!isLayoutAnimationEnabled) {
     return;
@@ -132,27 +133,18 @@ - (void)_manageChildren:(NSNumber *)containerTag
         return [(NSNumber *)obj1[0] compare:(NSNumber *)obj2[0]];
       }];
 
-  [_animationsManager reattachAnimatedChildren:permanentlyRemovedChildren
-                                   toContainer:container
-                                     atIndices:removeAtIndices];
-}
-
-- (void)callAnimationForTree:(UIView *)view parentTag:(NSNumber *)parentTag
-{
-  _parentMapper[view.reactTag] = parentTag;
-
-  for (UIView *subView in view.reactSubviews) {
-    [self callAnimationForTree:subView parentTag:view.reactTag];
-  }
+  [originalSelf.animationsManager reattachAnimatedChildren:permanentlyRemovedChildren
+                                               toContainer:container
+                                                 atIndices:removeAtIndices];
 }
 
-// Overrided https://github.com/facebook/react-native/blob/v0.65.0/React/Modules/RCTUIManager.m#L530
-- (RCTViewManagerUIBlock)uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *)rootShadowView
+- (RCTViewManagerUIBlock)reanimated_uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *)rootShadowView
 {
   if (!reanimated::FeaturesConfig::isLayoutAnimationEnabled()) {
-    return [super uiBlockWithLayoutUpdateForRootView:rootShadowView];
+    return [self reanimated_uiBlockWithLayoutUpdateForRootView:rootShadowView];
   }
 
+  RCTUIManager *originalSelf = (RCTUIManager *)self;
   NSHashTable<RCTShadowView *> *affectedShadowViews = [NSHashTable weakObjectsHashTable];
   [rootShadowView layoutWithAffectedShadowViews:affectedShadowViews];
 
@@ -266,7 +258,8 @@ - (RCTViewManagerUIBlock)uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *
       }
 
       // Reanimated changes /start
-      REASnapshot *snapshotBefore = isNew ? nil : [self->_animationsManager prepareSnapshotBeforeMountForView:view];
+      REASnapshot *snapshotBefore =
+          isNew ? nil : [originalSelf.animationsManager prepareSnapshotBeforeMountForView:view];
       snapshotsBefore[reactTag] = snapshotBefore;
       // Reanimated changes /end
 
@@ -327,7 +320,7 @@ - (RCTViewManagerUIBlock)uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *
       REASnapshot *snapshotBefore = snapshotsBefore[reactTag];
 
       if (isNew || snapshotBefore != nil) {
-        [self->_animationsManager viewDidMount:view withBeforeSnapshot:snapshotBefore withNewFrame:frame];
+        [originalSelf.animationsManager viewDidMount:view withBeforeSnapshot:snapshotBefore withNewFrame:frame];
       }
     }
 
@@ -336,52 +329,9 @@ - (RCTViewManagerUIBlock)uiBlockWithLayoutUpdateForRootView:(RCTRootShadowView *
     // private field
     [uiManager setNextLayoutAnimationGroup:nil];
 
-    [self->_animationsManager viewsDidLayout];
+    [originalSelf.animationsManager viewsDidLayout];
     // Reanimated changes /end
   };
 }
 
-- (Class)class
-{
-  return [RCTUIManager class];
-}
-
-+ (Class)class
-{
-  return [RCTUIManager class];
-}
-
-- (void)setUp:(REAAnimationsManager *)animationsManager
-{
-  _animationsManager = animationsManager;
-  _toBeRemovedRegister = [[NSMutableDictionary<NSNumber *, NSMutableSet<id<RCTComponent>> *> alloc] init];
-  _parentMapper = [[NSMutableDictionary<NSNumber *, NSNumber *> alloc] init];
-}
-
-- (void)unregisterView:(id<RCTComponent>)view
-{
-  NSNumber *tag = _parentMapper[view.reactTag];
-  if (tag == nil) {
-    return;
-  }
-
-  [_toBeRemovedRegister[tag] removeObject:view];
-  if (_toBeRemovedRegister[tag].count == 0) {
-    [_toBeRemovedRegister removeObjectForKey:tag];
-  }
-  NSMutableDictionary<NSNumber *, id<RCTComponent>> *viewRegistry = [self valueForKey:@"_viewRegistry"];
-  [view.reactSuperview removeReactSubview:view];
-  id<RCTComponent> parentView = viewRegistry[tag];
-  @try {
-    [parentView removeReactSubview:view];
-  } @catch (id anException) {
-  }
-#if __has_include(<RNScreens/RNSScreen.h>)
-  if ([view isKindOfClass:[RNSScreenView class]]) {
-    [parentView didUpdateReactSubviews];
-  }
-#endif
-  [viewRegistry removeObjectForKey:view.reactTag];
-}
-
 @end

ios/REAModule.h

@@ -5,11 +5,13 @@
 #import <React/RCTUIManagerObserverCoordinator.h>
 #import <React/RCTUIManagerUtils.h>
 
+#import <RNReanimated/REAAnimationsManager.h>
 #import <RNReanimated/REANodesManager.h>
 
 @interface REAModule : RCTEventEmitter <RCTBridgeModule, RCTEventDispatcherObserver, RCTUIManagerObserver>
 
 @property (nonatomic, readonly) REANodesManager *nodesManager;
+@property REAAnimationsManager *animationsManager;
 
 #ifdef RCT_NEW_ARCH_ENABLED
 - (void)installReanimatedAfterReload;

ios/REAModule.mm

@@ -195,6 +195,7 @@ - (void)setBridge:(RCTBridge *)bridge
   _operations = [NSMutableArray new];
 
   [bridge.uiManager.observerCoordinator addObserver:self];
+  _animationsManager = [[REAAnimationsManager alloc] initWithUIManager:bridge.uiManager];
 }
 
 #pragma mark-- Batch handling