From 995e1ab13dc2f681439970e1779100149fbeff3a Mon Sep 17 00:00:00 2001
From: Don Turner <donturner@google.com>
Date: Mon, 7 Jul 2025 15:14:01 +0100
Subject: [PATCH] Mirroring Nav2 back stack state

---
 .idea/codeStyles/Project.xml                  |   1 +
 app/build.gradle.kts                          |   1 +
 .../nowinandroid/navigation/NiaNavHost.kt     |  74 ++++-----
 .../samples/apps/nowinandroid/ui/NiaApp.kt    |  13 +-
 .../apps/nowinandroid/ui/NiaAppState.kt       | 142 ++++++++++++++++++
 .../src/main/kotlin/NavigationModule.kt       |  20 +++
 feature/bookmarks/api/build.gradle.kts        |   1 +
 .../api/src/main/kotlin/BookmarksRoutes.kt    |  21 +++
 feature/bookmarks/impl/build.gradle.kts       |   2 +-
 .../feature/bookmarks/impl/BookmarksScreen.kt |   2 +-
 10 files changed, 236 insertions(+), 41 deletions(-)
 create mode 100644 core/navigation/src/main/kotlin/NavigationModule.kt
 create mode 100644 feature/bookmarks/api/src/main/kotlin/BookmarksRoutes.kt

diff --git a/.idea/codeStyles/Project.xml b/.idea/codeStyles/Project.xml
index 6ef06e25a..3fa0ef569 100644
--- a/.idea/codeStyles/Project.xml
+++ b/.idea/codeStyles/Project.xml
@@ -12,6 +12,7 @@
       </option>
       <option name="IMPORT_LAYOUT_TABLE">
         <value>
+          <package name="" withSubpackages="true" static="false" module="true" />
           <package name="" withSubpackages="true" static="false" />
           <emptyLine />
           <package name="javax" withSubpackages="true" static="false" />
