fix(iOS): header snapshots not working

apps/src/tests/Test556.tsx

@@ -1,13 +1,17 @@
 import * as React from 'react';
 import { Button, StyleSheet, Text, View } from 'react-native';
 import { NavigationContainer, ParamListBase } from '@react-navigation/native';
-import { NativeStackNavigationProp, createNativeStackNavigator } from '@react-navigation/native-stack';
+import {
+  NativeStackNavigationProp,
+  createNativeStackNavigator,
+} from '@react-navigation/native-stack';
+import { Square } from '../shared';
 
 const Stack = createNativeStackNavigator();
 
 type ScreenBaseProps = {
   navigation: NativeStackNavigationProp<ParamListBase>;
-}
+};
 
 export default function App() {
   return (
@@ -16,30 +20,48 @@ export default function App() {
         screenOptions={{
           animation: 'fade',
         }}>
-        <Stack.Screen name="First" component={First} options={{
-          headerTitle: () => (
-            <View style={[styles.container, { backgroundColor: 'goldenrod' }]}>
-              <Text>Hello there!</Text>
-            </View>
-          ),
-          headerRight: () => (
-            <View style={[styles.container, { backgroundColor: 'lightblue' }]}>
-              <Text>Right-1</Text>
-            </View>
-          ),
-        }} />
-        <Stack.Screen name="Second" component={Second} options={{
-          headerTitle: () => (
-            <View style={[styles.container, { backgroundColor: 'mediumseagreen' }]}>
-              <Text>General Kenobi</Text>
-            </View>
-          ),
-          headerRight: () => (
-            <View style={[styles.container, { backgroundColor: 'mediumvioletred' }]}>
-              <Text>Right-2</Text>
-            </View>
-          ),
-        }} />
+        <Stack.Screen
+          name="First"
+          component={First}
+          options={{
+            headerTitle: () => (
+              <View
+                style={[styles.container, { backgroundColor: 'goldenrod' }]}>
+                <Text>Hello there!</Text>
+              </View>
+            ),
+            headerRight: () => (
+              <View
+                style={[styles.container, { backgroundColor: 'lightblue' }]}>
+                <Text>Right-1</Text>
+              </View>
+            ),
+          }}
+        />
+        <Stack.Screen
+          name="Second"
+          component={Second}
+          options={{
+            headerTitle: () => (
+              <View
+                style={[
+                  styles.container,
+                  { backgroundColor: 'mediumseagreen' },
+                ]}>
+                <Text>General Kenobi</Text>
+              </View>
+            ),
+            headerRight: () => (
+              <View
+                style={[
+                  styles.container,
+                  { backgroundColor: 'mediumvioletred' },
+                ]}>
+                <Text>Right-2</Text>
+              </View>
+            ),
+          }}
+        />
       </Stack.Navigator>
     </NavigationContainer>
   );
@@ -55,11 +77,34 @@ function First({ navigation }: ScreenBaseProps) {
 }
 
 function Second({ navigation }: ScreenBaseProps) {
+  const [backButtonVisible, setBackButtonVisible] = React.useState(true);
+
   return (
-    <Button
-      title="Tap me for first screen"
-      onPress={() => navigation.popTo('First')}
-    />
+    <>
+      <Button
+        title="Toggle left subview"
+        onPress={() => {
+          setBackButtonVisible(prev => !prev);
+          navigation.setOptions({
+            headerLeft: backButtonVisible
+              ? () => (
+                  <View
+                    style={[
+                      styles.container,
+                      { backgroundColor: 'mediumblue' },
+                    ]}>
+                    <Text>Left</Text>
+                  </View>
+                )
+              : undefined,
+          });
+        }}
+      />
+      <Button
+        title="Tap me for first screen"
+        onPress={() => navigation.popTo('First')}
+      />
+    </>
   );
 }
 

ios/RNSScreen.h

@@ -100,6 +100,8 @@ namespace react = facebook::react;
 @property (nonatomic) react::LayoutMetrics newLayoutMetrics;
 @property (weak, nonatomic) RNSScreenStackHeaderConfig *config;
 @property (nonatomic, readonly) BOOL hasHeaderConfig;
+@property (nonatomic, readonly, getter=isMarkedForUnmountInCurrentTransaction)
+    BOOL markedForUnmountInCurrentTransaction;
 #else
 @property (nonatomic, copy) RCTDirectEventBlock onAppear;
 @property (nonatomic, copy) RCTDirectEventBlock onDisappear;
@@ -124,6 +126,10 @@ namespace react = facebook::react;
 - (void)updateBounds;
 - (void)notifyDismissedWithCount:(int)dismissCount;
 - (instancetype)initWithFrame:(CGRect)frame;
+/// Tell `Screen` that it will be unmounted in next transaction.
+/// The component needs this so that we can later decide whether to
+/// replace it with snapshot or not.
+- (void)willBeUnmountedInUpcomingTransaction;
 #else
 - (instancetype)initWithBridge:(RCTBridge *)bridge;
 #endif // RCT_NEW_ARCH_ENABLED

ios/RNSScreen.mm

@@ -123,6 +123,9 @@ - (void)initCommonProps
 #endif // !TARGET_OS_TV
   _sheetsScrollView = nil;
   _didSetSheetAllowedDetentsOnController = NO;
+#ifdef RCT_NEW_ARCH_ENABLED
+  _markedForUnmountInCurrentTransaction = NO;
+#endif // RCT_NEW_ARCH_ENABLED
 }
 
 - (UIViewController *)reactViewController
@@ -1082,6 +1085,11 @@ - (BOOL)hasHeaderConfig
   return _config != nil;
 }
 
