Scope ViewModel to NavBackStackEntry jetpack compose

[susonthapa]

I am currently integrating jetpack compose into an existing project. For this, I have a fragment that hosts two composables. One is the category screen and the other is the products list present in that category as navigation destinations. This is the code within fragment's onCreateView.

        val navController = rememberNavController()
        NavHost(navController = navController, startDestination = Screen.AdminCategory.route) {
            composable(route = Screen.AdminCategory.route) {
                AdminCategoryScreen(scaffoldState = scaffoldState) {
                    val json = moshiInstance.adapter(Category::class.java).toJson(it)
                    navController.navigate(Screen.AdminCategoryProduct.createRoute(json))
                }
            }

            composable(
                route = Screen.AdminCategoryProduct.route,
                arguments = listOf(
                    navArgument(
                        Screen.AdminCategoryProduct.category
                    ) { type = CategoryParamType() }
                )
            ) {
                val category =
                    it.arguments?.getParcelable<Category>(Screen.AdminCategoryProduct.category)
                requireNotNull(category) { "category should be specified!" }
                AdminProductScreen(
                    category = category,
                    scaffoldState = scaffoldState,
                    it,
                    onBackClick = {
                        navController.navigateUp()
                    },
                )
            }
        }

In AdminProductScreen I'm getting the viewModel using this

    val vm: AdminProductVM = mavericksViewModel()

From the documentation

By default, mavericksViewModel() will get or create a view model scoped to the nearest LocalLifecycleOwner. In most cases, this will be the closest NavBackStackEntry, Fragment, or Activity

So I expect the viewModel to be cleared when the composable is popped but this is not the case. When I press the back button and navigate to another product page with a different category id, I can see the products from the previous category.

Looking at the source code for mavericksViewModel() and debugging it, the viewModelContext is FragmentViewModelContext.

@Composable
inline fun <reified VM : MavericksViewModel<S>, reified S : MavericksState> mavericksViewModel(
    scope: LifecycleOwner = LocalLifecycleOwner.current,
    noinline keyFactory: (() -> String)? = null,
    noinline argsFactory: (() -> Any?)? = null,
): VM {
    val activity = extractActivityFromContext(LocalContext.current)
    checkNotNull(activity) {
        "Composable is not hosted in a ComponentActivity!"
    }

    val viewModelStoreOwner = scope as? ViewModelStoreOwner ?: error("LifecycleOwner must be a ViewModelStoreOwner!")
    val savedStateRegistryOwner = scope as? SavedStateRegistryOwner ?: error("LifecycleOwner must be a SavedStateRegistryOwner!")
    val savedStateRegistry = savedStateRegistryOwner.savedStateRegistry
    val viewModelClass = VM::class
    val view = LocalView.current

    val viewModelContext = remember(scope, activity, viewModelStoreOwner, savedStateRegistry) {
        val parentFragment = findFragmentFromView(view)

        if (parentFragment != null) {
            val args = argsFactory?.invoke() ?: parentFragment.arguments?.get(Mavericks.KEY_ARG)
            FragmentViewModelContext(activity, args, parentFragment)
        } else {
            val args = argsFactory?.invoke() ?: activity.intent.extras?.get(Mavericks.KEY_ARG)
            ActivityViewModelContext(activity, args, viewModelStoreOwner, savedStateRegistry)
        }
    }
    return remember(viewModelClass, viewModelContext) {
        MavericksViewModelProvider.get(
            viewModelClass = viewModelClass.java,
            stateClass = S::class.java,
            viewModelContext = viewModelContext,
            key = keyFactory?.invoke() ?: viewModelClass.java.name
        )
    }
}

I have even tried passing the NavBackStackEntry as scope but still the result is the same. I don't see the viewModelStoreOwner being used in the FragmentViewModelContext. Am I missing something? Any help would be appreciated. Thank you.

I am using dagger for dependency injection.
In VM

    @AssistedFactory
    interface Factory : AssistedViewModelFactory<AdminProductVM, AdminProductState> {
        override fun create(state: AdminProductState): AdminProductVM
    }

    companion object :
        MavericksViewModelFactory<AdminProductVM, AdminProductState> by daggerMavericksViewModelFactory()

Dagger Module

    @Binds
    @IntoMap
    @MavericksViewModelKey(AdminProductVM::class)
    fun bindAdminProductVM(factory: AdminProductVM.Factory): AssistedViewModelFactory<*, *>

There are some changes in the navigation library with the recent updates. I'm currently using navigation 2.4.0-rc01 and compose 1.1.0-rc01.

[susonthapa]

After reading the documentation, I came up with this extension function.

@OptIn(InternalMavericksApi::class)
@Composable
inline fun <reified VM : MavericksViewModel<S>, reified S : MavericksState> composeMavericksViewModel(
    scope: LifecycleOwner = LocalLifecycleOwner.current,
    noinline keyFactory: (() -> String)? = null,
    noinline argsFactory: (() -> Any?)? = null,
): VM {
    val activity = extractActivityFromContext(LocalContext.current)
    checkNotNull(activity) {
        "Composable is not hosted in a ComponentActivity!"
    }

    val viewModelStoreOwner =
        scope as? ViewModelStoreOwner ?: error("LifecycleOwner must be a ViewModelStoreOwner!")
    val savedStateRegistryOwner = scope as? SavedStateRegistryOwner
        ?: error("LifecycleOwner must be a SavedStateRegistryOwner!")
    val savedStateRegistry = savedStateRegistryOwner.savedStateRegistry
    val viewModelClass = VM::class
    val view = LocalView.current

    val viewModelContext = remember(scope, activity, viewModelStoreOwner, savedStateRegistry) {
        val parentFragment = findFragmentFromView(view)

        if (parentFragment != null) {
            val args = argsFactory?.invoke() ?: parentFragment.arguments?.get(Mavericks.KEY_ARG)
            FragmentViewModelContext(
                activity,
                args,
                parentFragment,
                owner = viewModelStoreOwner,
                savedStateRegistry = savedStateRegistry
            )
        } else {
            val args = argsFactory?.invoke() ?: activity.intent.extras?.get(Mavericks.KEY_ARG)
            ActivityViewModelContext(activity, args, viewModelStoreOwner, savedStateRegistry)
        }
    }
    return remember(viewModelClass, viewModelContext) {
        MavericksViewModelProvider.get(
            viewModelClass = viewModelClass.java,
            stateClass = S::class.java,
            viewModelContext = viewModelContext,
            key = keyFactory?.invoke() ?: viewModelClass.java.name
        )
    }
}

The only change is this

            FragmentViewModelContext(
                activity,
                args,
                parentFragment,
                owner = viewModelStoreOwner,
                savedStateRegistry = savedStateRegistry
            )

Instead of using the owner and savedStateRegistry from the fragment, I'm using that from the scope. It's working for now, but I don't if this is the way to do this.