diff --git a/app/build.gradle.kts b/app/build.gradle.kts
index 9d43f184e..19acfb759 100644
--- a/app/build.gradle.kts
+++ b/app/build.gradle.kts
@@ -72,6 +72,7 @@ android {
 dependencies {
     implementation(projects.feature.interests)
     implementation(projects.feature.foryou)
+    implementation(projects.feature.bookmarks.api)
     implementation(projects.feature.bookmarks.impl)
     implementation(projects.feature.topic)
     implementation(projects.feature.search)
diff --git a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/navigation/NiaNavHost.kt b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/navigation/NiaNavHost.kt
index 21e9d3d2d..14d306c1b 100644
--- a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/navigation/NiaNavHost.kt
+++ b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/navigation/NiaNavHost.kt
@@ -19,11 +19,11 @@ package com.google.samples.apps.nowinandroid.navigation
 import androidx.compose.runtime.Composable
 import androidx.compose.ui.Modifier
 import androidx.navigation.compose.NavHost
-import androidx.navigation3.runtime.NavKey
+import androidx.navigation3.runtime.NavEntry
 import androidx.navigation3.runtime.entry
 import androidx.navigation3.runtime.entryProvider
-import androidx.navigation3.runtime.rememberNavBackStack
 import androidx.navigation3.ui.NavDisplay
+import com.google.samples.apps.nowinandroid.feature.bookmarks.impl.BookmarksScreen
 import com.google.samples.apps.nowinandroid.feature.bookmarks.impl.navigation.bookmarksScreen
 import com.google.samples.apps.nowinandroid.feature.foryou.navigation.ForYouBaseRoute
 import com.google.samples.apps.nowinandroid.feature.foryou.navigation.forYouSection
@@ -34,7 +34,6 @@ import com.google.samples.apps.nowinandroid.feature.topic.navigation.topicScreen
 import com.google.samples.apps.nowinandroid.navigation.TopLevelDestination.INTERESTS
 import com.google.samples.apps.nowinandroid.ui.NiaAppState
 import com.google.samples.apps.nowinandroid.ui.interests2pane.interestsListDetailScreen
-import kotlinx.serialization.Serializable
 
 /**
  * Top-level navigation graph. Navigation is organized as explained at
@@ -50,38 +49,41 @@ fun NiaNavHost(
     modifier: Modifier = Modifier,
 ) {
     val navController = appState.navController
-    val backStack = rememberNavBackStack(LegacyRoute)
-
-    NavDisplay(backStack = backStack, entryProvider = entryProvider {
-        entry<LegacyRoute> {
-            NavHost(
-                navController = navController,
-                startDestination = ForYouBaseRoute,
-                modifier = modifier,
-            ) {
-                forYouSection(
-                    onTopicClick = navController::navigateToTopic,
-                ) {
-                    topicScreen(
-                        showBackButton = true,
-                        onBackClick = navController::popBackStack,
-                        onTopicClick = navController::navigateToTopic,
-                    )
+    NavDisplay(
+        backStack = listOf(Unit),
+        onBack = { },
+        entryProvider = entryProvider(
+            fallback = { key ->
+                NavEntry(key = key) {
+                    NavHost(
+                        navController = navController,
+                        startDestination = ForYouBaseRoute,
+                        modifier = modifier,
+                    ) {
+                        forYouSection(
+                            onTopicClick = navController::navigateToTopic,
+                        ) {
+                            topicScreen(
+                                showBackButton = true,
+                                onBackClick = navController::popBackStack,
+                                onTopicClick = navController::navigateToTopic,
+                            )
+                        }
+                        bookmarksScreen(
+                            onTopicClick = navController::navigateToInterests,
+                            onShowSnackbar = onShowSnackbar,
+                        )
+                        searchScreen(
+                            onBackClick = navController::popBackStack,
+                            onInterestsClick = { appState.navigateToTopLevelDestination(INTERESTS) },
+                            onTopicClick = navController::navigateToInterests,
+                        )
+                        interestsListDetailScreen()
+                    }
                 }
-                bookmarksScreen(
-                    onTopicClick = navController::navigateToInterests,
-                    onShowSnackbar = onShowSnackbar,
-                )
-                searchScreen(
-                    onBackClick = navController::popBackStack,
-                    onInterestsClick = { appState.navigateToTopLevelDestination(INTERESTS) },
-                    onTopicClick = navController::navigateToInterests,
-                )
-                interestsListDetailScreen()
-            }
-        }
-    })
-}
+            },
+        ) {
 
-@Serializable
-data object LegacyRoute : NavKey
+        },
+    )
+}
diff --git a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaApp.kt b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaApp.kt
index 5a32d4585..06d52fefa 100644
--- a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaApp.kt
+++ b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaApp.kt
@@ -136,6 +136,9 @@ internal fun NiaApp(
 ) {
     val unreadDestinations by appState.topLevelDestinationsWithUnreadResources
         .collectAsStateWithLifecycle()
+
+    // TODO: This is only used to determine whether the bottom nav item is selected
+    // this can be improved by just getting
     val currentDestination = appState.currentDestination
 
     if (showSettingsDialog) {
@@ -144,6 +147,10 @@ internal fun NiaApp(
         )
     }
 
+    // TODO: figure out how have the existing navigation bar support both old and new routes
+    //  essentially we need a migration path from what's below to the common UI recipe
+    //  this probably comes down to how the back stack is modelled.
+
     NiaNavigationSuiteScaffold(
         navigationSuiteItems = {
             appState.topLevelDestinations.forEach { destination ->
@@ -167,9 +174,9 @@ internal fun NiaApp(
                     },
                     label = { Text(stringResource(destination.iconTextId)) },
                     modifier =
-                    Modifier
-                        .testTag("NiaNavItem")
-                        .then(if (hasUnread) Modifier.notificationDot() else Modifier),
+                        Modifier
+                            .testTag("NiaNavItem")
+                            .then(if (hasUnread) Modifier.notificationDot() else Modifier),
                 )
             }
         },
diff --git a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaAppState.kt b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaAppState.kt
index 5673a0bac..1c459aa7d 100644
--- a/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaAppState.kt
+++ b/app/src/main/kotlin/com/google/samples/apps/nowinandroid/ui/NiaAppState.kt
@@ -16,9 +16,11 @@
 
 package com.google.samples.apps.nowinandroid.ui
 
+import android.os.Bundle
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.Stable
 import androidx.compose.runtime.collectAsState
+import androidx.compose.runtime.mutableStateListOf
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
 import androidx.compose.runtime.rememberCoroutineScope
@@ -43,11 +45,13 @@ import com.google.samples.apps.nowinandroid.navigation.TopLevelDestination.BOOKM
 import com.google.samples.apps.nowinandroid.navigation.TopLevelDestination.FOR_YOU
 import com.google.samples.apps.nowinandroid.navigation.TopLevelDestination.INTERESTS
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
 import kotlinx.datetime.TimeZone
 
 @Composable
@@ -59,6 +63,11 @@ fun rememberNiaAppState(
     navController: NavHostController = rememberNavController(),
 ): NiaAppState {
     NavigationTrackingSideEffect(navController)
+
+    val nav3Navigator = remember(navController) {
+        Nav3NavigatorSimple(navController)
+    }
+
     return remember(
         navController,
         coroutineScope,
@@ -196,3 +205,136 @@ private fun NavigationTrackingSideEffect(navController: NavHostController) {
         }
     }
 }
+
+class Nav3NavigatorSimple(val navController: NavHostController){
+
+    val backStack = mutableStateListOf<Any>()
+
+    val coroutineScope = CoroutineScope(Job())
+
+    init {
+        coroutineScope.launch {
+            navController.currentBackStack.collect { nav2BackStack ->
+                println("Nav2 back stack changed")
+                backStack.clear()
+
+                for (nav2Entry in nav2BackStack){
+                    println("Adding destination: ${nav2Entry.destination}")
+                    backStack.add(nav2Entry.destination)
+                }
+            }
+        }
+    }
+}
+
+/*
+class Nav3Navigator<T: Any>(val navController: NavHostController, startRoute: T) {
+
+
+    init {
+        coroutineScope.launch {
+            navController.currentBackStack.collect { nav2BackStack ->
+                println("Nav2 back stack changed")
+                // TODO: Convert this into a nav3 back stack
+                for (nav2Entry in nav2BackStack){
+                    println("Destination: ${nav2Entry.destination}")
+                }
+            }
+        }
+
+        navController.addOnDestinationChangedListener(
+            listener = object : NavController.OnDestinationChangedListener {
+                override fun onDestinationChanged(
+                    controller: NavController,
+                    destination: NavDestination,
+                    arguments: Bundle?,
+                ) {
+                    println("NavController destination changed to $destination")
+                    // TODO: something! Or maybe we just listen to the back stack and mirror it here
+                }
+            }
+        )
+    }
+
+    // Keep track of the baseRoute - this is the route that is always at the bottom of the stack
+    private val baseRoute = startRoute
+
+    // Maintain a stack for each top level route
+    private var topLevelStacks : LinkedHashMap<T, SnapshotStateList<T>> = linkedMapOf(
+        baseRoute to mutableStateListOf(baseRoute)
+    )
+
+    // Expose the current top level route for consumers
+    var topLevelRoute by mutableStateOf(baseRoute)
+        private set
+
+    // Expose the back stack so it can be rendered by the NavDisplay
+    val backStack : SnapshotStateList<T> = mutableStateListOf(baseRoute)
+
+    private fun updateBackStack() {
+        backStack.apply {
+            // TODO: Could this be optimised?
+            clear()
+            addAll(topLevelStacks.flatMap { it.value })
+        }
+        println("Top level stacks: $topLevelStacks")
+        println("Backstack state: $backStack")
+    }
+
+    fun goTo(route: T, navOptions: NavOptions? = null){
+
+        backStack.add(route)
+        navController.navigate(route, navOptions)
+
+        *//*if (route is NavKey){
+            if (route is TopLevelRoute){
+                // Pop everything up to the base route stack
+                for (existingKey in topLevelStacks.keys.reversed()){
+                    if (existingKey != baseRoute) topLevelStacks.remove(existingKey)
+                }
+
+                if (route != baseRoute) {
+                    topLevelStacks.put(route, mutableStateListOf(route))
+                }
+                topLevelRoute = route
+
+            } else {
+                topLevelStacks[topLevelRoute]?.add(route)
+            }
+
+            topLevelStacks[topLevelRoute]?.add(route)
+
+            updateBackStack()
+        } else {
+            navController.navigate(route, navOptions)
+        }*//*
+    }
+
+*//*
+    fun removeLast(){
+        val removedKey = topLevelStacks[topLevelRoute]?.removeLastOrNull()
+        // If the removed key was a top level key, remove the associated top level stack
+        topLevelStacks.remove(removedKey)
+        topLevelRoute = topLevelStacks.keys.last()
+        updateBackStack()
+    }
+*//*
+
+    fun goBack(){
+        val removedKey = topLevelStacks[topLevelRoute]?.removeLastOrNull()
+        // If the removed key was a top level key, remove the associated top level stack
+        if (removedKey is TopLevelRoute){
+            topLevelStacks.remove(removedKey)
+            topLevelRoute = topLevelStacks.keys.last()
+        }
+
+        updateBackStack()
+
+        if (removedKey is LegacyRoute){
+            navController.popBackStack()
+        }
+    }
+
+
+
+}*/
diff --git a/core/navigation/src/main/kotlin/NavigationModule.kt b/core/navigation/src/main/kotlin/NavigationModule.kt
new file mode 100644
index 000000000..1086ac807
--- /dev/null
+++ b/core/navigation/src/main/kotlin/NavigationModule.kt
@@ -0,0 +1,20 @@
+/*
+ * Copyright 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.samples.apps.nowinandroid.core.navigation
+
+// Marker interface for top level routes
+// interface TopLevelRoute
\ No newline at end of file
diff --git a/feature/bookmarks/api/build.gradle.kts b/feature/bookmarks/api/build.gradle.kts
index 35287b9d7..6f7626dc7 100644
--- a/feature/bookmarks/api/build.gradle.kts
+++ b/feature/bookmarks/api/build.gradle.kts
@@ -24,5 +24,6 @@ android {
 }
 
 dependencies {
+    implementation(projects.core.navigation)
     implementation(libs.kotlinx.serialization.json)
 }
diff --git a/feature/bookmarks/api/src/main/kotlin/BookmarksRoutes.kt b/feature/bookmarks/api/src/main/kotlin/BookmarksRoutes.kt
new file mode 100644
index 000000000..2c13b3a61
--- /dev/null
+++ b/feature/bookmarks/api/src/main/kotlin/BookmarksRoutes.kt
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.samples.apps.nowinandroid.feature.bookmarks.api.navigation
+
+import androidx.navigation3.runtime.NavKey
+import kotlinx.serialization.Serializable
+
diff --git a/feature/bookmarks/impl/build.gradle.kts b/feature/bookmarks/impl/build.gradle.kts
index 07226c551..e1bff4470 100644
--- a/feature/bookmarks/impl/build.gradle.kts
+++ b/feature/bookmarks/impl/build.gradle.kts
@@ -26,7 +26,7 @@ android {
 
 dependencies {
     implementation(projects.core.data)
-
+    implementation(projects.core.navigation)
     testImplementation(projects.core.testing)
 
     androidTestImplementation(libs.bundles.androidx.compose.ui.test)
diff --git a/feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksScreen.kt b/feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksScreen.kt
index 2cf1c2b6e..5d6bc55ac 100644
--- a/feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksScreen.kt
+++ b/feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksScreen.kt
@@ -101,7 +101,7 @@ internal fun BookmarksRoute(
  */
 @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
 @Composable
-internal fun BookmarksScreen(
+fun BookmarksScreen(
     feedState: NewsFeedUiState,
     onShowSnackbar: suspend (String, String?) -> Boolean,
     removeFromBookmarks: (String) -> Unit,