feat(navigation): hold tab state outside the back stack

apps/flipcash/app/src/main/kotlin/com/flipcash/app/internal/ui/navigation/AppContent.kt

@@ -12,8 +12,10 @@ import androidx.compose.foundation.layout.Box
 import androidx.compose.foundation.layout.fillMaxSize
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.LaunchedEffect
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
+import androidx.compose.runtime.snapshotFlow
 import androidx.compose.ui.Alignment
 import androidx.compose.ui.Modifier
 import androidx.compose.ui.unit.dp
@@ -28,8 +30,10 @@ import com.flipcash.app.cardexpand.LocalCardExpansion
 import com.flipcash.app.core.AppRoute
 import com.flipcash.app.core.navigation.DeeplinkAction
 import com.flipcash.app.core.navigation.LocalTabBarVisibility
+import com.flipcash.app.core.navigation.NavBarButton
 import com.flipcash.app.core.navigation.TabBarVisibilityController
 import com.flipcash.app.core.navigation.asNavBarTab
+import com.flipcash.app.core.navigation.destinationRoute
 import com.flipcash.app.core.ui.transitions.CardExpandTransition
 import com.flipcash.app.internal.ui.AppNavigationBar
 import com.flipcash.app.internal.ui.navigation.decorators.rememberNavBillOverlayEntryDecorator
@@ -38,6 +42,7 @@ import com.flipcash.app.internal.ui.navigation.decorators.rememberNavMessagingEn
 import com.flipcash.app.internal.ui.navigation.decorators.rememberNavTabBarInsetEntryDecorator
 import com.getcode.navigation.AppNavHost
 import com.getcode.navigation.core.CodeNavigator
+import com.getcode.navigation.decorators.rememberRetainedEntryState
 import com.getcode.navigation.results.NavResultStateRegistry
 import com.getcode.navigation.scenes.ModalBottomSheetSceneStrategy
 import com.getcode.ui.components.bars.BarManager
