From 91d89c4b83258f0ab8ae7c1899541e466fb93beb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Alex=20Du=C5=BCy?=
 <91994767+alduzy@users.noreply.github.com>
Date: Mon, 7 Oct 2024 22:23:31 +0200
Subject: [PATCH] fix(iOS): header subviews layout on tab change (#2385)

## Description

This PR intents to fix header subviews incorrect layout when changing
tabs.

The previous solution did layout the subviews correctly in the test
cases, but triggered an undesirable `layoutIfNeeded` when going back
from tab to tab. In such case the navigation layout happened without
updating subview's layout metrics.

Moving the logic to subview resolves the issue as the re-layout is now
triggered only when subview's layout metrics are updated.

Related fixes from the past:
https://github.com/software-mansion/react-native-screens/pull/2316,
https://github.com/software-mansion/react-native-screens/pull/2248

## Changes

- combined `Test2231.tsx` with `Test432.tsx` to create comprehensive
test case
- moved re-layout logic to subview

## Screenshots / GIFs

### Before
![Screenshot 2024-10-04 at 10 10
15](https://github.com/user-attachments/assets/1a8a747e-fe1d-4b03-ba63-1891732d7d77)

### After
![Screenshot 2024-10-04 at 10 09
37](https://github.com/user-attachments/assets/68b3554f-d67d-47f4-946d-ace60e1bbf83)


## Test code and steps to reproduce

- Use `Test432.tsx` repro

## Checklist

- [x] Included code example that can be used to test this change
- [x] Ensured that CI passes
---
 apps/src/tests/Test2231.tsx        |  56 ------------
 apps/src/tests/Test432.tsx         | 137 +++++++++++++++++++----------
 apps/src/tests/index.ts            |   1 -
 ios/RNSScreenStackHeaderConfig.mm  |   3 -
 ios/RNSScreenStackHeaderSubview.mm |  24 +++++
 5 files changed, 115 insertions(+), 106 deletions(-)
 delete mode 100644 apps/src/tests/Test2231.tsx

diff --git a/apps/src/tests/Test2231.tsx b/apps/src/tests/Test2231.tsx
deleted file mode 100644
index 9449ffc3ff..0000000000
--- a/apps/src/tests/Test2231.tsx
+++ /dev/null
@@ -1,56 +0,0 @@
-import { createNativeStackNavigator } from '@react-navigation/native-stack';
-import React, { useLayoutEffect, useState } from 'react';
-import { View } from 'react-native';
-
-import { SettingsSwitch, Square } from '../shared';
-import { NavigationContainer } from '@react-navigation/native';
-
-const SettingsScreen = ({ navigation }: any) => {
-  const [hasLeftItem, setHasLeftItem] = useState(false);
-
-  const square1 = (props: { tintColor?: string }) => (
-    <View style={{ gap: 8, flexDirection: 'row' }}>
-      {hasLeftItem && <Square {...props} color="green" size={20} />}
-      <Square {...props} color="green" size={20} />
-    </View>
-  );
-
-  const square2 = (props: { tintColor?: string }) => (
-    <Square {...props} color="red" size={20} />
-  );
-
-  useLayoutEffect(() => {
-    navigation.setOptions({
-      headerRight: square1,
-      headerTitle: undefined,
-      headerLeft: hasLeftItem ? square2 : undefined,
-      headerBackTitleVisible: false,
-    });
-  }, [navigation, hasLeftItem]);
-
-  return (
-    <SettingsSwitch
-      label="Left item"
-      value={hasLeftItem}
-      onValueChange={setHasLeftItem}
-    />
-  );
-};
-
-const Stack = createNativeStackNavigator();
-
-const App = () => (
-  <NavigationContainer>
-    <Stack.Navigator>
-      <Stack.Screen
-        name="Settings"
-        component={SettingsScreen}
-        options={{
-          headerTintColor: 'hotpink',
-        }}
-      />
-    </Stack.Navigator>
-  </NavigationContainer>
-);
-
-export default App;
diff --git a/apps/src/tests/Test432.tsx b/apps/src/tests/Test432.tsx
index 8274bfd299..f58bd3b752 100644
--- a/apps/src/tests/Test432.tsx
+++ b/apps/src/tests/Test432.tsx
@@ -1,37 +1,36 @@
-import { Pressable, View, Button, Text } from 'react-native';
-
-import { NavigationContainer, useNavigation } from '@react-navigation/native';
+import { View, Button, Text } from 'react-native';
+import { NavigationContainer } from '@react-navigation/native';
 import {
   NativeStackScreenProps,
   createNativeStackNavigator,
 } from '@react-navigation/native-stack';
-import React, { useCallback } from 'react';
+import React, { useEffect, useLayoutEffect, useState } from 'react';
+import { createBottomTabNavigator } from '@react-navigation/bottom-tabs';
+import { Square } from '../shared';
 
-type RootStackParamList = {
+type StackParamList = {
   Home: undefined;
+  Details: undefined;
   Settings: undefined;
+  Info: undefined;
 };
-type RootStackScreenProps<T extends keyof RootStackParamList> =
-  NativeStackScreenProps<RootStackParamList, T>;
 
-const HomeScreen = ({ navigation }: RootStackScreenProps<'Home'>) => {
-  const [x, setX] = React.useState(false);
-  React.useEffect(() => {
-    navigation.setOptions({
-      headerBackVisible: !x,
-      headerRight: x
-        ? () => (
-            <View style={{ backgroundColor: 'green', width: 20, height: 20 }} />
-          )
-        : () => (
-            <View style={{ backgroundColor: 'green', width: 10, height: 10 }} />
-          ),
-    });
-  }, [navigation, x]);
+type StackScreenProps<T extends keyof StackParamList> = NativeStackScreenProps<
+  StackParamList,
+  T
+>;
 
+const HomeScreen = ({ navigation }: StackScreenProps<'Home'>) => {
   return (
     <View style={{ flex: 1, alignItems: 'center', justifyContent: 'center' }}>
-      <Button title="Tap me for header update" onPress={() => setX(!x)} />
+      <Button
+        title={'Go to details'}
+        onPress={() => navigation.navigate('Details')}
+      />
+      <Button
+        title={'Go to info'}
+        onPress={() => navigation.navigate('Info')}
+      />
       <Button
         title={'Show settings'}
         onPress={() => navigation.navigate('Settings')}
@@ -40,49 +39,95 @@ const HomeScreen = ({ navigation }: RootStackScreenProps<'Home'>) => {
   );
 };
 
+const DetailsScreen = ({ navigation }: StackScreenProps<'Details'>) => {
+  const [x, setX] = useState(false);
+  useEffect(() => {
+    navigation.setOptions({
+      headerBackVisible: !x,
+      headerRight: () =>
+        x ? <Square size={20} color="green" /> : <Square size={10} />,
+    });
+  }, [navigation, x]);
+
+  return <Button title="Toggle subviews" onPress={() => setX(prev => !prev)} />;
+};
+
 const SettingsScreen = () => {
-  return (
-    <View style={{ flex: 1, alignItems: 'center', justifyContent: 'center' }}>
-      <Text>Settings</Text>
+  return <Text>Settings</Text>;
+};
+
+const InfoScreen = ({ navigation }: StackScreenProps<'Info'>) => {
+  const [hasLeftItem, setHasLeftItem] = useState(false);
+
+  const square1 = (props: { tintColor?: string }) => (
+    <View style={{ gap: 8, flexDirection: 'row' }}>
+      {hasLeftItem && <Square {...props} color="green" size={20} />}
+      <Square {...props} color="green" size={20} />
     </View>
   );
-};
 
-const RootStack = createNativeStackNavigator<RootStackParamList>();
-const RootNavigator = () => {
-  const navigation = useNavigation();
-  const headerRight = useCallback(
-    () => (
-      <Pressable
-        onPress={() => {
-          navigation.goBack();
-        }}>
-        <Text>Close</Text>
-      </Pressable>
-    ),
-    [navigation],
+  const square2 = (props: { tintColor?: string }) => (
+    <Square {...props} color="red" size={20} />
   );
+
+  useLayoutEffect(() => {
+    navigation.setOptions({
+      headerRight: square1,
+      headerTitle: undefined,
+      headerLeft: hasLeftItem ? square2 : undefined,
+    });
+  }, [navigation, hasLeftItem]);
+
   return (
-    <RootStack.Navigator>
-      <RootStack.Screen name="Home" component={HomeScreen} />
-      <RootStack.Screen
+    <Button
+      title="Toggle subviews"
+      onPress={() => setHasLeftItem(prev => !prev)}
+    />
+  );
+};
+
+const Stack = createNativeStackNavigator<StackParamList>();
+
+const StackNavigator = () => {
+  return (
+    <Stack.Navigator>
+      <Stack.Screen
+        name="Home"
+        component={HomeScreen}
+        options={{
+          headerRight: () => <Square size={20} color="black" />,
+        }}
+      />
+      <Stack.Screen name="Details" component={DetailsScreen} />
+      <Stack.Screen
+        name="Info"
+        component={InfoScreen}
+        options={{
+          headerTintColor: 'hotpink',
+        }}
+      />
+      <Stack.Screen
         name="Settings"
         component={SettingsScreen}
         options={{
           presentation: 'modal',
           animation: 'slide_from_bottom',
-          headerShown: true,
-          headerRight: headerRight,
+          headerRight: () => <Square size={30} />,
         }}
       />
-    </RootStack.Navigator>
+    </Stack.Navigator>
   );
 };
 
+const Tabs = createBottomTabNavigator();
+
 export default function App() {
   return (
     <NavigationContainer>
-      <RootNavigator />
+      <Tabs.Navigator screenOptions={{ headerShown: false }}>
+        <Tabs.Screen name="First" component={StackNavigator} />
+        <Tabs.Screen name="Second" component={StackNavigator} />
+      </Tabs.Navigator>
     </NavigationContainer>
   );
 }
diff --git a/apps/src/tests/index.ts b/apps/src/tests/index.ts
index 94199a0fe4..f08e4f211c 100644
--- a/apps/src/tests/index.ts
+++ b/apps/src/tests/index.ts
@@ -105,7 +105,6 @@ export { default as Test2184 } from './Test2184';
 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 Test2232 } from './Test2232';
 export { default as Test2235 } from './Test2235';
 export { default as Test2252 } from './Test2252';
diff --git a/ios/RNSScreenStackHeaderConfig.mm b/ios/RNSScreenStackHeaderConfig.mm
index a6dbb1a041..e10ecf04b0 100644
--- a/ios/RNSScreenStackHeaderConfig.mm
+++ b/ios/RNSScreenStackHeaderConfig.mm
@@ -689,9 +689,6 @@ + (void)updateViewController:(UIViewController *)vc
         break;
       }
     }
-    // We're forcing a re-layout when the subviews change,
-    // see: https://github.com/software-mansion/react-native-screens/pull/2316
-    [navctr.view layoutIfNeeded];
   }
 
   // This assignment should be done after `navitem.titleView = ...` assignment (iOS 16.0 bug).
diff --git a/ios/RNSScreenStackHeaderSubview.mm b/ios/RNSScreenStackHeaderSubview.mm
index 8ee947d821..dca24e3800 100644
--- a/ios/RNSScreenStackHeaderSubview.mm
+++ b/ios/RNSScreenStackHeaderSubview.mm
@@ -1,5 +1,6 @@
 #import "RNSScreenStackHeaderSubview.h"
 #import "RNSConvert.h"
+#import "RNSScreenStackHeaderConfig.h"
 
 #ifdef RCT_NEW_ARCH_ENABLED
 #import <react/renderer/components/rnscreens/ComponentDescriptors.h>
@@ -22,6 +23,27 @@ @implementation RNSScreenStackHeaderSubview
 
 #pragma mark - Common
 
+- (nullable RNSScreenStackHeaderConfig *)getHeaderConfig
+{
+  RNSScreenStackHeaderConfig *headerConfig = (RNSScreenStackHeaderConfig *_Nullable)self.reactSuperview;
+#ifndef NDEBUG
+  if (headerConfig != nil && ![headerConfig isKindOfClass:[RNSScreenStackHeaderConfig class]]) {
+    RCTLogError(@"[RNScreens] Invalid view type, expecting RNSScreenStackHeaderConfig, got: %@", headerConfig);
+    return nil;
+  }
+#endif
+  return headerConfig;
+}
+
+// We're forcing the navigation controller's view to re-layout
+// see: https://github.com/software-mansion/react-native-screens/pull/2385
+- (void)layoutNavigationBarIfNeeded
+{
+  RNSScreenStackHeaderConfig *headerConfig = [self getHeaderConfig];
+  UINavigationController *navctr = headerConfig.screenView.reactViewController.navigationController;
+  [navctr.navigationBar layoutIfNeeded];
+}
+
 #ifdef RCT_NEW_ARCH_ENABLED
 
 #pragma mark - Fabric specific
@@ -78,6 +100,7 @@ - (void)updateLayoutMetrics:(const react::LayoutMetrics &)layoutMetrics
         self);
   } else {
     self.bounds = CGRect{CGPointZero, frame.size};
+    [self layoutNavigationBarIfNeeded];
   }
 }
 
@@ -102,6 +125,7 @@ - (void)reactSetFrame:(CGRect)frame
   // Block any attempt to set coordinates on RNSScreenStackHeaderSubview. This
   // makes UINavigationBar the only one to control the position of header content.
   [super reactSetFrame:CGRectMake(0, 0, frame.size.width, frame.size.height)];
+  [self layoutNavigationBarIfNeeded];
 }
 
 #endif // RCT_NEW_ARCH_ENABLED