rememberViewModelStoreOwner example implementation

RememberViewModelStoreOwner.kt

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 viewModelStoreOwner = 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 viewModelStoreOwner
}

@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 viewModelStoreOwner: 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()
			viewModelStoreOwner.viewModelStore.clear()
		}

		fun update(key: Any?, lifecycle: Lifecycle) {
			if (key != this.key) {
				this.key = key
				viewModelStoreOwner.viewModelStore.clear()
			}
			attachedLifecycle.update { lifecycle }
		}

		fun attachComposable() {
			scope.launch { isAttachedToComposable.emit(true) }
		}

		fun detachComposable() {
			scope.launch { isAttachedToComposable.emit(false) }
		}

	}

}

RememberViewModelStoreOwnerHilt.kt

import androidx.compose.runtime.*
import androidx.lifecycle.*
import androidx.lifecycle.compose.LocalLifecycleOwner
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 currentViewModelStoreOwner = LocalViewModelStoreOwner.current
	val viewModelStoreOwnerViewModel = viewModel<ViewModelStoreOwnerViewModel> {
		ViewModelStoreOwnerViewModel()
	}
	val viewModelStoreOwner = 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 remember(currentViewModelStoreOwner, viewModelStoreOwner) {
		if (currentViewModelStoreOwner is HasDefaultViewModelProviderFactory) {
			object : ViewModelStoreOwner by viewModelStoreOwner,
				HasDefaultViewModelProviderFactory by currentViewModelStoreOwner {}
		} else {
			viewModelStoreOwner
		}
	}
}

@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.viewModelStoreOwner
	} ?: run {
		val attachedComposable = AttachedComposable(hashKey, key, lifecycle)
		attachedComposables += hashKey to attachedComposable
		attachedComposable.viewModelStoreOwner
	}

	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 viewModelStoreOwner: 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()
			viewModelStoreOwner.viewModelStore.clear()
		}

		fun update(key: Any?, lifecycle: Lifecycle) {
			if (key != this.key) {
				this.key = key
				viewModelStoreOwner.viewModelStore.clear()
			}
			attachedLifecycle.update { lifecycle }
		}

		fun attachComposable() {
			scope.launch { isAttachedToComposable.emit(true) }
		}

		fun detachComposable() {
			scope.launch { isAttachedToComposable.emit(false) }
		}

	}

}

For context about why I think your solution is very neat 😉

This is how Resaca is able to detect container is resumed after configuration change PlatformLifecycleHandler

Original implementation came from LeakCanary and this article.

@sebaslogen Thanks for having a look! If you see the latest revision, we do delete viewmodel store owners once their composable have been detached, using the dispose(hashKey) method. So in the worst case, if all children are removed, we still keep the parent ViewModel that stores all the composable's viewmodelstoreowners in memory, but the rest can definitely be garbage collected. So if you spawn 100 composables with their own viewmodelstoreowner, and you remove them as well, then they should be cleared and removed from the attachedComposables map.
The only reason it can be too long in memory, is for example when the activity is recreated and the composable at the same time is removed. In that case, the previous lifecycle owner is destroyed and we will not reattach and be able to check the onresume. Of course, this is a bit of an edge case and it is all cleaned up once the parent ViewModelStoreOwner is cleared as well.

Other than that, the race condition you mentioned is working because first an element will be attached (through the RememberObserver) and that happens before the lifecycle goes through onResume. I don't have any hard documentation on this :) but I can guarantee that it "is working on my device" :). The viewmodelScope is dispatching the work immediate so if compose keeps this order of first remembering elements before it gets resumed, then it should be fine.

Great, I see now all the additions to the latest revision 👍

Having one parent ViewModelStoreOwnerViewModel per screen is not an issue for me, same as in the backstack, that's the cheap price of using these nice features 😉

I totally trust you that you see it working. As for the dispatcher, AFAIK the ViewModel uses Dispatchers.Main, instead of Dispatchers.Main.Immediate, so unless there are really extreme race conditions (like the problems we all have with TextField), the contract of attach and then resume should still work.

@sebaslogen as far as I know that was a mistake (or by design but they changed their minds :) ), it was changed in lifecycle 2.2.0 to use the dispatchers.Main.immediate:
https://developer.android.com/jetpack/androidx/releases/lifecycle#2.2.0-alpha04
https://issuetracker.google.com/issues/139740492
And it is also in the current implementation the immediate dispatcher