feat(Android, Stack v5): allow for native navigation in nested stacks (software-mansion#3601)

## Description

This commit refactors *StackContainer* code. Achieved functional effect
boils down to support for native-pop in nested stack, and in outer stack
after closing the nested one.

Closes
software-mansion/react-native-screens-labs#821

## Changes

I've added abstraction over *FragmentManager* operations. This is done
mostly to structure and better manage container update code complexity.
The *FragmentManager* facing code is now moved into
`FragmentOperationExecutor` class.

Now *FragmentManager* reference is cleaned up in `onDetachedFromWindow`.

The *StackContainer* is now `OnBackStackChangedListener`. This is done
to better time & detect that a screen has been dismissed. This callback
is invoked mid-transaction execution by fragment manager, not after all
animations finish & the UI disappears. It seems like much more
appropriate place to keep the *StackContainer* possibly up-to-date.

Please note, that *StackScreen* still emits `onDismiss` event much
later, after the UI hides & fragment is destroyed. I've changed it a bit
here. Earlier it has been emitted in `onDestroyView`, now I emit it in
`onDestroy`. I've changed that to prevent incorrect emitting when
fragment moves to *STARTED* lifecycle state but is not detached from
fragment manager - such situation might happen when we use
`FragmentTransaction.detach`, which is planned.

This commit also adds basic *primaryNavigationFragment* management. This
is necessary for native navigation in nested stacks to work. If we do
not set it correctly, then *childFragmentManager* (used by nested stack)
won't handle `backPressed` at all. This responsibility will be delegated
to the *supportFragmentManager* and whole nested stack will be popped
immediately, instead of only single screen. Important thing to note here
is that we need to run additional fragment manager operations when we
detect that a fragment has been natively popped.
`FragmentManager.primaryNavigationFragment` state seems to not be
updated automatically once written to, hence unless we update it, we
will encounter a crash where *FragmentManager* attempts to delegate back
handling to already detached fragment.

## Visual documentation

Push, Pop, Native pop work nicely, even when popping nested stack. The
same action via JS popping should not work yet & will be handled
separately.


https://github.com/user-attachments/assets/43f6f3e4-35f7-4bfc-bd23-674d97ceca76

There is no "before" video. This is initial implementation.

## Test plan

TestStackNesting

## Checklist

- [ ] Included code example that can be used to test this change.
- [ ] Updated / created local changelog entries in relevant test files.
- [ ] For visual changes, included screenshots / GIFs / recordings
documenting the change.
- [ ] For API changes, updated relevant public types.
- [ ] Ensured that CI passes

Author: kkafar

android/src/main/java/com/swmansion/rnscreens/gamma/stack/host/FragmentOperation.kt

@@ -0,0 +1,69 @@
+package com.swmansion.rnscreens.gamma.stack.host
+
+import androidx.fragment.app.FragmentManager
+import com.swmansion.rnscreens.gamma.stack.screen.StackScreenFragment
+
+internal sealed class FragmentOperation {
+    internal abstract fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    )
+}
+
+internal class AddOp(
+    val fragment: StackScreenFragment,
+    val containerViewId: Int,
+    val addToBackStack: Boolean,
+    val allowStateLoss: Boolean = true,
+) : FragmentOperation() {
+    override fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    ) {
+        executor.executeAddOp(fragmentManager, this)
+    }
+}
+
+internal class PopBackStackOp(
+    val fragment: StackScreenFragment,
+) : FragmentOperation() {
+    override fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    ) {
+        executor.executePopBackStackOp(fragmentManager, this)
+    }
+}
+
+internal class RemoveOp(
+    val fragment: StackScreenFragment,
+    val allowStateLoss: Boolean = true,
+    val flushSync: Boolean = false,
+) : FragmentOperation() {
+    override fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    ) {
+        executor.executeRemoveOp(fragmentManager, this)
+    }
+}
+
+internal class FlushNowOp : FragmentOperation() {
+    override fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    ) {
+        executor.executeFlushOp(fragmentManager, this)
+    }
+}
+
+internal class SetPrimaryNavFragmentOp(
+    val fragment: StackScreenFragment,
+) : FragmentOperation() {
+    override fun execute(
+        fragmentManager: FragmentManager,
+        executor: FragmentOperationExecutor,
+    ) {
+        executor.executeSetPrimaryNavFragmentOp(fragmentManager, this)
+    }
+}

