react-native-screens 3.32.0 → 3.33.0

package/common/cpp/react/renderer/components/rnscreens/RNSScreenShadowNode.cpp

@@ -3,13 +3,63 @@
 namespace facebook {
 namespace react {
 
+namespace yoga = facebook::yoga;
+using namespace rnscreens;
+
 extern const char RNSScreenComponentName[] = "RNSScreen";
 
-Point RNSScreenShadowNode::getContentOriginOffset() const {
+Point RNSScreenShadowNode::getContentOriginOffset(
+    bool /*includeTransform*/) const {
   auto stateData = getStateData();
   auto contentOffset = stateData.contentOffset;
   return {contentOffset.x, contentOffset.y};
 }
 
+void RNSScreenShadowNode::layout(facebook::react::LayoutContext layoutContext) {
+  YogaLayoutableShadowNode::layout(layoutContext);
+
+#ifdef ANDROID
+  applyFrameCorrections();
+#endif // ANDROID
+}
+
+#ifdef ANDROID
+void RNSScreenShadowNode::applyFrameCorrections() {
+  ensureUnsealed();
+
+  // On the very first layout we want to correct both Y offset and frame size.
+  // On consecutive layouts we want to correct only Y offset, as the frame size
+  // is received from JVM side. This is done so if the Screen dimensions are
+  // read from ShadowTree (e.g by reanimated) they have chance of being
+  // accurate. On JVM side we do ignore this frame anyway, because
+  // ScreenStackViewManager.needsCustomLayoutForChildren() == true.
+  const auto &stateData = getStateData();
+  const float lastKnownHeaderHeight = stateData.getLastKnownHeaderHeight();
+  const auto &headerCorrectionModes = stateData.getHeaderCorrectionModes();
+  layoutMetrics_.frame.origin.y += lastKnownHeaderHeight *
+      headerCorrectionModes.check(
+          FrameCorrectionModes::Mode::FrameOriginCorrection);
+  layoutMetrics_.frame.size.height -= lastKnownHeaderHeight *
+      headerCorrectionModes.check(
+          FrameCorrectionModes::Mode::FrameHeightCorrection);
+}
+
+void RNSScreenShadowNode::setHeaderHeight(float headerHeight) {
+  getStateDataMutable().setHeaderHeight(headerHeight);
+}
+
+FrameCorrectionModes &RNSScreenShadowNode::getFrameCorrectionModes() {
+  return getStateDataMutable().getFrameCorrectionModes();
+}
+
+RNSScreenShadowNode::StateData &RNSScreenShadowNode::getStateDataMutable() {
+  // We assume that this method is called to mutate the data, so we ensure
+  // we're unsealed.
+  ensureUnsealed();
+  return const_cast<RNSScreenShadowNode::StateData &>(getStateData());
+}
+
+#endif // ANDROID
+
 } // namespace react
 } // namespace facebook

package/common/cpp/react/renderer/components/rnscreens/RNSScreenShadowNode.h

@@ -4,11 +4,15 @@
 #include <react/renderer/components/rnscreens/EventEmitters.h>
 #include <react/renderer/components/rnscreens/Props.h>
 #include <react/renderer/components/view/ConcreteViewShadowNode.h>
+#include <react/renderer/core/LayoutContext.h>
+#include "FrameCorrectionModes.h"
 #include "RNSScreenState.h"
 
 namespace facebook {
 namespace react {
 
+using namespace rnscreens;
+
 JSI_EXPORT extern const char RNSScreenComponentName[];
 
 class JSI_EXPORT RNSScreenShadowNode final : public ConcreteViewShadowNode<
@@ -18,8 +22,26 @@ class JSI_EXPORT RNSScreenShadowNode final : public ConcreteViewShadowNode<
                                                  RNSScreenState> {
  public:
   using ConcreteViewShadowNode::ConcreteViewShadowNode;
+  using StateData = ConcreteViewShadowNode::ConcreteStateData;
+
+#pragma mark - ShadowNode overrides
+
+  Point getContentOriginOffset(bool includeTransform) const override;
+
+  void layout(LayoutContext layoutContext) override;
+
+#pragma mark - Custom interface
+
+  void setHeaderHeight(float headerHeight);
+
+  FrameCorrectionModes &getFrameCorrectionModes();
+
+ private:
+#ifdef ANDROID
+  void applyFrameCorrections();
 
-  Point getContentOriginOffset() const override;
+  StateData &getStateDataMutable();
+#endif // ANDROID
 };
 
 } // namespace react

package/common/cpp/react/renderer/components/rnscreens/RNSScreenState.cpp

@@ -3,12 +3,32 @@
 namespace facebook {
 namespace react {
 
+using namespace rnscreens;
+
 #ifdef ANDROID
 folly::dynamic RNSScreenState::getDynamic() const {
   return folly::dynamic::object("frameWidth", frameSize.width)(
       "frameHeight", frameSize.height)("contentOffsetX", contentOffset.x)(
       "contentOffsetY", contentOffset.y);
 }
+
+void RNSScreenState::setHeaderHeight(float headerHeight) {
+  lastKnownHeaderHeight_ = headerHeight;
+}
+
+float RNSScreenState::getLastKnownHeaderHeight() const noexcept {
+  return lastKnownHeaderHeight_;
+}
+
+FrameCorrectionModes &RNSScreenState::getFrameCorrectionModes() noexcept {
+  return headerCorrectionModes_;
+}
+
+const FrameCorrectionModes &RNSScreenState::getHeaderCorrectionModes()
+    const noexcept {
+  return headerCorrectionModes_;
+}
+
 #endif
 
 } // namespace react

package/common/cpp/react/renderer/components/rnscreens/RNSScreenState.h

@@ -9,9 +9,13 @@
 #include <react/renderer/mapbuffer/MapBufferBuilder.h>
 #endif
 
+#include "FrameCorrectionModes.h"
+
 namespace facebook {
 namespace react {
 
+using namespace rnscreens;
+
 class JSI_EXPORT RNSScreenState final {
  public:
   using Shared = std::shared_ptr<const RNSScreenState>;
@@ -27,7 +31,9 @@ class JSI_EXPORT RNSScreenState final {
             (Float)data["frameHeight"].getDouble()}),
         contentOffset(Point{
             (Float)data["contentOffsetX"].getDouble(),
-            (Float)data["contentOffsetY"].getDouble()}){};
+            (Float)data["contentOffsetY"].getDouble()}),
+        lastKnownHeaderHeight_{previousState.lastKnownHeaderHeight_},
+        headerCorrectionModes_{previousState.headerCorrectionModes_} {};
 #endif
 
   const Size frameSize{};
@@ -39,8 +45,24 @@ class JSI_EXPORT RNSScreenState final {
     return MapBufferBuilder::EMPTY();
   };
 
+  void setHeaderHeight(float headerHeight);
+
+  float getLastKnownHeaderHeight() const noexcept;
+
+  FrameCorrectionModes &getFrameCorrectionModes() noexcept;
+
+  const FrameCorrectionModes &getHeaderCorrectionModes() const noexcept;
+
 #endif
 
+ private:
+#ifdef ANDROID
+  // Header height as measured on dummy layout
+  float lastKnownHeaderHeight_{0.f};
+
+  FrameCorrectionModes headerCorrectionModes_{};
+#endif // ANDROID
+
 #pragma mark - Getters
 };
 

package/common/cpp/react/renderer/components/rnscreens/utils/RectUtil.h

@@ -0,0 +1,36 @@
+#pragma once
+
+#include <react/renderer/graphics/Float.h>
+#include <react/renderer/graphics/Size.h>
+#include <cmath>
+#include <concepts>
+
+namespace rnscreens {
+
+namespace react = facebook::react;
+
+template <typename T>
+  requires std::is_floating_point_v<T>
+inline constexpr bool equalWithRespectToEps(const T a, const T b, const T eps) {
+  return std::abs(a - b) <= eps;
+}
+
+/**
+ * Compares given two frame sizes with respect to the epsilon.
+ *
+ * @param first first frame size
+ * @param second second frame size
+ * @param eps comparison precision, defaults to 0.01, which should ensure that
+ * precision of comparison is under 1px
+ * @return whether the frame dimensions are the same with respect to given
+ * epsilon
+ */
+inline constexpr bool checkFrameSizesEqualWithEps(
+    const react::Size &first,
+    const react::Size &second,
+    const react::Float eps = 0.01) {
+  return equalWithRespectToEps(first.width, second.width, eps) &&
+      equalWithRespectToEps(first.height, second.height, eps);
+}
+
+} // namespace rnscreens

package/cpp/RNSScreenRemovalListener.cpp

@@ -0,0 +1,25 @@
+#include "RNSScreenRemovalListener.h"
+#include <react/renderer/mounting/ShadowViewMutation.h>
+using namespace facebook::react;
+
+std::optional<MountingTransaction> RNSScreenRemovalListener::pullTransaction(
+    SurfaceId surfaceId,
+    MountingTransaction::Number transactionNumber,
+    const TransactionTelemetry &telemetry,
+    ShadowViewMutationList mutations) const {
+  for (const ShadowViewMutation &mutation : mutations) {
+    if (mutation.type == ShadowViewMutation::Type::Remove &&
+        mutation.oldChildShadowView.componentName != nullptr &&
+        strcmp(mutation.parentShadowView.componentName, "RNSScreenStack") ==
+            0) {
+      listenerFunction_(mutation.oldChildShadowView.tag);
+    }
+  }
+
+  return MountingTransaction{
+      surfaceId, transactionNumber, std::move(mutations), telemetry};
+}
+
+bool RNSScreenRemovalListener::shouldOverridePullTransaction() const {
+  return true;
+}

package/cpp/RNSScreenRemovalListener.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <react/renderer/componentregistry/ComponentDescriptorFactory.h>
+#include <react/renderer/mounting/MountingOverrideDelegate.h>
+#include <react/renderer/mounting/ShadowView.h>
+
+using namespace facebook::react;
+
+struct RNSScreenRemovalListener : public MountingOverrideDelegate {
+  std::function<void(int)> listenerFunction_;
+  RNSScreenRemovalListener(std::function<void(int)> &&listenerFunction_)
+      : listenerFunction_(std::move(listenerFunction_)) {}
+
+  bool shouldOverridePullTransaction() const override;
+  std::optional<MountingTransaction> pullTransaction(
+      SurfaceId surfaceId,
+      MountingTransaction::Number number,
+      const TransactionTelemetry &telemetry,
+      ShadowViewMutationList mutations) const override;
+};

package/ios/RNSConvert.h

@@ -1,4 +1,5 @@
 #ifdef RCT_NEW_ARCH_ENABLED
+#import <UIKit/UIKit.h>
 #import <react/renderer/components/rnscreens/Props.h>
 #import "RNSEnums.h"
 

package/ios/RNSModalScreen.mm

@@ -7,6 +7,28 @@
 
 @implementation RNSModalScreen
 
+// When using UIModalPresentationStyleFullScreen the whole view hierarchy mounted under primary `UITransitionView` is
+// removed, including React's root view, which observes for trait collection changes & sends it to `Appearance` module
+// via system notification centre. To workaround this detached-root-view-situation we emit the event to React's
+// `Appearance` module ourselves. For the RCTRootView observer, visit
+// https://github.com/facebook/react-native/blob/d3e0430deac573fd44792e6005d5de20e9ad2797/packages/react-native/React/Base/RCTRootView.m#L362
+// For more information, see https://github.com/software-mansion/react-native-screens/pull/2211.
+- (void)traitCollectionDidChange:(UITraitCollection *)previousTraitCollection
+{
+  [super traitCollectionDidChange:previousTraitCollection];
+  if (RCTSharedApplication().applicationState == UIApplicationStateBackground ||
+      self.stackPresentation != RNSScreenStackPresentationFullScreenModal) {
+    return;
+  }
+
+  [[NSNotificationCenter defaultCenter]
+      postNotificationName:RCTUserInterfaceStyleDidChangeNotification
+                    object:self
+                  userInfo:@{
+                    RCTUserInterfaceStyleDidChangeNotificationTraitCollectionKey : self.traitCollection,
+                  }];
+}
+
 #ifdef RCT_NEW_ARCH_ENABLED
 + (react::ComponentDescriptorProvider)componentDescriptorProvider
 {

package/ios/RNSModule.mm

@@ -116,7 +116,7 @@ RCT_EXPORT_MODULE()
 {
   RNSScreenStackView *view = [self getScreenStackView:stackTag];
   if (view != nil && ![view isKindOfClass:[RNSScreenStackView class]]) {
-    RCTLogError(@"Invalid view type, expecting RNSScreenStackView, got: %@", view);
+    RCTLogError(@"[RNScreens] Invalid view type, expecting RNSScreenStackView, got: %@", view);
     return nil;
   }
   return view;

package/ios/RNSScreen.h

@@ -40,7 +40,7 @@ namespace react = facebook::react;
 - (void)notifyFinishTransitioning;
 - (RNSScreenView *)screenView;
 #ifdef RCT_NEW_ARCH_ENABLED
-- (void)setViewToSnapshot:(UIView *)snapshot;
+- (void)setViewToSnapshot;
 - (CGFloat)calculateHeaderHeightIsModal:(BOOL)isModal;
 #endif
 
@@ -56,6 +56,7 @@ namespace react = facebook::react;
 #endif
 
 @property (nonatomic) BOOL fullScreenSwipeEnabled;
+@property (nonatomic) BOOL fullScreenSwipeShadowEnabled;
 @property (nonatomic) BOOL gestureEnabled;
 @property (nonatomic) BOOL hasStatusBarHiddenSet;
 @property (nonatomic) BOOL hasStatusBarStyleSet;

package/ios/RNSScreen.mm

@@ -25,6 +25,8 @@
 #import "RNSScreenStack.h"
 #import "RNSScreenStackHeaderConfig.h"
 
+#import "UIView+RNSUtility.h"
+
 #ifdef RCT_NEW_ARCH_ENABLED
 namespace react = facebook::react;
 #endif // RCT_NEW_ARCH_ENABLED
@@ -414,7 +416,7 @@ namespace react = facebook::react;
       [[RNSHeaderHeightChangeEvent alloc] initWithEventName:@"onHeaderHeightChange"
                                                    reactTag:[NSNumber numberWithInt:self.tag]
                                                headerHeight:headerHeight];
-  [[RCTBridge currentBridge].eventDispatcher notifyObserversOfEvent:event];
+  [self postNotificationForEventDispatcherObserversWithEvent:event];
 #else
   if (self.onHeaderHeightChange) {
     self.onHeaderHeightChange(@{
@@ -473,22 +475,13 @@ namespace react = facebook::react;
   }
 }
 
-#ifdef RCT_NEW_ARCH_ENABLED
-- (RCTSurfaceTouchHandler *)touchHandler
-#else
-- (RCTTouchHandler *)touchHandler
-#endif
+- (nullable RNS_TOUCH_HANDLER_ARCH_TYPE *)touchHandler
 {
   if (_touchHandler != nil) {
     return _touchHandler;
   }
-  UIView *parent = [self superview];
-  while (parent != nil && ![parent respondsToSelector:@selector(touchHandler)])
-    parent = parent.superview;
-  if (parent != nil) {
-    return [parent performSelector:@selector(touchHandler)];
-  }
-  return nil;
+
+  return [self rnscreens_findTouchHandlerInAncestorChain];
 }
 
 - (void)notifyFinishTransitioning
@@ -509,7 +502,7 @@ namespace react = facebook::react;
                                                                    progress:progress
                                                                     closing:closing
                                                                goingForward:goingForward];
-  [[RCTBridge currentBridge].eventDispatcher notifyObserversOfEvent:event];
+  [self postNotificationForEventDispatcherObserversWithEvent:event];
 #else
   if (self.onTransitionProgress) {
     self.onTransitionProgress(@{
@@ -667,6 +660,14 @@ namespace react = facebook::react;
 #pragma mark - Fabric specific
 #ifdef RCT_NEW_ARCH_ENABLED
 
+- (void)postNotificationForEventDispatcherObserversWithEvent:(NSObject<RCTEvent> *)event
+{
+  NSDictionary *userInfo = [NSDictionary dictionaryWithObjectsAndKeys:event, @"event", nil];
+  [[NSNotificationCenter defaultCenter] postNotificationName:@"RCTNotifyEventDispatcherObserversOfEvent_DEPRECATED"
+                                                      object:nil
+                                                    userInfo:userInfo];
+}
+
 - (BOOL)hasHeaderConfig
 {
   return _config != nil;
@@ -710,6 +711,8 @@ namespace react = facebook::react;
 
   [self setFullScreenSwipeEnabled:newScreenProps.fullScreenSwipeEnabled];
 
+  [self setFullScreenSwipeShadowEnabled:newScreenProps.fullScreenSwipeShadowEnabled];
+
   [self setGestureEnabled:newScreenProps.gestureEnabled];
 
   [self setTransitionDuration:[NSNumber numberWithInt:newScreenProps.transitionDuration]];
@@ -793,7 +796,7 @@ namespace react = facebook::react;
   _newLayoutMetrics = layoutMetrics;
   _oldLayoutMetrics = oldLayoutMetrics;
   UIViewController *parentVC = self.reactViewController.parentViewController;
-  if (parentVC != nil && ![parentVC isKindOfClass:[RNSNavigationController class]]) {
+  if (parentVC == nil || ![parentVC isKindOfClass:[RNSNavigationController class]]) {
     [super updateLayoutMetrics:layoutMetrics oldLayoutMetrics:oldLayoutMetrics];
   }
   // when screen is mounted under RNSNavigationController it's size is controller
@@ -1388,12 +1391,16 @@ Class<RCTComponentViewProtocol> RNSScreenCls(void)
 #ifdef RCT_NEW_ARCH_ENABLED
 #pragma mark - Fabric specific
 
-- (void)setViewToSnapshot:(UIView *)snapshot
+- (void)setViewToSnapshot
 {
   UIView *superView = self.view.superview;
-  [self.view removeFromSuperview];
-  self.view = snapshot;
-  [superView addSubview:self.view];
+  // if we dismissed the view natively, it will already be detached from view hierarchy
+  if (self.view.window != nil) {
+    UIView *snapshot = [self.view snapshotViewAfterScreenUpdates:NO];
+    [self.view removeFromSuperview];
+    self.view = snapshot;
+    [superView addSubview:snapshot];
+  }
 }
 
 #else
@@ -1434,6 +1441,7 @@ RCT_EXPORT_MODULE()
 RCT_REMAP_VIEW_PROPERTY(activityState, activityStateOrNil, NSNumber)
 RCT_EXPORT_VIEW_PROPERTY(customAnimationOnSwipe, BOOL);
 RCT_EXPORT_VIEW_PROPERTY(fullScreenSwipeEnabled, BOOL);
+RCT_EXPORT_VIEW_PROPERTY(fullScreenSwipeShadowEnabled, BOOL);
 RCT_EXPORT_VIEW_PROPERTY(gestureEnabled, BOOL)
 RCT_EXPORT_VIEW_PROPERTY(gestureResponseDistance, NSDictionary)
 RCT_EXPORT_VIEW_PROPERTY(hideKeyboardOnSwipe, BOOL)

package/ios/RNSScreenStack.mm

@@ -26,6 +26,8 @@
 #import "RNSScreenStackHeaderConfig.h"
 #import "RNSScreenWindowTraits.h"
 
+#import "UIView+RNSUtility.h"
+
 #ifdef RCT_NEW_ARCH_ENABLED
 namespace react = facebook::react;
 #endif // RCT_NEW_ARCH_ENABLED
@@ -116,12 +118,8 @@ namespace react = facebook::react;
   BOOL _invalidated;
   BOOL _isFullWidthSwiping;
   UIPercentDrivenInteractiveTransition *_interactionController;
-  BOOL _hasLayout;
   __weak RNSScreenStackManager *_manager;
   BOOL _updateScheduled;
-#ifdef RCT_NEW_ARCH_ENABLED
-  UIView *_snapshot;
-#endif
 }
 
 #ifdef RCT_NEW_ARCH_ENABLED
@@ -140,7 +138,6 @@ namespace react = facebook::react;
 - (instancetype)initWithManager:(RNSScreenStackManager *)manager
 {
   if (self = [super init]) {
-    _hasLayout = NO;
     _invalidated = NO;
     _manager = manager;
     [self initCommonProps];
@@ -270,18 +267,9 @@ namespace react = facebook::react;
 - (void)maybeAddToParentAndUpdateContainer
 {
   BOOL wasScreenMounted = _controller.parentViewController != nil;
-#ifdef RCT_NEW_ARCH_ENABLED
-  BOOL isScreenReadyForShowing = self.window;
-#else
-  BOOL isScreenReadyForShowing = self.window && _hasLayout;
-#endif
-  if (!isScreenReadyForShowing && !wasScreenMounted) {
-    // We wait with adding to parent controller until the stack is mounted and has its initial
-    // layout done.
-    // If we add it before layout, some of the items (specifically items from the navigation bar),
-    // won't be able to position properly. Also the position and size of such items, even if it
-    // happens to change, won't be properly updated (this is perhaps some internal issue of UIKit).
-    // If we add it when window is not attached, some of the view transitions will be bloced (i.e.
+  if (!self.window && !wasScreenMounted) {
+    // We wait with adding to parent controller until the stack is mounted.
+    // If we add it when window is not attached, some of the view transitions will be blocked (i.e.
     // modal transitions) and the internal view controler's state will get out of sync with what's
     // on screen without us knowing.
     return;
@@ -735,43 +723,8 @@ namespace react = facebook::react;
   // item is close to an edge and we start pulling from edge we want the Touchable to be cancelled.
   // Without the below code the Touchable will remain active (highlighted) for the duration of back
   // gesture and onPress may fire when we release the finger.
-#ifdef RCT_NEW_ARCH_ENABLED
-  // On Fabric there is no view that exposes touchHandler above us in the view hierarchy, however it is still
-  // utilised. `RCTSurfaceView` should be present above us, which hosts `RCTFabricSurface` instance, which in turn
-  // hosts `RCTSurfaceTouchHandler` as a private field. When initialised, `RCTSurfaceTouchHandler` is attached to the
-  // surface view as a gestureRecognizer <- and this is where we can lay our hands on it.
-  UIView *parent = _controller.view;
-  while (parent != nil && ![parent isKindOfClass:RCTSurfaceView.class]) {
-    parent = parent.superview;
-  }
 
-  // This could be possible in modal context
-  if (parent == nil) {
-    return;
-  }
-
-  RCTSurfaceTouchHandler *touchHandler = nil;
-  // Experimentation shows that RCTSurfaceTouchHandler is the only gestureRecognizer registered here,
-  // so we should not be afraid of any performance hit here.
-  for (UIGestureRecognizer *recognizer in parent.gestureRecognizers) {
-    if ([recognizer isKindOfClass:RCTSurfaceTouchHandler.class]) {
-      touchHandler = static_cast<RCTSurfaceTouchHandler *>(recognizer);
-    }
-  }
-
-  [touchHandler rnscreens_cancelTouches];
-#else
-  // On Paper we can access touchHandler hosted by `RCTRootContentView` which should be above ScreenStack
-  // in view hierarchy.
-  UIView *parent = _controller.view;
-  while (parent != nil && ![parent respondsToSelector:@selector(touchHandler)]) {
-    parent = parent.superview;
-  }
-  if (parent != nil) {
-    RCTTouchHandler *touchHandler = [parent performSelector:@selector(touchHandler)];
-    [touchHandler rnscreens_cancelTouches];
-  }
-#endif // RCT_NEW_ARCH_ENABLED
+  [[self rnscreens_findTouchHandlerInAncestorChain] rnscreens_cancelTouches];
 }
 
 - (BOOL)gestureRecognizerShouldBegin:(UIGestureRecognizer *)gestureRecognizer
@@ -905,6 +858,7 @@ namespace react = facebook::react;
         [_interactionController cancelInteractiveTransition];
       }
       _interactionController = nil;
+      _isFullWidthSwiping = NO;
     }
     default: {
       break;
@@ -1109,7 +1063,6 @@ namespace react = facebook::react;
   // set yet, however the layout call is already enqueued on ui thread. Enqueuing update call on the
   // ui queue will guarantee that the update will run after layout.
   dispatch_async(dispatch_get_main_queue(), ^{
-    self->_hasLayout = YES;
     [self maybeAddToParentAndUpdateContainer];
   });
 }
@@ -1161,14 +1114,14 @@ namespace react = facebook::react;
 
   [_reactSubviews insertObject:(RNSScreenView *)childComponentView atIndex:index];
   ((RNSScreenView *)childComponentView).reactSuperview = self;
-  dispatch_async(dispatch_get_main_queue(), ^{
-    [self maybeAddToParentAndUpdateContainer];
-  });
+  // Container update is done after all mount operations are executed in
+  // `- [RNSScreenStackView mountingTransactionDidMount: withSurfaceTelemetry:]`
 }
 
 - (void)unmountChildComponentView:(UIView<RCTComponentViewProtocol> *)childComponentView index:(NSInteger)index
 {
   RNSScreenView *screenChildComponent = (RNSScreenView *)childComponentView;
+
   // We should only do a snapshot of a screen that is on the top.
   // We also check `_presentedModals` since if you push 2 modals, second one is not a "child" of _controller.
   // Also, when dissmised with a gesture, the screen already is not under the window, so we don't need to apply
@@ -1176,7 +1129,7 @@ namespace react = facebook::react;
   if (screenChildComponent.window != nil &&
       ((screenChildComponent == _controller.visibleViewController.view && _presentedModals.count < 2) ||
        screenChildComponent == [_presentedModals.lastObject view])) {
-    [screenChildComponent.controller setViewToSnapshot:_snapshot];
+    [screenChildComponent.controller setViewToSnapshot];
   }
 
   RCTAssert(
@@ -1196,28 +1149,22 @@ namespace react = facebook::react;
   screenChildComponent.reactSuperview = nil;
   [_reactSubviews removeObject:screenChildComponent];
   [screenChildComponent removeFromSuperview];
-  dispatch_async(dispatch_get_main_queue(), ^{
-    [self maybeAddToParentAndUpdateContainer];
-  });
 }
 
-- (void)takeSnapshot
+- (void)mountingTransactionDidMount:(const facebook::react::MountingTransaction &)transaction
+               withSurfaceTelemetry:(const facebook::react::SurfaceTelemetry &)surfaceTelemetry
 {
-  if (_presentedModals.count < 2) {
-    _snapshot = [_controller.visibleViewController.view snapshotViewAfterScreenUpdates:NO];
-  } else {
-    _snapshot = [[_presentedModals.lastObject view] snapshotViewAfterScreenUpdates:NO];
-  }
-}
-
-- (void)mountingTransactionWillMount:(react::MountingTransaction const &)transaction
-                withSurfaceTelemetry:(react::SurfaceTelemetry const &)surfaceTelemetry
-{
-  for (auto &mutation : transaction.getMutations()) {
-    if (mutation.type == react::ShadowViewMutation::Type::Remove && mutation.parentShadowView.componentName != nil &&
-        strcmp(mutation.parentShadowView.componentName, "RNSScreenStack") == 0) {
-      [self takeSnapshot];
-      return;
+  for (const auto &mutation : transaction.getMutations()) {
+    if (mutation.parentShadowView.tag == self.tag &&
+        (mutation.type == react::ShadowViewMutation::Type::Insert ||
+         mutation.type == react::ShadowViewMutation::Type::Remove)) {
+      // we need to wait until children have their layout set. At this point they don't have the layout
+      // set yet, however the layout call is already enqueued on ui thread. Enqueuing update call on the
+      // ui queue will guarantee that the update will run after layout.
+      dispatch_async(dispatch_get_main_queue(), ^{
+        [self maybeAddToParentAndUpdateContainer];
+      });
+      break;
     }
   }
 }