Fix: layout shift on orientation change

FabricExample/App.tsx

@@ -4,5 +4,6 @@ import { featureFlags } from '../src';
 featureFlags.experiment.synchronousScreenUpdatesEnabled = false
 featureFlags.experiment.synchronousHeaderConfigUpdatesEnabled = false
 featureFlags.experiment.synchronousHeaderSubviewUpdatesEnabled = false
+featureFlags.experiment.earlyScreenOrientationChangeEnabled = true
 
 export default App;

android/src/main/java/com/swmansion/rnscreens/ScreenViewManager.kt

@@ -342,6 +342,11 @@ open class ScreenViewManager :
 
     // END mark: iOS-only
 
+    override fun setEarlyScreenOrientationChangeEnabled(
+        view: Screen?,
+        value: Boolean,
+    ) = Unit // represents a feature flag and is checked via getProps() in RNSScreenComponentDescriptor.h
+
     @ReactProp(name = "sheetAllowedDetents")
     override fun setSheetAllowedDetents(
         view: Screen,

android/src/paper/java/com/facebook/react/viewmanagers/RNSScreenManagerDelegate.java

@@ -137,6 +137,9 @@ public void setProperty(T view, String propName, @Nullable Object value) {
       case "synchronousShadowStateUpdatesEnabled":
         mViewManager.setSynchronousShadowStateUpdatesEnabled(view, value == null ? false : (boolean) value);
         break;
+      case "earlyScreenOrientationChangeEnabled":
+        mViewManager.setEarlyScreenOrientationChangeEnabled(view, value == null ? false : (boolean) value);
+        break;
       default:
         super.setProperty(view, propName, value);
     }

android/src/paper/java/com/facebook/react/viewmanagers/RNSScreenManagerInterface.java

@@ -53,4 +53,5 @@ public interface RNSScreenManagerInterface<T extends View>  {
   void setRightScrollEdgeEffect(T view, @Nullable String value);
   void setTopScrollEdgeEffect(T view, @Nullable String value);
   void setSynchronousShadowStateUpdatesEnabled(T view, boolean value);
+  void setEarlyScreenOrientationChangeEnabled(T view, boolean value);
 }

apps/src/tests/Test2933.tsx

@@ -0,0 +1,138 @@
+import React, { useEffect, useState } from 'react';
+import type { PropsWithChildren } from 'react';
+import {
+  StyleSheet,
+  Text,
+  useColorScheme,
+  View,
+} from 'react-native';
+
+import { Screen, ScreenStack } from 'react-native-screens';
+
+type SectionProps = PropsWithChildren<{
+  title: string;
+}>;
+
+function Section({ children, title }: SectionProps): React.JSX.Element {
+  return (
+    <View style={styles.sectionContainer}>
+      <Text>{title}</Text>
+      <Text>{children}</Text>
+    </View>
+  );
+}
+
+function floodJsThread() {
+  setInterval(() => {
+    const end = Date.now() + 10;
+    while (Date.now() < end) {
+      // Intentionally do nothing; just burn CPU cycles.
+      Math.sqrt(Math.random());
+    }
+  }, 12);
+}
+
+/*
+ * create artificial pressure in the JS thread to show off thep problem.
+ * */
+floodJsThread();
+floodJsThread();
+
+function AppMain(): React.JSX.Element {
+  const isDarkMode = useColorScheme() === 'dark';
+
+  const backgroundStyle = {
+    flex: 1,
+    backgroundColor: 'white',
+  };
+
+  const [num, setNum] = useState(0);
+
+  useEffect(() => {
+    let i = 0;
+    setInterval(() => {
+      i++;
+      setNum(i);
+    }, 2)
+  }, []);
+
+  /*
+   * To keep the template simple and small we're adding padding to prevent view
+   * from rendering under the System UI.
+   * For bigger apps the reccomendation is to use `react-native-safe-area-context`:
+   * https://github.com/AppAndFlow/react-native-safe-area-context
+   *
+   * You can read more about it here:
+   * https://github.com/react-native-community/discussions-and-proposals/discussions/827
+   */
+  const safePadding = '5%';
+
+  return (
+    <View style={backgroundStyle}>
+      {/* // TODO nowy task
+      <StatusBar
+        barStyle={isDarkMode ? 'light-content' : 'dark-content'}
+        backgroundColor={backgroundStyle.backgroundColor}
+      /> */}
+      <View
+        style={{
+          flex: 1,
+          paddingHorizontal: safePadding,
+          paddingBottom: safePadding,
+          justifyContent: 'center',
+          alignContent: 'center',
+        }}>
+        <Section title="Step One">
+          This test shows how the native layout update triggers a layout shift.
+        </Section>
+        <Section title="Step One">
+          There is a view with a blue background. We don't expect to ever see
+          flashes of the blue background.
+        </Section>
+        <Section title="Orientation"> {num} </Section>
+      </View>
+    </View>
+  );
+}
+
+const styles = StyleSheet.create({
+  sectionContainer: {
+    marginTop: 32,
+    paddingHorizontal: 24,
+  },
+  sectionTitle: {
+    fontSize: 24,
+    fontWeight: '600',
+  },
+  sectionDescription: {
+    marginTop: 8,
+    fontSize: 18,
+    fontWeight: '400',
+  },
+  highlight: {
+    fontWeight: '700',
+  },
+});
+
+function App() {
+  return (
+    <ScreenStack style={{ flex: 1, backgroundColor: 'red' }}>
+      <Screen
+        style={{
+          position: 'absolute',
+          top: 0,
+          left: 0,
+          bottom: 0,
+          right: 0,
+          backgroundColor: 'blue',
+          padding: 20,
+        }}
+        enabled
+        isNativeStack>
+        <AppMain />
+      </Screen>
+    </ScreenStack>
+  );
+}
+
+export default App;

apps/src/tests/index.ts

@@ -138,6 +138,7 @@ export { default as Test2855 } from './Test2855';
 export { default as Test2877 } from './Test2877'; // [E2E created](iOS): issue is related to formSheet on iOS
 export { default as Test2895 } from './Test2895';
 export { default as Test2899 } from './Test2899';
+export { default as Test2933 } from './Test2933';
 export { default as Test2926 } from './Test2926'; // [E2E created](iOS): PR related to iOS search bar
 export { default as Test2949 } from './Test2949'; // [E2E skipped]: can't check system bars styles
 export { default as Test2963 } from './Test2963'; // [E2E created](iOS): issue related to iOS

common/cpp/react/renderer/components/rnscreens/RNSScreenComponentDescriptor.h

@@ -2,10 +2,13 @@
 
 #ifdef ANDROID
 #include <fbjni/fbjni.h>
+#include <react/renderer/uimanager/UIManager.h>
+#include "RNSScreenShadowNodeCommitHook.h"
 #endif // ANDROID
 #include <react/debug/react_native_assert.h>
 #include <react/renderer/components/rnscreens/Props.h>
 #include <react/renderer/components/rnscreens/utils/RectUtil.h>
+#include <react/renderer/components/root/RootShadowNode.h>
 #include <react/renderer/core/ConcreteComponentDescriptor.h>
 #include "RNSScreenShadowNode.h"
 
@@ -16,6 +19,16 @@ using namespace rnscreens;
 
 class RNSScreenComponentDescriptor final
     : public ConcreteComponentDescriptor<RNSScreenShadowNode> {
+ private:
+#ifdef ANDROID
+  /*
+   * A commit hook that triggers on `shadowTreeWillCommit` event,
+   * and can read the properties of RootShadowNodes for determining screen
+   * orientation.
+   */
+  mutable std::shared_ptr<RNSScreenShadowNodeCommitHook> commitHook_;
+#endif // ANDROID
+
  public:
   using ConcreteComponentDescriptor::ConcreteComponentDescriptor;
 
@@ -26,14 +39,32 @@ class RNSScreenComponentDescriptor final
     react_native_assert(
         dynamic_cast<YogaLayoutableShadowNode *>(&screenShadowNode));
     auto &layoutableShadowNode =
-        dynamic_cast<YogaLayoutableShadowNode &>(screenShadowNode);
+        static_cast<YogaLayoutableShadowNode &>(screenShadowNode);
 
     auto state =
         std::static_pointer_cast<const RNSScreenShadowNode::ConcreteState>(
             shadowNode.getState());
     auto stateData = state->getData();
-
 #ifdef ANDROID
+    // get featureFlags.experiment.earlyScreenOrientationChangeEnabled from
+    // Screen props enable the commit hook only when the developer asks to
+    react_native_assert(
+        dynamic_cast<const RNSScreenProps *>(
+            screenShadowNode.getProps().get()));
+    auto props =
+        dynamic_cast<const RNSScreenProps *>(screenShadowNode.getProps().get());
+
+    if (!commitHook_ && props->earlyScreenOrientationChangeEnabled) {
+      // For the the application that needs to react to orientation change
+      // as early as possible, we attach a commit hook that checks for the
+      // change in the old vs new RootShadowNode. The hook cannot be attached in
+      // the constructor because UIManager is still missing from
+      // ContextContainer. Instead, we do it here, on the first call to the
+      // function.
+      commitHook_ =
+          std::make_shared<RNSScreenShadowNodeCommitHook>(contextContainer_);
+    }
+
     if (stateData.frameSize.width != 0 && stateData.frameSize.height != 0) {
       // When we receive dimensions from JVM side we can remove padding used for
       // correction, and we can stop applying height and offset corrections for
@@ -66,9 +97,15 @@ class RNSScreenComponentDescriptor final
           FrameCorrectionModes::Mode::FrameHeightCorrection);
       screenShadowNode.getFrameCorrectionModes().unset(
           FrameCorrectionModes::Mode::FrameOriginCorrection);
-
       layoutableShadowNode.setSize(
           Size{stateData.frameSize.width, stateData.frameSize.height});
+    } else if (
+        stateData.frameSize.width == 0 && stateData.frameSize.height == 0) {
+      // Reset YogaNode so it recalculates its layout. Useful for the case
+      // when native orientation changes and react has not been notified yet.
+      // The if condition holds true on first render and when it is reset inside
+      // RNSScreenShadowNodeCommitHook.
+      layoutableShadowNode.setSize({YGUndefined, YGUndefined});
     }
 #else
     if (stateData.frameSize.width != 0 && stateData.frameSize.height != 0) {

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

@@ -155,6 +155,10 @@ FrameCorrectionModes &RNSScreenShadowNode::getFrameCorrectionModes() {
   return getStateDataMutable().getFrameCorrectionModes();
 }
 
+void RNSScreenShadowNode::resetFrameSizeState() {
+  getStateDataMutable().frameSize = {0, 0};
+}
+
 RNSScreenShadowNode::StateData &RNSScreenShadowNode::getStateDataMutable() {
   // We assume that this method is called to mutate the data, so we ensure
   // we're unsealed.

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

@@ -38,8 +38,10 @@ class JSI_EXPORT RNSScreenShadowNode final : public ConcreteViewShadowNode<
 
   FrameCorrectionModes &getFrameCorrectionModes();
 
- private:
 #ifdef ANDROID
+  void resetFrameSizeState();
+
+ private:
   void applyFrameCorrections();
 
   StateData &getStateDataMutable();

common/cpp/react/renderer/components/rnscreens/RNSScreenShadowNodeCommitHook.cpp

@@ -0,0 +1,103 @@
+#ifdef ANDROID
+
+#include "RNSScreenShadowNodeCommitHook.h"
+#include <stack>
+
+namespace facebook {
+namespace react {
+
+RNSScreenShadowNodeCommitHook::RNSScreenShadowNodeCommitHook(
+    std::shared_ptr<const ContextContainer> contextContainer)
+    : contextContainer_(contextContainer) {
+  if (contextContainer_) {
+    auto fabricUIManager =
+        contextContainer_
+            ->at<jni::alias_ref<facebook::react::JFabricUIManager::javaobject>>(
+                "FabricUIManager");
+    fabricUIManager->getBinding()
+        ->getScheduler()
+        ->getUIManager()
+        ->registerCommitHook(*this);
+  }
+}
+
+RNSScreenShadowNodeCommitHook::~RNSScreenShadowNodeCommitHook() noexcept {
+  if (contextContainer_) {
+    auto fabricUIManager =
+        contextContainer_
+            ->at<jni::alias_ref<facebook::react::JFabricUIManager::javaobject>>(
+                "FabricUIManager");
+    fabricUIManager->getBinding()
+        ->getScheduler()
+        ->getUIManager()
+        ->unregisterCommitHook(*this);
+  }
+}
+
+RootShadowNode::Unshared RNSScreenShadowNodeCommitHook::shadowTreeWillCommit(
+    const facebook::react::ShadowTree &shadowTree,
+    const RootShadowNode::Shared &oldRootShadowNode,
+    const RootShadowNode::Unshared &newRootShadowNode,
+    const ShadowTreeCommitOptions &) noexcept {
+  auto oldRootProps =
+      std::static_pointer_cast<const RootProps>(oldRootShadowNode->getProps());
+  auto newRootProps =
+      std::static_pointer_cast<const RootProps>(newRootShadowNode->getProps());
+
+  const bool wasHorizontal = isHorizontal_(*oldRootProps.get());
+  const bool willBeHorizontal = isHorizontal_(*newRootProps.get());
+
+  const bool orientationDidChange = wasHorizontal != willBeHorizontal;
+
+  std::shared_ptr<ShadowNode> finalRootShadowNode = newRootShadowNode;
+  if (orientationDidChange) {
+    std::vector<const RNSScreenShadowNode *> screens;
+    findScreenNodes(newRootShadowNode, screens);
+
+    for (auto screen : screens) {
+      const auto rootShadowNodeClone = newRootShadowNode->cloneTree(
+          screen->getFamily(), [](const ShadowNode &oldShadowNode) {
+            auto clone =
+                oldShadowNode.clone({.state = oldShadowNode.getState()});
+            auto screenNode = static_pointer_cast<RNSScreenShadowNode>(clone);
+            auto yogaNode =
+                static_pointer_cast<YogaLayoutableShadowNode>(clone);
+
+            screenNode->resetFrameSizeState();
+            yogaNode->setSize({YGUndefined, YGUndefined});
+
+            return clone;
+          });
+
+      if (rootShadowNodeClone) {
+        finalRootShadowNode = rootShadowNodeClone;
+      }
+    }
+  }
+
+  return std::static_pointer_cast<RootShadowNode>(finalRootShadowNode);
+}
+
+void RNSScreenShadowNodeCommitHook::findScreenNodes(
+    const std::shared_ptr<const ShadowNode> &rootShadowNode,
+    std::vector<const RNSScreenShadowNode *> &screenNodes) {
+  std::stack<const ShadowNode *> shadowNodesToVisit;
+  shadowNodesToVisit.emplace(rootShadowNode.get());
+
+  while (!shadowNodesToVisit.empty()) {
+    auto node = shadowNodesToVisit.top();
+    shadowNodesToVisit.pop();
+
+    for (auto const &child : node->getChildren()) {
+      if (node->getComponentHandle() == RNSScreenShadowNode::Handle()) {
+        screenNodes.push_back(static_cast<const RNSScreenShadowNode *>(node));
+      }
+      shadowNodesToVisit.emplace(child.get());
+    }
+  }
+}
+
+} // namespace react
+} // namespace facebook
+
+#endif // ANDROID