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

	}

}

So in this case the ViewModel will not survive scrolling: In a LazyColumn the Composable is "forgotten" on scrolling so the ViewModel would be cleared as well. If you have a composable in the view (but don't scroll), it would survive configuration changes and navigation changes, but not scrolling!
If you want it to survive scrolling, you would have to remember the viewModelStoreOwner just before your lazy list and then use that in the children, in which case all view models will be remembered as long as the lazy list is part of the composition. If you want to also take into consideration adding/removing items, then you'd have to do a little more magic on top