android/src/main/java/com/swmansion/rnscreens/gamma/stack/host/FragmentOperationExecutor.kt

@@ -0,0 +1,82 @@
+package com.swmansion.rnscreens.gamma.stack.host
+
+import androidx.fragment.app.FragmentManager
+import androidx.fragment.app.FragmentTransaction
+import com.swmansion.rnscreens.gamma.helpers.createTransactionWithReordering
+
+internal class FragmentOperationExecutor {
+    internal fun executeOperations(
+        fragmentManager: FragmentManager,
+        ops: List<FragmentOperation>,
+        flushSync: Boolean = false
+    ) {
+        ops.forEach { it.execute(fragmentManager, this) }
+
+        if (flushSync) {
+            fragmentManager.executePendingTransactions()
+        }
+    }
+
+    internal fun executeAddOp(fragmentManager: FragmentManager, op: AddOp) {
+        fragmentManager.createTransactionWithReordering().let { tx ->
+            tx.add(op.containerViewId, op.fragment)
+            if (op.addToBackStack) {
+                tx.addToBackStack(op.fragment.stackScreen.screenKey)
+            }
+            commitTransaction(tx, op.allowStateLoss)
+        }
+    }
+
+    internal fun executePopBackStackOp(fragmentManager: FragmentManager, op: PopBackStackOp) {
+        fragmentManager.popBackStack(
+            op.fragment.stackScreen.screenKey,
+            FragmentManager.POP_BACK_STACK_INCLUSIVE
+        )
+    }
+
+    internal fun executeRemoveOp(fragmentManager: FragmentManager, op: RemoveOp) {
+        fragmentManager.createTransactionWithReordering().let { tx ->
+            tx.remove(op.fragment)
+            commitTransaction(tx, op.allowStateLoss, op.flushSync)
+        }
+    }
+
+    internal fun executeFlushOp(fragmentManager: FragmentManager, op: FlushNowOp) {
+        fragmentManager.executePendingTransactions()
+    }
+
+    internal fun executeSetPrimaryNavFragmentOp(fragmentManager: FragmentManager, op: SetPrimaryNavFragmentOp) {
+        fragmentManager.createTransactionWithReordering().let { tx ->
+            tx.setPrimaryNavigationFragment(op.fragment)
+            commitTransaction(tx, allowStateLoss = true, flushSync = false)
+        }
+    }
+
+    private fun commitTransaction(
+        tx: FragmentTransaction,
+        allowStateLoss: Boolean,
+        flushSync: Boolean = false
+    ) {
+        if (flushSync) {
+            commitSync(tx, allowStateLoss)
+        } else {
+            commitAsync(tx, allowStateLoss)
+        }
+    }
+
+    private fun commitAsync(tx: FragmentTransaction, allowStateLoss: Boolean) {
+        if (allowStateLoss) {
+            tx.commitAllowingStateLoss()
+        } else {
+            tx.commit()
+        }
+    }
+
+    private fun commitSync(tx: FragmentTransaction, allowStateLoss: Boolean) {
+        if (allowStateLoss) {
+            tx.commitNowAllowingStateLoss()
+        } else {
+            tx.commitNow()
+        }
+    }
+}

android/src/main/java/com/swmansion/rnscreens/gamma/stack/host/StackContainer.kt

@@ -4,10 +4,10 @@ import android.annotation.SuppressLint
 import android.content.Context
 import android.util.Log
 import androidx.coordinatorlayout.widget.CoordinatorLayout
+import androidx.fragment.app.Fragment
 import androidx.fragment.app.FragmentManager
 import com.swmansion.rnscreens.gamma.helpers.FragmentManagerHelper
 import com.swmansion.rnscreens.gamma.helpers.ViewIdGenerator
-import com.swmansion.rnscreens.gamma.helpers.createTransactionWithReordering
 import com.swmansion.rnscreens.gamma.stack.screen.StackScreen
 import com.swmansion.rnscreens.gamma.stack.screen.StackScreenFragment
 import com.swmansion.rnscreens.utils.RNSLog
