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

	}

}

@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