From 5eb15d848c8225658601dd2e63f5767d546ae04b Mon Sep 17 00:00:00 2001
From: Joe Mun <12057449+joemun@users.noreply.github.com>
Date: Fri, 27 Oct 2023 04:41:13 -0700
Subject: [PATCH] fix(Android): Set stateWrapper in ScreenViewManager in Fabric
 (#1944)

## Description

We are in the process of migrating to New Architecture. We noticed a
suspicious error message in LogCat, and discovered that
`FabricEnabledViewGroup`'s state updates in `updateScreenSizeFabric()`
would fail with "setState called without a StateWrapper", due to:


https://github.com/facebook/react-native/blob/79d620dff7561e6c9370184f9f19e79643486716/packages/react-native/ReactAndroid/src/main/java/com/facebook/react/uimanager/FabricViewStateManager.java#L65

## Changes

Call `fabricViewStateManager.setStateWrapper()` in
`ScreenViewManager.kt`.

<img width="1285" alt="Screen Shot 2023-10-23 at 11 20 09 AM"
src="https://github.com/software-mansion/react-native-screens/assets/12057449/3411e42d-fba3-4b11-b601-66cfb557cc62">

<img width="1808" alt="Screen Shot 2023-10-23 at 11 32 28 AM"
src="https://github.com/software-mansion/react-native-screens/assets/12057449/28d1f518-da4a-4971-bfb2-5c2c5b7458e0">

## Test code and steps to reproduce

Run `FabricExample` and `Example`. View LogCat logs.

## Checklist

- [ ] Included code example that can be used to test this change
- [ ] Updated TS types
- [ ] Updated documentation: <!-- For adding new props to native-stack
-->
- [ ]
https://github.com/software-mansion/react-native-screens/blob/main/guides/GUIDE_FOR_LIBRARY_AUTHORS.md
- [ ]
https://github.com/software-mansion/react-native-screens/blob/main/native-stack/README.md
- [ ]
https://github.com/software-mansion/react-native-screens/blob/main/src/types.tsx
- [ ]
https://github.com/software-mansion/react-native-screens/blob/main/src/native-stack/types.tsx
- [ ] Ensured that CI passes

---------

Co-authored-by: Tymoteusz Boba <Tymoteusz.Boba@gmail.com>
---
 .../com/swmansion/rnscreens/ScreenViewManager.kt   | 14 ++++++++++++++
 .../swmansion/rnscreens/FabricEnabledViewGroup.kt  |  4 ++++
 2 files changed, 18 insertions(+)

diff --git a/android/src/main/java/com/swmansion/rnscreens/ScreenViewManager.kt b/android/src/main/java/com/swmansion/rnscreens/ScreenViewManager.kt
index 5d92778f65..8c70876ef4 100644
--- a/android/src/main/java/com/swmansion/rnscreens/ScreenViewManager.kt
+++ b/android/src/main/java/com/swmansion/rnscreens/ScreenViewManager.kt
@@ -4,6 +4,8 @@ import com.facebook.react.bridge.JSApplicationIllegalArgumentException
 import com.facebook.react.bridge.ReadableMap
 import com.facebook.react.common.MapBuilder
 import com.facebook.react.module.annotations.ReactModule
+import com.facebook.react.uimanager.ReactStylesDiffMap
+import com.facebook.react.uimanager.StateWrapper
 import com.facebook.react.uimanager.ThemedReactContext
 import com.facebook.react.uimanager.ViewGroupManager
 import com.facebook.react.uimanager.ViewManagerDelegate
@@ -35,6 +37,18 @@ class ScreenViewManager : ViewGroupManager<Screen>(), RNSScreenManagerInterface<
         setActivityState(view, activityState.toInt())
     }
 
+    override fun updateState(
+        view: Screen,
+        props: ReactStylesDiffMap?,
+        stateWrapper: StateWrapper?
+    ): Any? {
+        if (BuildConfig.IS_NEW_ARCHITECTURE_ENABLED) {
+            // fabricViewStateManager should never be null in Fabric. The null check is only for Paper's empty impl.
+            view.fabricViewStateManager?.setStateWrapper(stateWrapper)
+        }
+        return super.updateState(view, props, stateWrapper)
+    }
+
     @ReactProp(name = "activityState")
     fun setActivityState(view: Screen, activityState: Int) {
         if (activityState == -1) {
diff --git a/android/src/paper/java/com/swmansion/rnscreens/FabricEnabledViewGroup.kt b/android/src/paper/java/com/swmansion/rnscreens/FabricEnabledViewGroup.kt
index 106638858d..6c1c238bd6 100644
--- a/android/src/paper/java/com/swmansion/rnscreens/FabricEnabledViewGroup.kt
+++ b/android/src/paper/java/com/swmansion/rnscreens/FabricEnabledViewGroup.kt
@@ -2,8 +2,12 @@ package com.swmansion.rnscreens
 
 import android.view.ViewGroup
 import com.facebook.react.bridge.ReactContext
+import com.facebook.react.uimanager.FabricViewStateManager
 
 abstract class FabricEnabledViewGroup constructor(context: ReactContext?) : ViewGroup(context) {
+
+    val fabricViewStateManager get() = null as FabricViewStateManager?
+
     protected fun updateScreenSizeFabric(width: Int, height: Int) {
         // do nothing
     }