Scope ViewModelStoreOwner to Composable lifecycle

VMScope.kt

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