Navigation Compose scoped view models

[Tolriq]

Hi again :)

So I'm trying to have scoped viewModels to screens with hilt and this is quite an headache due to the current hilt limitations.

Do you have any plans to add support for something equivalent to hiltViewModel from compose navigation?

[KaustubhPatange]

This would've been so simple if hilt supports multiple viewModels based on different keys with which you can do something like this,

val viewmodel = viewModel<MyViewModel>(key = "value1")

And in other composable, you can provide a different key (say key = value2) which will give you different instances of the viewmodel. Is this the limitation you are talking about?

I have a workaround until hilt in the future supports this limitation but my workaround is not battle-tested. So what you would do is maintain an activity scope SavedStateViewModelFactory which you can pass through CompositionLocalProvider down the compose hierarchy & use the factory param on viewModel(key = ..., factory = Local**.current) to supply this factory with the key to create your viewModel. This way you can maintain multiple viewmodels of the same class.

This has one problem you cannot inject hilt dependency through this method, so you need to provide custom HiltViewModelFactory.

I also heard that the compose team is also going to provide a lifecycleOwner per composable? If this is true then you just need to do it when creating viewmodel which will give you multiple instances of the same viewmodel per composable.

[Tolriq]

Do you have some sample code? While I manage to get multiple viewmodels and proper scoping I can't get the SavedState to work :(

[KaustubhPatange]

If you use SavedStateViewModelFactory you will get additional parameter saveStateHandle in ViewModel. Is this you are talking about?

[Tolriq]

Yes but how do you integrate it is with hilt? HiltViewModelFactory.createInternal if using the activity SavedStateRegistryOwner you ends up with SavedStateProvider with the given key is already registered even when passing keys to the viewmodel.

The goal is to have scoped viewmodels that are removed and cleared when the composable is removed. Hence the issue here. Passing a key when using the activity scope is easy but for large applications I can't keep all viewmodels in memory.

Edit: So the goal is really to scope the viewmodels to navigation screens.

[KaustubhPatange]

Even if you are using navigation-compose from Jetpack the view models are still kept in memory even if the composables are abandoned. They are scoped to navigation graph i.e NavHost so until the NavHost exist they are still kept in memory which is mostly the case.

Scoping ViewModel to each composable in the real sense (not the Jetpack navigation way) is hard because there is no way to track the current composable lifecycle. That's why there is an ongoing issue about providing lifecycleOwner per composition scope.

If you really need scoping then Fragment is the way to go which itself is a lifecycleOwner so you can scope to it. I can't update the library to provide proper scoping since they will still remain in the memory like Jetpack navigation's implementation.

[Tolriq]

No the viewmodels are properly disposed when the screen is not on the backstack with jetpack navigation.

https://issuetracker.google.com/issues/165642391 does offer some way that do not support config change but in a full compose you can have compose handle the config change and then it's no more an issue.