@@ -86,6 +91,24 @@ internal fun AppContent(
     // full screen in place, so there's no route change for the visibility rule below to notice.
     val tabBarVisibility = remember { TabBarVisibilityController() }
 
+    // A tab press replaces the whole back stack (tab-bar semantics — see AppNavigationBar), so every
+    // tab home was destroyed and rebuilt on each switch: the wallet re-fetched its balances and the
+    // chat list scrolled back to the top. Hold each tab home's ViewModels and scroll state outside the
+    // back stack so a switch shows what the tab last had. Only the four routes a tab press produces
+    // are held, so variants of a tab route (a resumed Tips, say) still get a fresh screen.
+    val tabHomeKeys = remember {
+        NavBarButton.entries.map { it.destinationRoute().toString() }.toSet()
+    }
+    val entryState = rememberRetainedEntryState { it in tabHomeKeys }
+
+    // Leaving the tabs altogether — signing out — is where held state stops being the user's own. A
+    // push keeps its tab home on the stack, so this fires on a replaceAll away from the tabs.
+    LaunchedEffect(entryState, codeNavigator.backStack) {
+        snapshotFlow {
+            codeNavigator.backStack.any { (it as? AppRoute)?.asNavBarTab() != null }
+        }.collect { onTabs -> if (!onTabs) entryState.releaseAll() }
+    }
+
     // Card-expand (iOS #587): the wallet requests an expansion (via LocalCardExpansion); the detail is
     // drawn by CardExpandHost inside the wallet entry, driven by one progress scalar, so the deck stays
     // composed and reorganises behind it. See CardExpansionController / CurrencyInfoExpansion.
@@ -101,6 +124,7 @@ internal fun AppContent(
             AppNavHost(
                 navigator = codeNavigator,
                 resultStateRegistry = resultStateRegistry,
+                entryState = entryState,
                 decorators = listOf(
                     // First = outermost, and outermost draws last: a bottom bar message is a prompt
                     // that has to be answered, so it sits above everything the entry draws — the

ui/navigation/build.gradle.kts

@@ -27,6 +27,7 @@ dependencies {
     implementation(libs.bundles.hilt)
     ksp(libs.bundles.hilt.compiler)
 
+    api(libs.compose.view.models)
     api(libs.navigation3.runtime)
     api(libs.navigation3.ui)
     api(libs.lifecycle.viewmodel.navigation3)

ui/navigation/src/main/kotlin/com/getcode/navigation/AppNavHost.kt

@@ -22,18 +22,18 @@ import androidx.compose.runtime.snapshots.Snapshot
 import androidx.compose.ui.graphics.Color
 import androidx.compose.ui.graphics.luminance
 import androidx.compose.ui.graphics.toArgb
-import androidx.lifecycle.viewmodel.navigation3.rememberViewModelStoreNavEntryDecorator
 import androidx.navigation3.runtime.NavEntry
 import androidx.navigation3.runtime.NavEntryDecorator
 import androidx.navigation3.runtime.NavKey
-import androidx.navigation3.runtime.rememberSaveableStateHolderNavEntryDecorator
 import androidx.navigation3.scene.Scene
 import androidx.navigation3.scene.SceneStrategy
 import androidx.navigation3.scene.SinglePaneSceneStrategy
 import androidx.navigation3.ui.NavDisplay
 import com.getcode.animation.LocalSharedTransitionScope
 import com.getcode.navigation.core.CodeNavigator
+import com.getcode.navigation.decorators.RetainedEntryState
 import com.getcode.navigation.decorators.rememberNavResultScopeEntryDecorator
+import com.getcode.navigation.decorators.rememberRetainedEntryState
 import com.getcode.navigation.results.NavResultStateRegistry
 import com.getcode.navigation.results.rememberNavResultStateRegistry
 import com.getcode.theme.CodeTheme
@@ -62,6 +62,10 @@ fun AppNavHost(
         popTransitionSpec()
     },
     onBack: (() -> Unit)? = null,
+    // Owns each entry's ViewModel store and rememberSaveable state. Defaults to retaining nothing,
+    // which is Nav3's own behaviour; a host passes one that retains keys whose state should survive
+    // their entry (the tab homes).
+    entryState: RetainedEntryState = rememberRetainedEntryState(),
     entryProvider: (key: NavKey) -> NavEntry<NavKey>,
     decorators: List<NavEntryDecorator<NavKey>> = emptyList(),
 ) {
@@ -102,9 +106,7 @@ fun AppNavHost(
         transitionSpec = transitionSpec,
         popTransitionSpec = popTransitionSpec,
         predictivePopTransitionSpec = predictivePopTransitionSpec,
-        entryDecorators = listOf(
-            rememberSaveableStateHolderNavEntryDecorator(),
-            rememberViewModelStoreNavEntryDecorator(),
+        entryDecorators = entryState.decorators + listOf(
             rememberNavResultScopeEntryDecorator(
                 backStack = navigator.backStack,
                 navResultStore = navigator.resultStore,

ui/navigation/src/main/kotlin/com/getcode/navigation/decorators/RetainedEntryState.kt

@@ -0,0 +1,150 @@
+package com.getcode.navigation.decorators
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.Stable
+import androidx.compose.runtime.remember
+import androidx.compose.runtime.saveable.SaveableStateHolder
+import androidx.compose.runtime.saveable.rememberSaveableStateHolder
+import androidx.lifecycle.viewmodel.ViewModelStoreProvider
+import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
+import androidx.lifecycle.viewmodel.compose.rememberViewModelStoreOwner
+import androidx.lifecycle.viewmodel.compose.rememberViewModelStoreProvider
+import androidx.navigation3.runtime.NavEntryDecorator
+import androidx.navigation3.runtime.NavKey
+import androidx.savedstate.compose.LocalSavedStateRegistryOwner
+
+/**
+ * Book-keeping for [RetainedEntryState]: which content keys keep their state after their entry is
+ * gone, and when that stops being true.
+ *
+ * Split out from the decorators because the ordering here is the whole difficulty, and none of it
+ * needs Compose to be exercised. A pop is reported only once the entry's content has left
+ * composition, so it lands a transition later than the back stack change that caused it — which
+ * means [release] routinely runs *before* the pops it is meant to discard. Marking the ledger
+ * released rather than only emptying it is what makes those late arrivals clear instead of being
+ * held onto for the next account.
+ *
+ * @param retains whether a content key is one whose state should outlive its entry.
+ */
+internal class RetentionLedger(private val retains: (contentKey: Any) -> Boolean) {
+    private val held = linkedSetOf<Any>()
+    private var released = false
+
+    /** A retainable entry is on screen again, so retention resumes. */
+    fun onRendered(contentKey: Any) {
+        if (retains(contentKey)) released = false
+    }
+
+    /** True when [contentKey]'s state should survive the pop that just happened. */
+    fun onPopped(contentKey: Any): Boolean {
+        val keep = !released && retains(contentKey)
+        if (keep) held += contentKey
+        return keep
+    }
+
+    /**
+     * Everything currently held, dropped. Pops still in flight are dropped too, until the next
+     * retainable entry renders.
+     */
+    fun release(): List<Any> {
+        released = true
+        val dropped = held.toList()
+        held.clear()
+        return dropped
+    }
+}
+
+/**
+ * Per-entry ViewModel and `rememberSaveable` state for a [androidx.navigation3.ui.NavDisplay], with
+ * an opt-in for the entries whose state should survive them.
+ *
+ * Nav3 scopes both to the entry and destroys both with it, which is right for a screen that was
+ * pushed and popped. It is wrong for the tab homes: a tab press replaces the whole back stack, so
+ * every tab came back cold — a fresh ViewModel with its default state, and a list scrolled back to
+ * the top. Retaining the entries on the back stack instead would fix that, but it also puts them in
+ * front of back: [androidx.navigation3.ui.NavDisplay] enables its back handler on
+ * `scene.previousEntries.isNotEmpty()`, so anything left underneath for state reasons is also
+ * something back walks through. Holding the state here keeps the back stack — and back — as they
+ * were.
+ *
+ * The mechanism is the one AndroidX documents on
+ * [androidx.lifecycle.viewmodel.compose.rememberViewModelStoreOwner]: a store outlives the
+ * composition that used it and is destroyed only by an explicit
+ * [ViewModelStoreProvider.clearKey]. These decorators are the stock pair with that call, and
+ * [SaveableStateHolder.removeState], skipped for the retained keys.
+ *
+ * State is held against the entry's `contentKey`, so a tab comes back to what it had only if it
+ * comes back under the same key.
+ */
+@Stable
+class RetainedEntryState internal constructor(
+    private val saveableStateHolder: SaveableStateHolder,
+    private val viewModelStoreProvider: ViewModelStoreProvider,
+    private val ledger: RetentionLedger,
+) {
+    /**
+     * Outermost first. Saveable state wraps the ViewModel store because a store's
+     * [androidx.lifecycle.SavedStateHandle] is restored through the saved state registry the
+     * saveable decorator provides.
+     */
+    val decorators: List<NavEntryDecorator<NavKey>> = listOf(
+        NavEntryDecorator(
+            onPop = { contentKey ->
+                if (!ledger.onPopped(contentKey)) saveableStateHolder.removeState(contentKey)
+            },
+            decorate = { entry ->
+                ledger.onRendered(entry.contentKey)
+                saveableStateHolder.SaveableStateProvider(entry.contentKey) { entry.Content() }
+            },
+        ),
+        NavEntryDecorator(
+            onPop = { contentKey ->
+                if (!ledger.onPopped(contentKey)) viewModelStoreProvider.clearKey(contentKey)
+            },
+            decorate = { entry ->
+                val owner = rememberViewModelStoreOwner(
+                    key = entry.contentKey,
+                    provider = viewModelStoreProvider,
+                    savedStateRegistryOwner = LocalSavedStateRegistryOwner.current,
+                )
+                CompositionLocalProvider(LocalViewModelStoreOwner provides owner) { entry.Content() }
+            },
+        ),
+    )
+
+    /**
+     * Drop everything being held, and stop holding until a retainable entry renders again.
+     *
+     * The caller decides when retention has stopped meaning anything — for the tab homes, when the
+     * user is no longer on the tabs at all. Without this the ViewModels of a signed-out account
+     * would be waiting for whoever signs in next.
+     */
+    fun releaseAll() {
+        ledger.release().forEach { contentKey ->
+            saveableStateHolder.removeState(contentKey)
+            viewModelStoreProvider.clearKey(contentKey)
+        }
+    }
+}
+
+/**
+ * A [RetainedEntryState] remembered across recompositions.
+ *
+ * @param retains whether an entry's `contentKey` names state that should outlive the entry.
+ *   Retains nothing by default, which is Nav3's own behaviour.
+ */
+@Composable
+fun rememberRetainedEntryState(
+    retains: (contentKey: Any) -> Boolean = { false },
+): RetainedEntryState {
+    val saveableStateHolder = rememberSaveableStateHolder()
+    val viewModelStoreProvider = rememberViewModelStoreProvider()
+    return remember(saveableStateHolder, viewModelStoreProvider) {
+        RetainedEntryState(
+            saveableStateHolder = saveableStateHolder,
+            viewModelStoreProvider = viewModelStoreProvider,
+            ledger = RetentionLedger(retains),
+        )
+    }
+}

ui/navigation/src/test/kotlin/com/getcode/navigation/decorators/RetentionLedgerTest.kt

@@ -0,0 +1,91 @@
+package com.getcode.navigation.decorators
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFalse
+import kotlin.test.assertTrue
+
+class RetentionLedgerTest {
+
+    private val tabs = setOf("Wallet", "Chats", "TipCard", "Scanner")
+
+    private fun ledger() = RetentionLedger { it in tabs }
+
+    @Test
+    fun `a retained key survives its pop`() {
+        assertTrue(ledger().onPopped("Wallet"))
+    }
+
+    @Test
+    fun `an unretained key does not`() {
+        assertFalse(ledger().onPopped("CurrencyInfo(usdc)"))
+    }
+
+    @Test
+    fun `release drops everything held`() {
+        val ledger = ledger()
+        ledger.onPopped("Wallet")
+        ledger.onPopped("Scanner")
+        ledger.onPopped("CurrencyInfo(usdc)")
+
+        assertEquals(listOf("Wallet", "Scanner"), ledger.release())
+    }
+
+    @Test
+    fun `release drops each key once`() {
+        val ledger = ledger()
+        ledger.onPopped("Wallet")
+        ledger.release()
+
+        assertEquals(emptyList(), ledger.release())
+    }
+
+    @Test
+    fun `switching between the same two tabs holds one entry each`() {
+        val ledger = ledger()
+        repeat(3) {
+            ledger.onPopped("Wallet")
+            ledger.onPopped("Scanner")
+        }
+
+        assertEquals(listOf("Wallet", "Scanner"), ledger.release())
+    }
+
+    // A pop is reported only once the entry's content has left composition, so on logout the pops of
+    // the tab homes land after the back stack has already stopped containing any. Retaining them then
+    // would leave the signed-out account's ViewModels waiting for whoever signs in next.
+    @Test
+    fun `a pop arriving after release is not retained`() {
+        val ledger = ledger()
+        ledger.release()
+
+        assertFalse(ledger.onPopped("Wallet"))
+    }
+
+    @Test
+    fun `a pop arriving after release is not held for a later release`() {
+        val ledger = ledger()
+        ledger.release()
+        ledger.onPopped("Wallet")
+
+        assertEquals(emptyList(), ledger.release())
+    }
+
+    @Test
+    fun `rendering a retained key resumes retention`() {
+        val ledger = ledger()
+        ledger.release()
+        ledger.onRendered("Wallet")
+
+        assertTrue(ledger.onPopped("Wallet"))
+    }
+
+    @Test
+    fun `rendering an unretained key does not resume retention`() {
+        val ledger = ledger()
+        ledger.release()
+        ledger.onRendered("CurrencyInfo(usdc)")
+
+        assertFalse(ledger.onPopped("Wallet"))
+    }
+}