@@ -17,15 +17,17 @@ import java.lang.ref.WeakReference
 internal class StackContainer(
     context: Context,
     private val delegate: WeakReference<StackContainerDelegate>,
-) : CoordinatorLayout(context) {
+) : CoordinatorLayout(context),
+    FragmentManager.OnBackStackChangedListener {
     private var fragmentManager: FragmentManager? = null
 
+    private fun requireFragmentManager(): FragmentManager =
+        checkNotNull(fragmentManager) { "[RNScreens] Attempt to use nullish FragmentManager" }
+
     /**
      * Describes most up-to-date view of the stack. It might be different from
      * state kept by FragmentManager as this data structure is updated immediately,
      * while operations on fragment manager are scheduled.
-     *
-     * FIXME: In case of native-pop, this might be out of date!
      */
     private val stackModel: MutableList<StackScreenFragment> = arrayListOf()
 
@@ -34,6 +36,8 @@ internal class StackContainer(
     private val hasPendingOperations: Boolean
         get() = pendingPushOperations.isNotEmpty() || pendingPopOperations.isNotEmpty()
 
+    private val fragmentOpExecutor: FragmentOperationExecutor = FragmentOperationExecutor()
+
     init {
         id = ViewIdGenerator.generateViewId()
     }
@@ -42,16 +46,28 @@ internal class StackContainer(
         RNSLog.d(TAG, "StackContainer [$id] attached to window")
         super.onAttachedToWindow()
 
-        fragmentManager =
-            checkNotNull(FragmentManagerHelper.findFragmentManagerForView(this)) {
-                "[RNScreens] Nullish fragment manager - can't run container operations"
-            }
+        setupFragmentManger()
 
         // We run container update to handle any pending updates requested before container was
         // attached to window.
         performContainerUpdateIfNeeded()
     }
 
+    override fun onDetachedFromWindow() {
+        super.onDetachedFromWindow()
+        requireFragmentManager().removeOnBackStackChangedListener(this)
+        fragmentManager = null
+    }
+
+    internal fun setupFragmentManger() {
+        fragmentManager =
+            checkNotNull(FragmentManagerHelper.findFragmentManagerForView(this)) {
+                "[RNScreens] Nullish fragment manager - can't run container operations"
+            }.also {
+                it.addOnBackStackChangedListener(this)
+            }
+    }
+
     /**
      * Call this function to trigger container update
      */
@@ -60,9 +76,7 @@ internal class StackContainer(
         // the call because we don't have valid fragmentManager yet.
         // Update will be eventually executed in onAttachedToWindow().
         if (hasPendingOperations && isAttachedToWindow) {
-            val fragmentManager =
-                checkNotNull(fragmentManager) { "[RNScreens] Fragment manager was null during stack container update" }
-            performOperations(fragmentManager)
+            performOperations(requireFragmentManager())
         }
     }
 
@@ -75,64 +89,107 @@ internal class StackContainer(
     }
 
     private fun performOperations(fragmentManager: FragmentManager) {
-        // TODO: Handle case when we have pop & push of the same screen in single batch.
+        val fragmentOps = applyOperationsAndComputeFragmentManagerOperations()
+        fragmentOpExecutor.executeOperations(fragmentManager, fragmentOps, flushSync = false)
 
-        pendingPopOperations.forEach { performPopOperation(fragmentManager, it) }
-        pendingPushOperations.forEach { performPushOperation(fragmentManager, it) }
-
-        pendingPopOperations.clear()
-        pendingPushOperations.clear()
+        dumpStackModel()
     }
 
-    private fun performPushOperation(
-        fragmentManager: FragmentManager,
-        operation: PushOperation,
-    ) {
-        val transaction = fragmentManager.createTransactionWithReordering()
+    private fun applyOperationsAndComputeFragmentManagerOperations(): List<FragmentOperation> {
+        val fragmentOps = mutableListOf<FragmentOperation>()
 
-        val associatedFragment = StackScreenFragment(WeakReference(this), operation.screen)
-        stackModel.add(associatedFragment)
+        // Handle pop operations first.
+        // We don't care about pop/push duplicates, as long as we don't let the main loop progress
+        // before we commit all the transactions, FragmentManager will handle that for us.
 
-        transaction.add(this.id, associatedFragment)
+        pendingPopOperations.forEach { operation ->
+            val fragment =
+                checkNotNull(stackModel.find { it.stackScreen === operation.screen }) {
+                    "[RNScreens] Unable to find a fragment to pop"
+                }
 
-        // Don't add root screen to back stack to handle exiting from app.
-        if (fragmentManager.fragments.isNotEmpty()) {
-            transaction.addToBackStack(operation.screen.screenKey)
-        }
+            check(stackModel.size > 1) {
+                "[RNScreens] Attempt to pop last screen from the stack"
+            }
 
-        transaction.commitAllowingStateLoss()
-    }
+            fragmentOps.add(PopBackStackOp(fragment))
 
-    private fun performPopOperation(
-        fragmentManager: FragmentManager,
-        operation: PopOperation,
-    ) {
-        val associatedFragment = stackModel.find { it.stackScreen === operation.screen }
-        require(associatedFragment != null) {
-            "[RNScreens] Unable to find a fragment to pop."
+            check(stackModel.removeAt(stackModel.lastIndex) === fragment) {
+                "[RNScreens] Attempt to pop non-top screen"
+            }
+        }
+
+        pendingPushOperations.forEach { operation ->
+            val newFragment = createFragmentForScreen(operation.screen)
+            fragmentOps.add(
+                AddOp(
+                    newFragment,
+                    containerViewId = this.id,
+                    addToBackStack = stackModel.isNotEmpty(),
+                ),
+            )
+            stackModel.add(newFragment)
         }
 
-        val backStackEntryCount = fragmentManager.backStackEntryCount
-        if (backStackEntryCount > 0) {
-            fragmentManager.popBackStack(
-                operation.screen.screenKey,
-                FragmentManager.POP_BACK_STACK_INCLUSIVE,
+        check(stackModel.isNotEmpty()) { "[RNScreens] Stack should never be empty after updates" }
+
+        // Top fragment is the primary navigation fragment.
+        fragmentOps.add(SetPrimaryNavFragmentOp(stackModel.last()))
+
+        pendingPopOperations.clear()
+        pendingPushOperations.clear()
+
+        return fragmentOps
+    }
+
+    private fun onNativeFragmentPop(fragment: StackScreenFragment) {
+        Log.d(TAG, "StackContainer [$id] natively removed fragment ${fragment.stackScreen.screenKey}")
+        require(stackModel.remove(fragment)) { "[RNScreens] onNativeFragmentPop must be called with the fragment present in stack model" }
+        check(stackModel.isNotEmpty()) { "[RNScreens] Stack model should not be empty after a native pop" }
+
+        val fragmentManager = requireFragmentManager()
+        if (fragmentManager.primaryNavigationFragment === fragment) {
+            // We need to update the primary navigation fragment, otherwise the fragment manager
+            // will have invalid state, pointing to the dismissed fragment.
+            fragmentOpExecutor.executeOperations(
+                fragmentManager,
+                listOf(SetPrimaryNavFragmentOp(stackModel.last())),
             )
-        } else {
-            // When fast refresh is used on root screen, we need to remove the screen manually.
-            val transaction = fragmentManager.createTransactionWithReordering()
-            transaction.remove(associatedFragment)
-            transaction.commitNowAllowingStateLoss()
         }
+    }
 
-        stackModel.remove(associatedFragment)
+    private fun dumpStackModel() {
+        Log.d(TAG, "StackContainer [$id] MODEL BEGIN")
+        stackModel.forEach {
+            Log.d(TAG, "${it.stackScreen.screenKey}")
+        }
     }
 
-    internal fun onFragmentDestroyView(fragment: StackScreenFragment) {
-        if (stackModel.remove(fragment) && !fragment.stackScreen.isNativelyDismissed) {
-            Log.e(TAG, "[RNScreens] StackContainer natively popped a screen that was not in model!")
+    private fun createFragmentForScreen(screen: StackScreen): StackScreenFragment =
+        StackScreenFragment(screen).also {
+            Log.d(TAG, "Created Fragment $it for screen ${screen.screenKey}")
+        }
+
+    // This is called after special effects (animations) are dispatched
+    override fun onBackStackChanged() = Unit
+
+    // This is called before the special effects (animations) are dispatched, however mid transaction!
+    // Therefore make sure to not execute any action that might cause synchronous transaction synchronously
+    // from this callback.
+    override fun onBackStackChangeCommitted(
+        fragment: Fragment,
+        pop: Boolean,
+    ) {
+        if (fragment !is StackScreenFragment) {
+            Log.w(TAG, "[RNScreens] Unexpected type of fragment: ${fragment.javaClass.simpleName}")
+            return
+        }
+        if (pop) {
+            delegate.get()?.onScreenDismiss(fragment.stackScreen)
+            if (stackModel.contains(fragment)) {
+                onNativeFragmentPop(fragment)
+            }
         }
-        delegate.get()?.onScreenDismiss(fragment.stackScreen)
     }
 
     companion object {