Scope ViewModels to Composables

ProvideViewModels.kt

// PLEASE, READ
//
// This is a way to scope ViewModels to the Composition.
// However, this doesn't survive configuration changes or procress death on its own.
// You can handle all config changes in compose by making the activity handle those in the Manifest file
// e.g. android:configChanges="colorMode|density|fontScale|keyboard|keyboardHidden|layoutDirection|locale|mcc|mnc|navigation|orientation|screenLayout|screenSize|smallestScreenSize|touchscreen|uiMode">
//
// This is just an exploration to see what's possible in Compose. We don't encourage developers to copy-paste
// this code if they don't fully understand the implications of it and if this actually solves the use case to solve.

/*
 * Composition-aware ViewModelStoreOwner
 */
internal class CompositionScopedViewModelStoreOwner : ViewModelStoreOwner, RememberObserver {

    private val viewModelStore = ViewModelStore()

    override fun getViewModelStore(): ViewModelStore = viewModelStore

    override fun onAbandoned() {
        viewModelStore.clear()
    }

    override fun onForgotten() {
        viewModelStore.clear()
    }

    override fun onRemembered() {
        // Nothing to do here
    }
}

/*
 * Applies a [CompositionScopedViewModelStore] value to [LocalViewModelStoreOwner] 
 * to be able to scope ViewModels to a certain subtree of the Composition.
 *
 * Note: It's not a good idea to use `ProvideViewModels` at a screen level due to 
 * https://twitter.com/ianhlake/status/1395128325811494913. As when the store leaves
 * the Composition, all state and ViewModels are lost.
 * It's a good idea to use `ProvideViewModels` whenever it makes sense, for example,
 * before starting a certain flow of your app that have multiple screens with their own VM.
 * For example, a login flow, registration flow, or the main app flow.
 */
@Composable
fun ProvideViewModels(content: @Composable () -> Unit) {
    val viewModelStoreOwner = remember { CompositionScopedViewModelStoreOwner() }
    CompositionLocalProvider(LocalViewModelStoreOwner provides viewModelStoreOwner) {
        content()
    }
}


/*

Example of how you'd use this:

@Composable
fun MyApp() {
    // ...
    
    // Scope ViewModels in this part of the Composition.
    ProvideViewModels {
        val viewModel: SharedLoginViewModel = viewModel(factory = getSharedLoginViewModelFactory())
        LoginScreensFlow(viewModel)
    }
}

@Composable
fun LoginScreensFlow(viewModel: SharedLoginViewModel) { /* ... */ }
*/

Here is my solution to make it survive the configuration changes:

@Composable
internal fun ProvidesViewModelStoreOwner(
    ownerKey: String = rememberSaveable { UUID.randomUUID().toString() },
    content: @Composable () -> Unit,
) {
    val context = LocalContext.current

    remember {
        object : RememberObserver {
            override fun onRemembered() = Unit
            override fun onAbandoned() {
                clear()
            }

            override fun onForgotten() {
                clear()
            }

            private fun clear() {
                val isChangingConfigurations = context.findActivity().isChangingConfigurations
                if (!isChangingConfigurations) {
                    AppVMStoreOwnersHolder.remove(ownerKey)
                }
            }
        }
    }

    val viewModelStoreOwner: ViewModelStoreOwner = remember(ownerKey) {
        AppVMStoreOwnersHolder.getOwner(ownerKey)
    }

    CompositionLocalProvider(LocalViewModelStoreOwner provides viewModelStoreOwner) {
        content()
    }
}

internal object AppVMStoreOwnersHolder {
    private val lock = ReentrantLock()
    private val storeOwnerMap: MutableMap<String, ViewModelStoreOwner> = mutableMapOf()

    fun getOwner(key: String): ViewModelStoreOwner = lock.withLock {
        storeOwnerMap.getOrPut(key) {
            object : ViewModelStoreOwner {
                override val viewModelStore = ViewModelStore()
            }
        }
    }

    fun remove(key: String) = lock.withLock {
        Timber.d("Remove $key")
        storeOwnerMap[key]?.viewModelStore?.clear()
        storeOwnerMap.remove(key)
    }
}

/*
Example of how you'd use this:
@Composable
fun MyApp() {
    // ...
    
    ProvidesViewModelStoreOwner {
        LoginScreensFlow(viewModel)
    }
}
@Composable
fun LoginScreensFlow(viewModel: LoginViewModel = koinViewModel()) { /* ... */ }
*/

I'm not sure about this solution, but it helps scope view models without changing the LocalViewModelStoreOwner

Create a class that stores a map of ViewModelStore, which will be attached to a specific key, which is your scope
Extend ViewModel to make use of the onCleared(), which is called when LocalViewModelStoreOwner is cleared
Provide ViewModelStore for a key and clear it when it is no longer needed (scope destroyed), this will trigger clear of every ViewModel created by this store

class ViewModelStores : ViewModel() {
    val owner: ViewModelStoreOwner
        @Composable
        get() = checkNotNull(LocalViewModelStoreOwner.current) {
            "No ViewModelStoreOwner was provided via LocalViewModelStoreOwner"
        }

    private val stores = ConcurrentMap<String, ViewModelStore>()

    fun provideStore(key: String) = stores.getOrPut(key) {
        ViewModelStore()
    }

    fun clearStore(key: String) {
        val store = stores.remove(key)
        store?.clear()
    }

    override fun onCleared() {
        super.onCleared()
        stores.forEach { store ->
            store.value.clear()
        }

        stores.clear()
    }
}

Here's how I use it with koin:

@Composable
inline fun <reified T : IBaseViewModel, S : NavigationScreen> screenViewModel(
    screen: S,
    scope: Scope = currentKoinScope(),
    qualifier: Qualifier? = null
): T {
    val viewModelStores = koinViewModel<ViewModelStores>()
    val getViewModelClass: ((abstractClass: KClass<out T>) -> KClass<out ViewModel>) = koinInject()
    val concreteClass = getViewModelClass(T::class)

    val parameters = remember {
        parametersOf(screen)
    }
    
    LocalComposeNavigator.current.OnScreenRemoved(screen) {
        viewModelStores.clearStore(screen.key)
    }

    return resolveViewModel(
        concreteClass,
        viewModelStores.provideStore(screen.key),
        screen.key,
        defaultExtras(viewModelStores.owner),
        qualifier,
        scope,
        parameters = { parameters }
    ) as T
}

Basically, val viewModelStores = koinViewModel<ViewModelStores>() will always return the same viewModel (single), which is created by the main store owner (fragment or activity), which is safer than an object and complies to the lifecycle of the component.

Tested with Compose Multiplatform on both Android and iOS.