From d4c736b26ff86dd08cec4e0485b031e915e9b3c1 Mon Sep 17 00:00:00 2001
From: Tymoteusz Boba <Tymoteusz.Boba@gmail.com>
Date: Fri, 26 Jul 2024 14:17:15 +0200
Subject: [PATCH] fix(iOS): Crash while pushing n different screens at the same
 time (#2249)

## Description

On Paper, while user tries to navigate to 2 or more screens at the same
time, updates were batched (on async thread), resulting rerendering
stack hierarchy on the next layout. In the same scenario, on Fabric,
updates are not being batched, which causes crash:

```
"<RNSNavigationController> is pushing the same view controller instance (<RNSScreen>) more than once which is not supported and is most likely an error in the application"
```

since `maybeAddToParentAndUpdateContainer` is being called to early.
This PR fixes this bug by moving the logic of calling that method on
`mountingTransactionDidMount` method.

I've also checked if this change does not call
`maybeAddToParentAndUpdateContainer` too frequent and I can't see more
than one correct call there (every doubled call is being catched by
`return` in conditions).

Fixes #2235.

## Changes

- Moved updating containers to `mountingTransactionDidMount` method.

## Screenshots / GIFs

### Before


https://github.com/user-attachments/assets/24daf575-d06c-4972-9166-89454be60126

### After


https://github.com/user-attachments/assets/d377c527-0b2c-4850-b0b5-8e93dd5f25ac

## Test code and steps to reproduce

You can use `Test2235.tsx` to test this change.

## Checklist

- [X] Included code example that can be used to test this change
- [x] Ensured that CI passes

---------

Co-authored-by: Kacper Kafara <kacperkafara@gmail.com>
---
 apps/src/tests/Test2235.tsx | 107 ++++++++++++++++++++++++++++++++++++
 apps/src/tests/index.ts     |   1 +
 ios/RNSScreenStack.mm       |  33 ++++++++---
 3 files changed, 134 insertions(+), 7 deletions(-)
 create mode 100644 apps/src/tests/Test2235.tsx

diff --git a/apps/src/tests/Test2235.tsx b/apps/src/tests/Test2235.tsx
new file mode 100644
index 0000000000..4c00670dde
--- /dev/null
+++ b/apps/src/tests/Test2235.tsx
@@ -0,0 +1,107 @@
+import React from 'react';
+import { View, StyleSheet } from 'react-native';
+import {
+  createNativeStackNavigator,
+  NativeStackNavigationProp,
+} from '@react-navigation/native-stack';
+import { Button } from '../shared';
+import { NavigationContainer, ParamListBase } from '@react-navigation/native';
+
+interface RouteProps {
+  navigation: NativeStackNavigationProp<ParamListBase>;
+}
+
+const MainScreen = ({ navigation }: RouteProps): React.JSX.Element => (
+  <View style={{ ...styles.container, backgroundColor: 'moccasin' }}>
+    <Button
+      title="Push 2 screens at once"
+      onPress={() => {
+        navigation.navigate('Detail');
+        navigation.navigate('Detail2');
+      }}
+    />
+    <Button
+      title="Open nested stack"
+      onPress={() => {
+        navigation.navigate('NestedStack');
+      }}
+    />
+  </View>
+);
+
+const NestedMainScreen = ({ navigation }: RouteProps): React.JSX.Element => (
+  <View style={{ ...styles.container, backgroundColor: 'moccasin' }}>
+    <Button
+      title="Go to detail"
+      onPress={() => {
+        navigation.navigate('Detail');
+        navigation.navigate('Detail2');
+      }}
+    />
+    <Button
+      title="Pop stack"
+      onPress={() => {
+        navigation.goBack();
+      }}
+    />
+  </View>
+);
+
+
+const DetailScreen = ({ navigation }: RouteProps): React.JSX.Element => (
+  <View style={{ ...styles.container, backgroundColor: 'thistle' }}>
+    <Button
+      title="Go back"
+      onPress={() => navigation.goBack()}
+    />
+  </View>
+);
+
+const DetailScreen2 = ({
+  navigation,
+}: RouteProps): React.JSX.Element => (
+  <View style={{ ...styles.container, backgroundColor: 'yellow' }}>
+    <Button
+      title="Go back"
+      onPress={() => navigation.goBack()}
+    />
+  </View>
+);
+
+const NestedStackScreen = (): React.JSX.Element => (
+  <NestedStack.Navigator screenOptions={{ headerBackVisible: false }}>
+    <NestedStack.Screen name="NestedDetail" component={NestedMainScreen} />
+    <Stack.Screen name="Detail" component={DetailScreen} />
+    <Stack.Screen name="Detail2" component={DetailScreen2} />
+  </NestedStack.Navigator>
+)
+
+const Stack = createNativeStackNavigator<ParamListBase>();
+const NestedStack = createNativeStackNavigator<ParamListBase>();
+
+const App = (): React.JSX.Element => (
+  <NavigationContainer>
+    <Stack.Navigator
+      screenOptions={{
+        headerBackVisible: false,
+      }}>
+      <Stack.Screen
+        name="Main"
+        component={MainScreen}
+        options={{ title: 'Simple Native Stack' }}
+      />
+      <Stack.Screen name="Detail" component={DetailScreen} />
+      <Stack.Screen name="Detail2" component={DetailScreen2} />
+      <Stack.Screen name="NestedStack" component={NestedStackScreen} />
+    </Stack.Navigator>
+  </NavigationContainer>
+);
+
+const styles = StyleSheet.create({
+  container: {
+    flex: 1,
+    paddingTop: 10,
+  },
+});
+
+export default App;
diff --git a/apps/src/tests/index.ts b/apps/src/tests/index.ts
index 9233d3ba8b..81fe1806a6 100644
--- a/apps/src/tests/index.ts
+++ b/apps/src/tests/index.ts
@@ -105,5 +105,6 @@ export { default as Test2223 } from './Test2223';
 export { default as Test2227 } from './Test2227';
 export { default as Test2229 } from './Test2229';
 export { default as Test2231 } from './Test2231';
+export { default as Test2235 } from './Test2235';
 export { default as TestScreenAnimation } from './TestScreenAnimation';
 export { default as TestHeader } from './TestHeader';
diff --git a/ios/RNSScreenStack.mm b/ios/RNSScreenStack.mm
index 433130f9c4..428e3ca6cc 100644
--- a/ios/RNSScreenStack.mm
+++ b/ios/RNSScreenStack.mm
@@ -36,7 +36,12 @@ @interface RNSScreenStackView () <
     UINavigationControllerDelegate,
     UIAdaptivePresentationControllerDelegate,
     UIGestureRecognizerDelegate,
-    UIViewControllerTransitioningDelegate>
+    UIViewControllerTransitioningDelegate
+#ifdef RCT_NEW_ARCH_ENABLED
+    ,
+    RCTMountingTransactionObserving
+#endif
+    >
 
 @property (nonatomic) NSMutableArray<UIViewController *> *presentedModals;
 @property (nonatomic) BOOL updatingModals;
@@ -1112,9 +1117,8 @@ - (void)mountChildComponentView:(UIView<RCTComponentViewProtocol> *)childCompone
 
   [_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
@@ -1149,9 +1153,6 @@ - (void)unmountChildComponentView:(UIView<RCTComponentViewProtocol> *)childCompo
   screenChildComponent.reactSuperview = nil;
   [_reactSubviews removeObject:screenChildComponent];
   [screenChildComponent removeFromSuperview];
-  dispatch_async(dispatch_get_main_queue(), ^{
-    [self maybeAddToParentAndUpdateContainer];
-  });
 }
 
 - (void)takeSnapshot
@@ -1163,6 +1164,24 @@ - (void)takeSnapshot
   }
 }
 
+- (void)mountingTransactionDidMount:(const facebook::react::MountingTransaction &)transaction
+               withSurfaceTelemetry:(const facebook::react::SurfaceTelemetry &)surfaceTelemetry
+{
+  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;
+    }
+  }
+}
+
 - (void)prepareForRecycle
 {
   [super prepareForRecycle];