+- (void)willBeUnmountedInUpcomingTransaction
+{
+  _markedForUnmountInCurrentTransaction = YES;
+}
+
 + (react::ComponentDescriptorProvider)componentDescriptorProvider
 {
   return react::concreteComponentDescriptorProvider<react::RNSScreenComponentDescriptor>();

ios/RNSScreenStack.mm

@@ -1208,6 +1208,7 @@ - (void)mountingTransactionWillMount:(const facebook::react::MountingTransaction
     if (mutation.type == react::ShadowViewMutation::Delete) {
       RNSScreenView *_Nullable toBeRemovedChild = [self childScreenForTag:mutation.oldChildShadowView.tag];
       if (toBeRemovedChild != nil) {
+        [toBeRemovedChild willBeUnmountedInUpcomingTransaction];
         _toBeDeletedScreens.push_back(toBeRemovedChild);
       }
     }

ios/RNSScreenStackHeaderConfig.mm

@@ -823,12 +823,17 @@ - (void)mountChildComponentView:(UIView<RCTComponentViewProtocol> *)childCompone
 
 - (void)unmountChildComponentView:(UIView<RCTComponentViewProtocol> *)childComponentView index:(NSInteger)index
 {
-  // For explanation of why we can make a snapshot here despite the fact that our children are already
-  // unmounted see https://github.com/software-mansion/react-native-screens/pull/2261
-  [self replaceNavigationBarViewsWithSnapshotOfSubview:(RNSScreenStackHeaderSubview *)childComponentView];
+  BOOL isGoingToBeRemoved = _screenView.isMarkedForUnmountInCurrentTransaction;
+  if (isGoingToBeRemoved) {
+    // For explanation of why we can make a snapshot here despite the fact that our children are already
+    // unmounted see https://github.com/software-mansion/react-native-screens/pull/2261
+    [self replaceNavigationBarViewsWithSnapshotOfSubview:(RNSScreenStackHeaderSubview *)childComponentView];
+  }
   [_reactSubviews removeObject:(RNSScreenStackHeaderSubview *)childComponentView];
   [childComponentView removeFromSuperview];
-  [self updateViewControllerIfNeeded];
+  if (!isGoingToBeRemoved) {
+    [self updateViewControllerIfNeeded];
+  }
 }
 
 - (void)replaceNavigationBarViewsWithSnapshotOfSubview:(RNSScreenStackHeaderSubview *)childComponentView