Composable scoped viewmodel

ViewModels.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.CompositionLocalProvider
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.remember
import androidx.compose.runtime.saveable.rememberSaveable
import androidx.compose.ui.platform.LocalContext
import androidx.lifecycle.ViewModelStore
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import timber.log.Timber
import java.util.UUID
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

@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)
    }
}

/* How to use it:
@Composable
fun SubscriptionPromptDialog(
    onDismissRequest: () -> Unit = {},
) {
    Dialog(
        onDismissRequest = onDismissRequest,
    ) {
        Surface(modifier.fillMaxSize()) {
            ProvidesViewModelStoreOwner { // Makes LocalViewModelStoreOwner.current return StoreOwner that is bound to Dialog lifecycle
                DialogContent()
            }
        }
    }
}

@Composable
@Preview
private fun DialogContent(
    modifier: Modifier = Modifier,
    viewModel: SubscriptionViewModel = koinViewModel(), // Uses the store owner provided by `ProvidesViewModelStoreOwner` that is bound to Dialog @Composable lifecycle
) { ... }
*/