Getting arg from ViewModel with savedStateHandle

composables.kt

// Both screens look the same, but now we have a ViewModel and some differences on the App composable
@HiltViewModel
class DetailsViewModel @Inject constructor(
  savedStateHandle: SavedStateHandle,
) : ViewModel() {
  private val state = MutableStateFlow(requireNotNull(savedStateHandle.get("arg")))
  fun state() = state.asStateFlow()
}

@Composable
fun App(
  modifier: Modifier = Modifier,
) {
  Box(
    modifier = modifier
      .fillMaxSize()
      .background(color = colorScheme.background),
  ) {
    val navController = rememberNavController()
    NavHost(navController = navController, startDestination = "start") {
      composable(route = "start") {
        StartScreen(
          onNavigateClick = { arg -> navController.navigate(route = "details/$arg") },
        )
      }
      // ---> interesting changes are here <---
      composable(route = "details/{arg}") { navBackStackEntry ->
        val arg = requireNotNull(navBackStackEntry.arguments?.getString("arg"))
        navBackStackEntry.savedStateHandle.set("arg", arg)
        val viewModel = hiltViewModel<DetailsViewModel>()
        val state by viewModel.state().collectAsStateWithLifecycle(lifecycleOwner = navBackStackEntry)
        DetailsScreen(arg = state.arg)
      }
    }
  }
}