Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save belinwu/c64d86129f25f0b2c5079f97470dee94 to your computer and use it in GitHub Desktop.
Save belinwu/c64d86129f25f0b2c5079f97470dee94 to your computer and use it in GitHub Desktop.
rememberViewModelStoreOwner example implementation
import androidx.compose.runtime.*
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.lifecycle.*
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.viewModel
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.*
@Composable
fun rememberViewModelStoreOwner(
key: Any?,
): ViewModelStoreOwner {
val viewModelKey = "rememberViewModelStoreOwner#" + currentCompositeKeyHash.toString(36)
val localLifecycle = LocalLifecycleOwner.current.lifecycle
val viewModelStoreOwnerViewModel = viewModel<ViewModelStoreOwnerViewModel> {
ViewModelStoreOwnerViewModel()
}
val viewModelStore = viewModelStoreOwnerViewModel.get(viewModelKey, key, localLifecycle)
remember {
object : RememberObserver {
override fun onAbandoned() {
viewModelStoreOwnerViewModel.detachComposable(viewModelKey)
}
override fun onForgotten() {
viewModelStoreOwnerViewModel.detachComposable(viewModelKey)
}
override fun onRemembered() {
viewModelStoreOwnerViewModel.attachComposable(viewModelKey)
}
}
}
return viewModelStore
}
@Composable
fun WithViewModelStoreOwner(
key: Any?,
content: @Composable () -> Unit,
) {
CompositionLocalProvider(
value = LocalViewModelStoreOwner provides rememberViewModelStoreOwner(key),
content = content,
)
}
@Stable
private class ViewModelStoreOwnerViewModel : ViewModel() {
private var attachedComposables = mapOf<String, AttachedComposable>()
fun get(hashKey: String, key: Any?, lifecycle: Lifecycle) = attachedComposables[hashKey]?.let {
it.update(key, lifecycle)
it.viewModelStore
} ?: run {
val attachedComposable = AttachedComposable(hashKey, key, lifecycle)
attachedComposables += hashKey to attachedComposable
attachedComposable.viewModelStore
}
fun attachComposable(hashKey: String) {
attachedComposables[hashKey]?.attachComposable()
}
fun detachComposable(hashKey: String) {
attachedComposables[hashKey]?.detachComposable()
}
override fun onCleared() {
attachedComposables.keys.forEach {
dispose(it)
}
super.onCleared()
}
private fun dispose(hashKey: String) {
attachedComposables[hashKey]?.let {
it.clear()
attachedComposables -= hashKey
}
}
private inner class AttachedComposable(
private val hashKey: String,
private var key: Any?,
initialLifecycle: Lifecycle,
) {
val viewModelStore: ViewModelStoreOwner = object : ViewModelStoreOwner {
override val viewModelStore = ViewModelStore()
}
private val supervisorJob: CompletableJob = SupervisorJob()
private val scope: CoroutineScope = viewModelScope + supervisorJob
private val attachedLifecycle = MutableStateFlow<Lifecycle?>(initialLifecycle)
private val isAttachedToComposable = MutableSharedFlow<Boolean>()
init {
scope.launch {
attachedLifecycle
.flatMapLatest { lifecycle -> lifecycle?.eventFlow ?: emptyFlow() }
.collectLatest {
if (it == Lifecycle.Event.ON_DESTROY) {
attachedLifecycle.update { null }
}
}
}
scope.launch {
isAttachedToComposable.collectLatest { isAttached ->
when {
// If we are attached or we are destroyed, we do not need to do anything
isAttached || attachedLifecycle.value == null -> return@collectLatest
// If we are detached and the lifecycle state is resumed, we should reset the view model store
attachedLifecycle.value?.currentState == Lifecycle.State.RESUMED -> {
dispose(hashKey)
}
else -> {
// We wait for the lifecycle event ON_RESUME to be triggered before resetting the ViewModelStore
// If in the mean time we are attached again, this work is cancelled
attachedLifecycle
.flatMapLatest { lifecycle -> lifecycle?.eventFlow ?: emptyFlow() }
// Wait for first event that matches ON_RESUME.
.firstOrNull { it == Lifecycle.Event.ON_RESUME } ?: return@collectLatest
dispose(hashKey)
}
}
}
}
}
fun clear() {
supervisorJob.cancelChildren()
viewModelStore.viewModelStore.clear()
}
fun update(key: Any?, lifecycle: Lifecycle) {
if (key != this.key) {
this.key = key
viewModelStore.viewModelStore.clear()
}
attachedLifecycle.update { lifecycle }
}
fun attachComposable() {
scope.launch { isAttachedToComposable.emit(true) }
}
fun detachComposable() {
scope.launch { isAttachedToComposable.emit(false) }
}
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment