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

	}

}

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