refactor(navigation): hierarchical, route-dispatching CodeNavigator (Phase 1)

ui/navigation/build.gradle.kts

@@ -32,4 +32,7 @@ dependencies {
     api(libs.lifecycle.viewmodel.navigation3)
     api(libs.hilt.nav.compose)
     api(libs.rinku)
+
+    testImplementation(kotlin("test"))
+    testImplementation(libs.bundles.unit.testing)
 }

ui/navigation/src/main/kotlin/com/getcode/navigation/core/CodeNavigator.kt

@@ -12,6 +12,7 @@ import androidx.compose.runtime.staticCompositionLocalOf
 import androidx.navigation3.runtime.NavBackStack
 import androidx.navigation3.runtime.NavKey
 import com.getcode.navigation.Sheet
+import com.getcode.navigation.flow.FlowStep
 import com.getcode.navigation.results.NavResultStateRegistry
 import com.getcode.navigation.results.NavResultStore
 import com.getcode.navigation.results.NavResultStoreImpl
@@ -37,17 +38,33 @@ fun rememberCodeNavigator(
     backStack: NavBackStack<NavKey>,
     resultStateRegistry: NavResultStateRegistry,
     onRootReached: () -> Unit,
+    parent: CodeNavigator? = null,
+    isFlowNavigator: Boolean = false,
+    flowScope: FlowScope? = null,
 ): CodeNavigator {
     val navResultStore = rememberNavResultStore(resultStateRegistry = resultStateRegistry)
-    return remember(navResultStore, onRootReached) {
-        CodeNavigator(backStack = backStack, resultStore = navResultStore, onRootReached = onRootReached)
+    return remember(navResultStore, onRootReached, flowScope) {
+        CodeNavigator(
+            backStack = backStack,
+            resultStore = navResultStore,
+            onRootReached = onRootReached,
+            parent = parent,
+            isFlowNavigator = isFlowNavigator,
+            flowScope = flowScope,
+        )
     }
 }
 
-data class CodeNavigator(
+class CodeNavigator(
     val backStack: NavBackStack<NavKey>,
     val resultStore: NavResultStore,
     val onRootReached: () -> Unit,
+    /** The navigator this one is nested inside (the app-level navigator for a flow). Null at the app root. */
+    val parent: CodeNavigator? = null,
+    /** True when this navigator owns a flow's inner back stack (its stack holds [com.getcode.navigation.flow.FlowStep]s). */
+    val isFlowNavigator: Boolean = false,
+    /** Non-null only when this is a flow's inner navigator. See [FlowScope]. */
+    val flowScope: FlowScope? = null,
 ) {
     /**
      * When set, signals the active sheet to animate its dismiss. The callback
@@ -81,9 +98,29 @@ data class CodeNavigator(
     val currentRouteKey: NavKey?
         get() = backStack.lastOrNull()
 
+    /**
+     * Route-type-aware navigation. [com.getcode.navigation.flow.FlowStep]s land on the nearest
+     * flow stack; every other route bubbles up to the app-root stack. A screen calls this the same
+     * way whether it is top-level, in a flow, or in a sheet.
+     */
     fun navigate(
         route: NavKey,
         options: NavOptions = NavOptions(),
+    ) {
+        val belongsHere = if (route is FlowStep) isFlowNavigator else parent == null
+        when {
+            belongsHere -> navigateHere(route, options)
+            parent != null -> parent.navigate(route, options)
+            else -> trace(
+                "Dropped navigation to FlowStep $route: no flow host on the stack",
+                type = TraceType.Error,
+            )
+        }
+    }
+
+    private fun navigateHere(
+        route: NavKey,
+        options: NavOptions = NavOptions(),
     ) {
         if (options.debugRouting) {
             trace("Navigating to $route from ${backStack.lastOrNull()} with $options", type = TraceType.Navigation)
@@ -156,14 +193,14 @@ data class CodeNavigator(
     fun restoreRouting(routes: List<NavKey>) {
         val list = routes.toMutableList()
         val base = list.removeAt(0)
-        navigate(
+        navigateHere(
             route = base,
             options = NavOptions(
                 popUpTo = NavOptions.PopUpTo.ClearAll
             ),
         )
         list.forEach {
-            navigate(it)
+            navigateHere(it)
         }
     }
 
@@ -180,6 +217,41 @@ data class CodeNavigator(
         backStack.removeAt(backStack.lastIndex)
     }
 
+    /**
+     * Delivers [result] to the enclosing flow's caller by exiting via [FlowScope.exitWithResult].
+     * The flow host tears down the flow after receiving this signal.
+     * No-op (with a trace) when called outside a flow — there is no caller to return to.
+     */
+    fun navigateBackWithResult(result: android.os.Parcelable) {
+        val scope = flowScope
+        if (scope != null) {
+            scope.exitWithResult(result)
+        } else {
+            trace(
+                "navigateBackWithResult($result) called outside a flow scope; ignored",
+                type = TraceType.Error,
+            )
+        }
+    }
+
+    /**
+     * Leave the current bounded scope, regardless of depth:
+     *  - inside a flow -> delegates to [FlowScope.dismiss]; FlowHost animates the sheet out first
+     *    when the flow is a sheet root;
+     *  - a plain sheet -> animate the sheet out via [pendingSheetDismiss]; the scene pops it on completion;
+     *  - otherwise -> [navigateBack].
+     */
+    fun dismiss() {
+        val scope = flowScope
+        when {
+            scope != null -> scope.dismiss()
+            // Empty lambda = "animate the sheet out, no post-dismiss action"; the sheet scene
+            // observes pendingSheetDismiss, animates to Hidden, then pops the entry.
+            currentRouteKey is Sheet -> pendingSheetDismiss = {}
+            else -> navigateBack()
+        }
+    }
+
     fun <T : NavKey> clearToFirst(routeClass: KClass<T>): Boolean {
         Timber.d("Clearing backstack to first instance of $routeClass")
         var routeFound = false

ui/navigation/src/main/kotlin/com/getcode/navigation/core/FlowScope.kt

@@ -0,0 +1,29 @@
+package com.getcode.navigation.core
+
+import android.os.Parcelable
+
+/**
+ * The bounded navigation scope a [CodeNavigator] belongs to when it is a flow's inner navigator.
+ *
+ * A step never touches this directly — it calls CodeNavigator.navigateBackWithResult /
+ * CodeNavigator.dismiss, which delegate here. The implementation lives in
+ * [com.getcode.navigation.flow.FlowHost], where the flow-exit callback and enclosing-sheet
+ * context are already available. This keeps [CodeNavigator] free of flow/sheet mechanics.
+ */
+interface FlowScope {
+    /** True when this flow is the root content of a bottom sheet. */
+    val isSheetRoot: Boolean
+
+    /** Complete the flow, delivering [result] to whoever launched it. */
+    fun exitWithResult(result: Parcelable)
+
+    /**
+     * Leave the entire scope (the whole flow, animating the sheet out first when the flow is a
+     * sheet root). Equivalent to a user pressing the close-X.
+     *
+     * For a sheet-root flow, the enclosing sheet entry is removed from the outer backstack by the
+     * animation handler before the flow's exit callback runs — so the host's onExit must not pop
+     * the sheet entry again.
+     */
+    fun dismiss()
+}

ui/navigation/src/main/kotlin/com/getcode/navigation/flow/FlowHost.kt

@@ -40,6 +40,7 @@ import com.getcode.navigation.core.rememberCodeNavigator
 import com.getcode.navigation.results.NavResultOrCanceled
 import com.getcode.navigation.results.NavResultStateRegistry
 import com.getcode.navigation.results.asKey
+import com.getcode.navigation.core.FlowScope
 import com.getcode.navigation.scenes.LocalSheetNavigator
 
 /**
@@ -247,10 +248,43 @@ private fun <S : FlowStep, R : Parcelable> FlowHostImpl(
     // onRootReached and onExit read through rememberUpdatedState so the references
     // never change — preventing unnecessary recompositions of children that read
     // the composition locals.
+
+    // FlowScope exposes flow-exit + sheet-aware dismiss to the inner CodeNavigator, so steps can
+    // call navigator.navigateBackWithResult(...) / navigator.dismiss() instead of reaching for
+    // FlowNavigator or the outer navigator.
+    val currentSheetNavigator = rememberUpdatedState(sheetNavigator)
+    val flowScope = remember(isSheetRoot) {
+        object : FlowScope {
+            override val isSheetRoot: Boolean = isSheetRoot
+
+            override fun exitWithResult(result: Parcelable) {
+                @Suppress("UNCHECKED_CAST")
+                currentOnExit.value(FlowExitReason.Completed(result as R), isSheetRoot)
+            }
+
+            override fun dismiss() {
+                val nav = currentSheetNavigator.value
+                // The sheet scene removes the sheet entry from the outer backstack when the
+                // animation completes, then runs this lambda — so onExit must not pop the
+                // sheet entry again.
+                if (isSheetRoot && nav != null) {
+                    nav.pendingSheetDismiss = {
+                        currentOnExit.value(FlowExitReason.BackedOutOfRoot, isSheetRoot)
+                    }
+                } else {
+                    currentOnExit.value(FlowExitReason.BackedOutOfRoot, isSheetRoot)
+                }
+            }
+        }
+    }
+
     val innerNavigator = rememberCodeNavigator(
         backStack = innerBackStack,
         resultStateRegistry = resultStateRegistry,
         onRootReached = remember { { currentOnExit.value(FlowExitReason.BackedOutOfRoot, isSheetRoot) } },
+        parent = outerNavigator,
+        isFlowNavigator = true,
+        flowScope = flowScope,
     )
 
     val flowNavigator = remember(innerNavigator) {

ui/navigation/src/test/kotlin/com/getcode/navigation/TestNav.kt

@@ -0,0 +1,70 @@
+package com.getcode.navigation
+
+import android.os.Parcelable
+import androidx.navigation3.runtime.NavBackStack
+import androidx.navigation3.runtime.NavKey
+import com.getcode.navigation.core.CodeNavigator
+import com.getcode.navigation.core.EmptyCodeNavigator
+import com.getcode.navigation.core.FlowScope
+import com.getcode.navigation.flow.FlowRouteWithResult
+import com.getcode.navigation.flow.FlowStep
+import kotlinx.parcelize.Parcelize
+
+// --- App-level routes (NOT FlowSteps) ---
+data object AppHome : NavKey
+data object AppRegion : NavKey
+data object CallerScreen : NavKey
+
+// --- A sheet route ---
+data object DemoSheet : Sheet
+
+// --- Flow steps ---
+data object StepOne : FlowStep
+data object StepTwo : FlowStep
+
+// --- A flow-host route that returns DemoResult ---
+@Parcelize
+data class DemoResult(val value: String) : Parcelable
+
+data class DemoFlow(
+    override val initialStack: List<NavKey> = listOf(StepOne),
+) : FlowRouteWithResult<DemoResult>
+
+/** Records what a flow scope was asked to do, without any real flow/sheet plumbing. */
+class RecordingFlowScope(
+    override val isSheetRoot: Boolean = false,
+) : FlowScope {
+    val calls = mutableListOf<String>()
+    var deliveredResult: Parcelable? = null
+
+    override fun exitWithResult(result: Parcelable) {
+        deliveredResult = result
+        calls += "exitWithResult"
+    }
+
+    override fun dismiss() {
+        calls += "dismiss"
+    }
+}
+
+/** Builds a [CodeNavigator] over a real [NavBackStack] seeded with [keys]. */
+fun testNavigator(
+    vararg keys: NavKey,
+    parent: CodeNavigator? = null,
+    isFlow: Boolean = false,
+    scope: FlowScope? = null,
+    onRootReached: () -> Unit = {},
+): CodeNavigator {
+    require(keys.isNotEmpty()) { "seed the navigator with at least one key" }
+    val backStack = NavBackStack<NavKey>(keys.first()).apply {
+        keys.drop(1).forEach { add(it) }
+    }
+    return CodeNavigator(
+        backStack = backStack,
+        resultStore = EmptyCodeNavigator.resultStore,
+        onRootReached = onRootReached,
+        parent = parent,
+        isFlowNavigator = isFlow,
+        flowScope = scope,
+    )
+}

ui/navigation/src/test/kotlin/com/getcode/navigation/core/CodeNavigatorDispatchTest.kt

@@ -0,0 +1,78 @@
+package com.getcode.navigation.core
+
+import com.getcode.navigation.AppHome
+import com.getcode.navigation.AppRegion
+import com.getcode.navigation.StepOne
+import com.getcode.navigation.StepTwo
+import com.getcode.navigation.testNavigator
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class CodeNavigatorDispatchTest {
+
+    @Test
+    fun `flow step lands on the flow stack`() {
+        val root = testNavigator(AppHome)
+        val flow = testNavigator(StepOne, parent = root, isFlow = true)
+
+        flow.navigate(StepTwo)
+
+        assertEquals(listOf(StepOne, StepTwo), flow.backStack.toList())
+        assertEquals(listOf(AppHome), root.backStack.toList())
+    }
+
+    @Test
+    fun `app route from a flow bubbles up to the root stack`() {
+        val root = testNavigator(AppHome)
+        val flow = testNavigator(StepOne, parent = root, isFlow = true)
+
+        flow.navigate(AppRegion)
+
+        assertEquals(listOf(StepOne), flow.backStack.toList())
+        assertEquals(listOf(AppHome, AppRegion), root.backStack.toList())
+    }
+
+    @Test
+    fun `app route from a nested flow bubbles all the way to the root`() {
+        val root = testNavigator(AppHome)
+        val outerFlow = testNavigator(StepOne, parent = root, isFlow = true)
+        val innerFlow = testNavigator(StepOne, parent = outerFlow, isFlow = true)
+
+        innerFlow.navigate(AppRegion)
+
+        assertEquals(listOf(AppHome, AppRegion), root.backStack.toList())
+        assertEquals(listOf(StepOne), outerFlow.backStack.toList())
+        assertEquals(listOf(StepOne), innerFlow.backStack.toList())
+    }
+
+    @Test
+    fun `app route at the root lands on the root stack`() {
+        val root = testNavigator(AppHome)
+
+        root.navigate(AppRegion)
+
+        assertEquals(listOf(AppHome, AppRegion), root.backStack.toList())
+    }
+
+    @Test
+    fun `flow step at the root is dropped`() {
+        val root = testNavigator(AppHome)
+
+        root.navigate(StepTwo)
+
+        assertEquals(listOf(AppHome), root.backStack.toList())
+    }
+
+    @Test
+    fun `restoreRouting rebuilds this stack directly without type dispatch`() {
+        val root = testNavigator(AppHome)
+        val flow = testNavigator(StepOne, parent = root, isFlow = true)
+
+        // Even though these are FlowSteps and the parent exists, restoreRouting must
+        // rebuild the flow navigator's OWN stack, not bubble to the root.
+        flow.replaceAll(listOf(StepTwo, StepOne))
+
+        assertEquals(listOf(StepTwo, StepOne), flow.backStack.toList())
+        assertEquals(listOf(AppHome), root.backStack.toList())
+    }
+}

ui/navigation/src/test/kotlin/com/getcode/navigation/core/CodeNavigatorHarnessTest.kt

@@ -0,0 +1,21 @@
+package com.getcode.navigation.core
+
+import com.getcode.navigation.AppHome
+import com.getcode.navigation.AppRegion
+import com.getcode.navigation.testNavigator
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+class CodeNavigatorHarnessTest {
+
+    @Test
+    fun `push and pop mutate the backstack headlessly`() {
+        val nav = testNavigator(AppHome)
+
+        nav.navigate(AppRegion)
+        assertEquals(listOf(AppHome, AppRegion), nav.backStack.toList())
+
+        nav.navigateBack()
+        assertEquals(listOf(AppHome), nav.backStack.toList())
+    }
+}