Compose - get the same instance of ViewModel inside and outside of Navigation Graph

Question

The top-level composable in my Compose app is structured like this:

ModalBottomSheetLayout(/*...*/) {
    Scaffold(
        topBar = {
            when (currentScreen) {
                /*...*/
            }
        },
        content = {
            AppNavigation(navController)
        },
        bottomBar = {
            // Bottom navigation
        }
    )
}

Since I am using BottomSheet together with BottomNavigation I cannot delegate handling the former to the screens from AppNavigation because it will break Material Design guidelines (BottomSheet will be shown above the BottomNavigation).

Top bars of my screens are also separate from the screens themselves as per my other post

And of course since bottom sheet that is tied to a screen must reflect some state of this screen, we would like to share a ViewModel between them. Same goes for the top bar.

The above restrictions necessarily mean that the instance of a desired ViewModel must be created outside the scope of all mentioned composables and as such outside the NavHost and its scoped ViewModelStores, which is a huge problem because the only other ViewModelStore is the one owned by the Activity, which in a single-activity pattern is never cleared! So every shared ViewModel effectively becomes a singleton and a memory leak, aside from weird bugs it causes (sate that needed to be updated by creating a new instance but wasn't). What I actually want is for shared view models to be destroyed when user leaves the screen it is tied to, but it doesn't seem possible.

So is this even possible to fix this problem without changing the structure of the UI (and such dealing with other problems that led to it in the first place)? Does it even needs fixing? Are singleton view models OK provided state-related bugs are fixed?

Answer 1

After almost half a year of nobody answering this question the issue became critical in my project, so I started looking for solutions. After days of googling I was almost ready for a battle to the death with Android Framework through reflection to manually add and remove view models to/from Activity's ViewModelStore, but thankfully I found a one-liner solution.
I use Hilt in my project, so to get destination-scoped ViewModels I use the following extension function:

@Composable
inline fun <reified T : ViewModel> NavBackStackEntry.hiltViewModel() =
    ViewModelProvider(
        this.viewModelStore,
        HiltViewModelFactory(LocalContext.current, this)
    )[T::class.java]

private fun NavGraphBuilder.projectsGraph() {
    navigation(
        startDestination = AppTab.Projects.startDestination,
        route = AppTab.Projects.route
    ) {
        composable(AppScreen.Projects.route) {
            Projects()
        }

        composable(AppScreen.ProjectDetails.route) {
            ProjectDetails(projectViewModel = it.hiltViewModel()) // <--
        }
    }
}

Hilt creates a ViewModel instance based on the information inside a NavBackStackEntry object inside its own scoped ViewModelStore, and we can get this object through NavHostContoller.currentBackStackEntry!

I already had a CompositionLocal for NavHostContoller for convenience:

val LocalNavController = compositionLocalOf<NavHostController> {
    throw IllegalStateException("NavController does not yet exist")
}

CompositionLocalProvider(
    LocalNavController provides rememberNavController()
) {
    RootComposable()
}

So the only thing left to do is to create another function that gets a ViewModel created for a given NavBackStackEntry:

/**
 * Returns an existing [ViewModel] for the current navigation destination
 * regardless of where in the composition you want it :)
 */
@Composable
internal inline fun <reified VM: ViewModel> screenViewModel() =
    LocalNavController.current.currentBackStackEntryAsState().value?.hiltViewModel<VM>()

As a result, provided your LocalNavController provides to your app bars and bottom sheet, you can simply do the following and get the same instance of a ViewModel inside and outside of navigation graph.

@Composable
fun ProjectTopAppBar(
    viewModel: ProjectViewModel = screenViewModel()
) {
    viewModel ?: return
    // ...
}

Even better, it won't be a memory leak because it is still attached to a destination.