fix: apply same logic to ScreenContainer (#1116)

Added delayed update logic to ScreenContainer in order to resolve issues with too early detached nested navigators when going back from a screen in native-stack with nested different navigators. This change was part of synchronous fragment updates introduced in #1066. Going back like that triggers removeAllScreens code of the nested navigator and it resolved in detaching all of the child screens before the update of native-stack was dispatched, leading to blank screen from the start of navigation.

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

@@ -20,7 +20,8 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
     protected val mScreenFragments = ArrayList<T>()
     @JvmField
     protected var mFragmentManager: FragmentManager? = null
-    protected var mIsAttached = false
+    private var mIsAttached = false
+    private var mNeedUpdate = false
     private var mLayoutEnqueued = false
     private val mLayoutCallback: ChoreographerCompat.FrameCallback = object : ChoreographerCompat.FrameCallback() {
         override fun doFrame(frameTimeNanos: Long) {
@@ -81,7 +82,7 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
         get() = mParentScreenFragment != null
 
     fun notifyChildUpdate() {
-        onScreenChanged()
+        performUpdatesNow()
     }
 
     protected open fun adapt(screen: Screen): T {
@@ -128,9 +129,9 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
             return null
         }
 
-    protected open fun setFragmentManager(fm: FragmentManager) {
+    private fun setFragmentManager(fm: FragmentManager) {
         mFragmentManager = fm
-        onScreenChanged()
+        performUpdatesNow()
     }
 
     private fun setupFragmentManager() {
@@ -242,7 +243,7 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
         // delayed lifecycle (due to transitions). As a result due to ongoing transitions the fragment
         // may choose not to remove the view despite the parent container being completely detached
         // from the view hierarchy until the transition is over. In such a case when the container gets
-        // re-attached while tre transition is ongoing, the child view would still be there and we'd
+        // re-attached while the transition is ongoing, the child view would still be there and we'd
         // attempt to re-attach it to with a misconfigured fragment. This would result in a crash. To
         // avoid it we clear all the children here as we attach all the child fragments when the
         // container is reattached anyways. We don't use `removeAllViews` since it does not check if the
@@ -265,11 +266,36 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
         }
     }
 
-    protected open fun onScreenChanged() {
-        if (!mIsAttached || mFragmentManager == null) {
+    private fun onScreenChanged() {
+        // we perform update in `onBeforeLayout` of `ScreensShadowNode` by adding an UIBlock
+        // which is called after updating children of the ScreenContainer.
+        // We do it there because `onUpdate` logic requires all changes of children to be already
+        // made in order to provide proper animation for fragment transition for ScreenStack
+        // and this in turn makes nested ScreenContainers detach too early and disappear
+        // before transition if also not dispatched after children updates.
+        // The exception to this rule is `updateImmediately` which is triggered by actions
+        // not connected to React view hierarchy changes, but rather internal events
+        mNeedUpdate = true
+    }
+
+    protected fun performUpdatesNow() {
+        // we want to update the immediately when the fragment manager is set or native back button
+        // dismiss is dispatched or Screen's activityState changes since it is not connected to React
+        // view hierarchy changes and will not trigger `onBeforeLayout` method of `ScreensShadowNode`
+        mNeedUpdate = true
+        performUpdates()
+    }
+
+    fun performUpdates() {
+        if (!mNeedUpdate || !mIsAttached || mFragmentManager == null) {
             return
         }
+        mNeedUpdate = false
+        onUpdate()
+        notifyContainerUpdate()
+    }
 
+    open fun onUpdate() {
         // detach screens that are no longer active
         val orphaned: MutableSet<Fragment> = HashSet(requireNotNull(mFragmentManager, { "mFragmentManager is null when performing update in ScreenContainer" }).fragments)
         for (screenFragment in mScreenFragments) {
@@ -308,7 +334,6 @@ open class ScreenContainer<T : ScreenFragment>(context: Context?) : ViewGroup(co
             }
             screenFragment.screen.setTransitioning(transitioning)
         }
-        notifyContainerUpdate()
     }
 
     protected open fun notifyContainerUpdate() {

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

@@ -1,7 +1,9 @@
 package com.swmansion.rnscreens
 
 import android.view.View
+import com.facebook.react.bridge.ReactApplicationContext
 import com.facebook.react.module.annotations.ReactModule
+import com.facebook.react.uimanager.LayoutShadowNode
 import com.facebook.react.uimanager.ThemedReactContext
 import com.facebook.react.uimanager.ViewGroupManager
 
@@ -36,6 +38,10 @@ class ScreenContainerViewManager : ViewGroupManager<ScreenContainer<*>>() {
         return parent.getScreenAt(index)
     }
 
+    override fun createShadowNodeInstance(context: ReactApplicationContext): LayoutShadowNode {
+        return ScreensShadowNode(context)
+    }
+
     override fun needsCustomLayoutForChildren(): Boolean {
         return true
     }

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

@@ -20,7 +20,6 @@ class ScreenStack(context: Context?) : ScreenContainer<ScreenStackFragment>(cont
     private val drawingOpPool: MutableList<DrawingOp> = ArrayList()
     private val drawingOps: MutableList<DrawingOp> = ArrayList()
     private var mTopScreen: ScreenStackFragment? = null
-    private var mNeedUpdate = false
     private val mBackStackListener = FragmentManager.OnBackStackChangedListener {
         if (mFragmentManager?.backStackEntryCount == 0) {
             // when back stack entry count hits 0 it means the user's navigated back using hw back
@@ -129,35 +128,7 @@ class ScreenStack(context: Context?) : ScreenContainer<ScreenStackFragment>(cont
         return super.hasScreen(screenFragment) && !mDismissed.contains(screenFragment)
     }
 
-    public override fun onScreenChanged() {
-        // we perform update in `onBeforeLayout` of `ScreensShadowNode` by adding an UIBlock
-        // which is called after updating children of the ScreenStack.
-        // We do it there because `onUpdate` logic requires all changes of children to be already
-        // made in order to provide proper animation for fragment transition.
-        // The exception to this rule is `updateImmediately` which is triggered by actions
-        // not connected to React view hierarchy changes, but rather internal events
-        mNeedUpdate = true
-    }
-
-    public override fun setFragmentManager(fm: FragmentManager) {
-        mFragmentManager = fm
-        performUpdatesNow()
-    }
-
-    private fun performUpdatesNow() {
-        // we want to update the immediately when the fragment manager is set or native back button
-        // dismiss is dispatched since it is not connected to React view hierarchy changes and will
-        // not trigger `onBeforeLayout` method of `ScreensShadowNode`
-        mNeedUpdate = true
-        performUpdates()
-    }
-
-    fun performUpdates() {
-        if (!mNeedUpdate || !mIsAttached || mFragmentManager == null) {
-            return
-        }
-        mNeedUpdate = false
-
+    override fun onUpdate() {
         // When going back from a nested stack with a single screen on it, we may hit an edge case
         // when all screens are dismissed and no screen is to be displayed on top. We need to gracefully
         // handle the case of newTop being NULL, which happens in several places below
@@ -309,7 +280,6 @@ class ScreenStack(context: Context?) : ScreenContainer<ScreenStackFragment>(cont
             it.commitNowAllowingStateLoss()
             mTopScreen?.let { screen -> setupBackHandlerIfNeeded(screen) }
         }
-        notifyContainerUpdate()
     }
 
     override fun notifyContainerUpdate() {

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

@@ -11,7 +11,7 @@ internal class ScreensShadowNode(private var mContext: ReactContext) : LayoutSha
         super.onBeforeLayout(nativeViewHierarchyOptimizer)
         (mContext.getNativeModule(UIManagerModule::class.java))?.addUIBlock { nativeViewHierarchyManager: NativeViewHierarchyManager ->
             val view = nativeViewHierarchyManager.resolveView(reactTag)
-            if (view is ScreenStack) {
+            if (view is ScreenContainer<*>) {
                 view.performUpdates()
             }
         }