Refactor conditional navigation recipe to make it saveable.

app/src/main/java/com/example/nav3recipes/conditional/ConditionalActivity.kt

@@ -22,45 +22,88 @@ import androidx.activity.compose.setContent
 import androidx.compose.foundation.layout.Column
 import androidx.compose.material3.Button
 import androidx.compose.material3.Text
+import androidx.compose.runtime.Composable
 import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateListOf
+import androidx.compose.runtime.setValue
 import androidx.compose.runtime.mutableStateOf
 import androidx.compose.runtime.remember
-import androidx.compose.runtime.setValue
+import androidx.compose.runtime.saveable.rememberSaveable
+import androidx.compose.runtime.saveable.rememberSerializable
+import androidx.navigation3.runtime.NavBackStack
+import androidx.navigation3.runtime.NavKey
 import androidx.navigation3.runtime.entryProvider
+import androidx.navigation3.runtime.serialization.NavBackStackSerializer
+import androidx.navigation3.runtime.serialization.NavKeySerializer
 import androidx.navigation3.ui.NavDisplay
 import com.example.nav3recipes.content.ContentBlue
 import com.example.nav3recipes.content.ContentGreen
 import com.example.nav3recipes.content.ContentYellow
+import kotlinx.serialization.Serializable
+
+
+/**
+ * Class for representing navigation keys in the app.
+ *
+ * Note: We use a sealed class because KotlinX Serialization handles
+ * polymorphic serialization of sealed classes automatically.
+ *
+ * @param requiresLogin - true if the navigation key requires that the user is logged in
+ * to navigate to it
+ */
+@Serializable
+sealed class ConditionalNavKey(val requiresLogin: Boolean = false) : NavKey
+
+/**
+ * Key representing home screen
+ */
+@Serializable
+private data object Home : ConditionalNavKey()
 
-private data object Home
+/**
+ * Key representing profile screen that is only accessible once the user has logged in
+ */
+@Serializable
+private data object Profile : ConditionalNavKey(requiresLogin = true)
 
