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) { /* ... */ }
*/

What do you think about this? I'm storing ViewModelStoreOwner inside the parent-scoped ViewModel map. Every time the activity gets recreated I get the owner from the map.

@Composable
internal fun ProvidesViewModelStoreOwner(
    ownerKey: String = rememberSaveable { UUID.randomUUID().toString() },
    ownerHolder: CompositionScopedVMStoreOwnerHolder = viewModel(),
    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) {
                    ownerHolder.remove(ownerKey)
                }
            }
        }
    }

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

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

internal class CompositionScopedVMStoreOwnerHolder : ViewModel() {
    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 {
        storeOwnerMap[key]?.viewModelStore?.clear()
        storeOwnerMap.remove(key)
    }

    override fun onCleared() = lock.withLock {
        storeOwnerMap.forEach { it.value.viewModelStore.clear() }
        storeOwnerMap.clear()
        super.onCleared()
    }
}