My problem is just about how to properly glue the SavedStateRegistryOwner and the SavedStateHandle. I'm lost at trying to see how it's done in jetpack navigation for the moment :(

[KaustubhPatange]

I like the idea listed in the issue! We can make it to support config change basically we can wrap it in rememberSaveable to provide a custom saver (but a different one) for temporarily storing them into a static hashmap & consuming it on configuration change, this will also prevent memory leaks.

I can't help you with your problem I'll see if I can provide scoping in navigator-compose in future updates. If possible I'll let you know & if you find any solution to your problem let me know as well :)

[Tolriq]

Ok so here's a "working" solution for scoped view model with save states and hilt support. Vastly taken from the issue sample + hilt navigation + how rememberSaveable works :)

Require more testing and verifications but concept should be ok. (Except no support for config change but not a major issue in full compose apps with proper android:configChange in manifest) You can easily expand hiltViewModel to support a viewmodel level key too.

You can scope the viewmodel with

ViewModelLifecycleScope { }

And inject the viewmodel with a simple

val viewModel: ViewModel = hiltViewModel()

@Composable
inline fun <reified VM : ViewModel> hiltViewModel(
    viewModelStoreOwner: ViewModelStoreOwner = checkNotNull(LocalViewModelStoreOwner.current) {
        "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
    }
): VM {
    val factory = createHiltViewModelFactory(viewModelStoreOwner)
    val key = if (viewModelStoreOwner is ScopedViewModelStoreOwner) {
        viewModelStoreOwner.key
    } else null
    return viewModel(viewModelStoreOwner, key = key, factory = factory)
}

@Composable
@PublishedApi
internal fun createHiltViewModelFactory(
    viewModelStoreOwner: ViewModelStoreOwner
): ViewModelProvider.Factory? = if (viewModelStoreOwner is ScopedViewModelStoreOwner) {
    val activity = LocalContext.current.let {
        if (it is ComponentActivity) return@let it
        var ctx = it
        while (ctx is ContextWrapper) {
            if (ctx is ComponentActivity) {
                return@let ctx
            }
            ctx = ctx.baseContext
        }
        throw IllegalStateException(
            "Expected an activity context for creating a HiltViewModelFactory for a ScopedViewModelStoreOwner but instead found: $ctx"
        )
    }
    HiltViewModelFactory.createInternal(
        activity,
        viewModelStoreOwner,
        null,
        viewModelStoreOwner.defaultViewModelProviderFactory,
    )
} else {
    null
}

/**
 *
 * Does not support configuration change.
 * Ensure manifest contains: android:configChange with all options.
 *
 */
@Composable
fun ViewModelLifecycleScope(key: String? = null, content: @Composable () -> Unit) = CompositionLocalProvider(
    LocalViewModelStoreOwner provides rememberScopedViewModelStore(key = key),
    content = content
)

@Composable
internal fun rememberScopedViewModelStore(
    key: String? = null,
): ScopedViewModelStoreOwner {
    // key is the one provided by the user or the one generated by the compose runtime
    val finalKey = key.takeUnless { it.isNullOrEmpty() } ?: currentCompositeKeyHash.toString()

    val registry = LocalSaveableStateRegistry.current
    val ctx = LocalContext.current
    val value = remember {
        // TODO not restore when the input values changed (use hashKeys?) b/152014032
        ScopedViewModelStoreOwner(key, ctx, registry?.consumeRestored(finalKey) as Bundle?)
    }

    // re-register if the registry or key has been changed
    if (registry != null) {
        DisposableEffect(registry, finalKey, value) {
            val entry = registry.registerProvider(finalKey) {
                value.performSave(Bundle())
            }
            onDispose {
                entry.unregister()
            }
        }
    }
    return value
}

class ScopedViewModelStoreOwner(val key: String? = null, context: Context? = null, var savedState: Bundle?) :
    ViewModelStoreOwner, RememberObserver, HasDefaultViewModelProviderFactory, SavedStateRegistryOwner {

    private val viewModelStore = ViewModelStore()

    private var lifecycle = LifecycleRegistry(this)
    private val savedStateRegistryController = SavedStateRegistryController.create(this)

    private val defaultFactory by lazy {
        SavedStateViewModelFactory((context?.applicationContext as? Application), this, null)
    }

    init {
        lifecycle.currentState = Lifecycle.State.INITIALIZED
        savedStateRegistryController.performRestore(savedState)
        lifecycle.currentState = Lifecycle.State.CREATED
        lifecycle.currentState = Lifecycle.State.STARTED
    }

    override fun getViewModelStore(): ViewModelStore = viewModelStore

    override fun onAbandoned() {
        lifecycle.currentState = Lifecycle.State.CREATED
    }

    override fun onForgotten() {
        viewModelStore.clear()
        lifecycle.currentState = Lifecycle.State.DESTROYED
    }

    override fun onRemembered() {
        lifecycle.currentState = Lifecycle.State.RESUMED
    }

    override fun getDefaultViewModelProviderFactory(): ViewModelProvider.Factory {
        return defaultFactory
    }

    override fun getLifecycle(): Lifecycle {
        return lifecycle
    }

    fun performSave(outBundle: Bundle): Bundle {
        savedStateRegistryController.performSave(outBundle)
        return outBundle
    }

    override fun getSavedStateRegistry(): SavedStateRegistry {
        return savedStateRegistryController.savedStateRegistry
    }
}

[KaustubhPatange]

Look's great I thought of something similar with the library for eg: Each Setup or CreateDialog will provide a scoped lifecycleOwner through which you can create scoped viewModels, etc. The idea is roughly in my mind but let's see if this is possible or not.

[Tolriq]

No reason it would not work. But what about scoping to Setup and also scoping to reach Route ?

Scoping to setup is nice for like NavBar to keep the viewmodel around when changing screen that can happen often.

But scoping to route is also nice for deep navigation. There's no reason to keep the lower viewmodels when going back and navigating to another detail screen for example. You could ends up with tons of kept view models for no reasons or gains.

[KaustubhPatange]

My mistake I was saying scoping to Route for Setup and DialogRoute for CreateDialog.

They should work but again I want them to survive configuration change since even though we hypothetically say that in compose full app you should ignore all config change it's not generally true. For eg: if we have resources based on different locale or density the app will crash if you ignore them because now R.string points to different id which is not in the table anymore.

So yes we require configuration change & that's what I will aim to solve. If it works then the logic will be part of the library.

[Tolriq]

I'm sorry but I do not follow you.

Compose perfectly handle the config changes itself without impact and crashes and properly recompose the needed things when changing locale orientation without restarting the activity. (Same for the other things).

[KaustubhPatange]

I don't think so I tried with different strings.xml based on different locale it produces crash or sometime the same value of the string (on config change) not the updated one from different local strings.xml.

[Tolriq]

I use that in prod and just tested again I have no issue, you are probably accessing the strings in an edge case that require fixing rather than it being a global issue.

[KaustubhPatange]

an edge case This could be true but no problem in my code. I'll see what is causing this.

Also, if you are ignoring configChanges in manifest you don't even require viewmodels. You can use a service locator pattern to yourself manage some scoped instance of an object (that you use in place of viewmodels) that would be more better instead of providing a lifecycleOwner scoped to each Route. That's what simple-stack does in my opinion.

[Tolriq]

There's many things we can do manually ;) But then you are error prone, process death and save states are hard to manage + testing and everything.

And even with configChange the viewmodel is not kept but it's then the same as recreation after death so should be handled properly anyway.

Anyway this is not a very big deal I can reuse my current solution with your routes just need to wrap each route inside the scope. I hope you'll find a way to support all cases.

[KaustubhPatange]

Follow up,

I found this POC pull request on chrisbane's tivi app where he implemented a state holder like functionality (ViewModel) that will be canceled/removed when the backstack entry doesn't exist.

The idea is to use a coroutine scope from rememberCoroutineScope & track the state of this scope which will return a viewModelStore like object from where we can create our state holders (ViewModel). You can also use hilt to inject dependencies into this state holders using EntryPointAccessors.

When the entry from the backstack doesn't exist the scope containing ViewModelStore will be notified that one of its child scopes is canceled or became inactive which will eventually remove viewModel from the store.

It has the same issue that it will not survive configuration change but as you said in the full compose app you can use manifest's config change attribute to exclude all configuration changes. I like this idea because now you don't need ViewModel you can use any class as a state holder where you will be provided a coroutine scope through hilt's @AssistedInject (which is the same scope to track the state of this ViewModel, so if canceled the ViewModel is removed).

[Tolriq]

Yes I've seen this one but unfortunately no more SavedStateHandle too that are important for process death.

While I see many benefits of not using AAC ViewModel for KMM, I'm not sure the world is ready for that, there's really a lot of out of the box benefits from normal ViewModel and hilt injections without factories boilerplate.

[Tolriq]

Not related to this, but to avoid opening another one :) Any chance you can release current master with the breaking change as I'll start really implementing things and that would avoid a later refactor?

