IssueTracker

Component used: navigation-compose Version used: 2.4.0-alpha05

Steps to reproduce:

1. Create a new compose project using Android Studio Arctic Fox rc01 wizard.
2. Replace the code below in MainActivity
3. Run the app, click the button "Go to B" then click the button "Go back".

Ref to Slack convo: https://kotlinlang.slack.com/archives/CJLTWPH7S/p1626963158280100?thread_ts=1626960240.270300&cid=CJLTWPH7S

Sample code to reproduce the issue:

class MainActivity : ComponentActivity() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContent {
            val navController = rememberNavController()
            NavHost(
                navController = navController,
                startDestination = "a",
                modifier = Modifier.fillMaxSize(),
                route = "mainGraph",
            ) {
                composable(route = "a") {
                    ScreenA(navToB = { navController.navigate("b") })
                }
                navigation(
                    startDestination = "b",
                    route = "subGraph",
                ) {
                    composable(route = "b") {
                        val vm = viewModel<SomeViewModel>(
                            viewModelStoreOwner = navController.getBackStackEntry("subGraph")
                        )
                        ScreenB(navBack = { navController.navigateUp() })
                    }
                }
            }
        }
    }
}

@Composable
fun ScreenA(
    navToB: () -> Unit
) {
    Column {
        Text("I am screen A")
        Button(onClick = navToB) {
            Text("Go to B")
        }
    }
}

@Composable
fun ScreenB(
    navBack: () -> Unit
) {
    Column {
        Text("I am screen B")
        Button(onClick = navBack) {
            Text("Go back")
        }
    }
}