-// A marker interface is used to mark any routes that require login
-private data object Profile : AppBackStack.RequiresLogin
-private data object Login
+/**
+ * Key representing login screen
+ *
+ * @param redirectToKey - navigation key to redirect to after successful login
+ */
+@Serializable
+private data class Login(
+    val redirectToKey: ConditionalNavKey? = null
+) : ConditionalNavKey()
 
 class ConditionalActivity : ComponentActivity() {
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
         setContent {
 
-
-            val appBackStack = remember {
-                AppBackStack(startRoute = Home, loginRoute = Login)
+            val backStack = rememberNavBackStack<ConditionalNavKey>(Home)
+            var isLoggedIn by rememberSaveable {
+                mutableStateOf(false)
+            }
+            val navigator = remember {
+                Navigator(
+                    backStack = backStack,
+                    onNavigateToRestrictedKey = { redirectToKey -> Login(redirectToKey) },
+                    isLoggedIn = { isLoggedIn }
+                )
             }
 
             NavDisplay(
-                backStack = appBackStack.backStack,
-                onBack = { appBackStack.remove() },
+                backStack = backStack,
+                onBack = { navigator.goBack() },
                 entryProvider = entryProvider {
                     entry<Home> {
-                        ContentGreen("Welcome to Nav3. Logged in? ${appBackStack.isLoggedIn}") {
+                        ContentGreen("Welcome to Nav3. Logged in? ${isLoggedIn}") {
                             Column {
-                                Button(onClick = { appBackStack.add(Profile) }) {
+                                Button(onClick = { navigator.navigate(Profile) }) {
                                     Text("Profile")
                                 }
-                                Button(onClick = { appBackStack.add(Login) }) {
+                                Button(onClick = { navigator.navigate(Login()) }) {
                                     Text("Login")
                                 }
                             }
@@ -69,16 +112,21 @@ class ConditionalActivity : ComponentActivity() {
                     entry<Profile> {
                         ContentBlue("Profile screen (only accessible once logged in)") {
                             Button(onClick = {
-                                appBackStack.logout()
+                                isLoggedIn = false
+                                navigator.navigate(Home)
                             }) {
                                 Text("Logout")
                             }
                         }
                     }
-                    entry<Login> {
-                        ContentYellow("Login screen. Logged in? ${appBackStack.isLoggedIn}") {
+                    entry<Login> { key ->
+                        ContentYellow("Login screen. Logged in? $isLoggedIn") {
                             Button(onClick = {
-                                appBackStack.login()
+                                isLoggedIn = true
+                                key.redirectToKey?.let { targetKey ->
+                                    backStack.remove(key)
+                                    navigator.navigate(targetKey)
+                                }
                             }) {
                                 Text("Login")
                             }
@@ -90,52 +138,14 @@ class ConditionalActivity : ComponentActivity() {
     }
 }
 
-/**
- * A back stack that can handle routes that require login.
- *
- * @param startRoute The starting route
- * @param loginRoute The route that users should be taken to when they attempt to access a route
- * that requires login
- */
-class AppBackStack<T : Any>(startRoute: T, private val loginRoute: T) {
-
-    interface RequiresLogin
-
-    private var onLoginSuccessRoute: T? = null
-
-    var isLoggedIn by mutableStateOf(false)
-        private set
-
-    val backStack = mutableStateListOf(startRoute)
-
-    fun add(route: T) {
-        if (route is RequiresLogin && !isLoggedIn) {
-            // Store the intended destination and redirect to login
-            onLoginSuccessRoute = route
-            backStack.add(loginRoute)
-        } else {
-            backStack.add(route)
-        }
-
-        // If the user explicitly requested the login route, don't redirect them after login
-        if (route == loginRoute) {
-            onLoginSuccessRoute = null
-        }
-    }
-
-    fun remove() = backStack.removeLastOrNull()
-
-    fun login() {
-        isLoggedIn = true
-
-        onLoginSuccessRoute?.let {
-            backStack.add(it)
-            backStack.remove(loginRoute)
-        }
-    }
 
-    fun logout() {
-        isLoggedIn = false
-        backStack.removeAll { it is RequiresLogin }
+// An overload of `rememberNavBackStack` that returns a subtype of `NavKey`.
+// See https://issuetracker.google.com/issues/463382671 for a discussion of this function
+@Composable
+fun <T : NavKey> rememberNavBackStack(vararg elements: T): NavBackStack<T> {
+    return rememberSerializable(
+        serializer = NavBackStackSerializer(elementSerializer = NavKeySerializer())
+    ) {
+        NavBackStack(*elements)
     }
 }

app/src/main/java/com/example/nav3recipes/conditional/Navigator.kt

@@ -0,0 +1,48 @@
+/*
+ * 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
+ *
+ *     http://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.example.nav3recipes.conditional
+
+import androidx.navigation3.runtime.NavBackStack
+
+/**
+ * Provides navigation events with built-in support for conditional access. If the user attempts to
+ * navigate to a [ConditionalNavKey] that requires login ([ConditionalNavKey.requiresLogin] is true)
+ * but is not currently logged in, the Navigator will redirect the user to a login key.
+ *
+ * @property backStack The back stack that is modified by this class
+ * @property onNavigateToRestrictedKey A lambda that is called when the user attempts to navigate
+ * to a key that requires login. This should return the key that represents the login screen. The
+ * user's target key is supplied as a parameter so that after successful login the user can be
+ * redirected to their target destination.
+ * @property isLoggedIn A lambda that returns whether the user is logged in.
+ */
+class Navigator(
+    private val backStack: NavBackStack<ConditionalNavKey>,
+    private val onNavigateToRestrictedKey: (targetKey: ConditionalNavKey?) -> ConditionalNavKey,
+    private val isLoggedIn: () -> Boolean,
+) {
+    fun navigate(key: ConditionalNavKey) {
+        if (key.requiresLogin && !isLoggedIn()) {
+            val loginKey = onNavigateToRestrictedKey(key)
+            backStack.add(loginKey)
+        } else {
+            backStack.add(key)
+        }
+    }
+
+    fun goBack() = backStack.removeLastOrNull()
+}