[KaustubhPatange]

Hey, why not use the snapshot versions for now https://github.com/KaustubhPatange/navigator/wiki/Using-a-Snapshot-Version? There are some changes I need to make before releasing it to the official maven.

[Tolriq]

Because snapshots are auto purged to the last X versions and I end up with not reproductible builds or forced update at not wanted timing. I do use snapshots to test stuff but not when building more "production" blocks.

Anyway I still have many other areas to play with so I can change my planning and delay the navigation part thanks for the update.

[KaustubhPatange]

If you properly hardcode a snapshot version eg: implementation "io.github.kaustubhpatange:navigator-compose:0.1-alpha21-2d3574d-SNAPSHOT" then new builds (to snapshots) will not be auto purged for you.

Anyway, I'll try to publish (to maven) as soon as possible.

[Tolriq]

What I mean is that Sonatype have auto purge of snapshot versions IIRC by default it's 10. So if you push 11 commits that snapshot build is not more available on their servers and I can't build anymore with that pinned version.

[KaustubhPatange]

Oh, I got it now what you meant but I don't think there is a hard limit like 10. From what I know it depends on the tagged version eg: 0.1-alpha21, if a certain number passes then the delete task is scheduled but that means I need to enable it first from the nexus admin panel, strange!

No worries I've released the version.