rememberViewModelStoreOwner example implementation

RememberViewModelStoreOwner.kt

import androidx.compose.runtime.Composable
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.Stable
import androidx.compose.runtime.currentCompositeKeyHash
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.hilt.navigation.compose.hiltViewModel
import androidx.lifecycle.Lifecycle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.ViewModelStore
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.eventFlow
import androidx.lifecycle.viewModelScope
import androidx.lifecycle.viewmodel.compose.viewModel
import kotlinx.coroutines.CompletableJob
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancelChildren
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.collectLatest
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.flow.firstOrNull
import kotlinx.coroutines.flow.flatMapLatest
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import kotlinx.coroutines.plus

@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.Composable
import androidx.compose.runtime.RememberObserver
import androidx.compose.runtime.Stable
import androidx.compose.runtime.currentCompositeKeyHash
import androidx.compose.runtime.remember
import androidx.compose.ui.platform.LocalLifecycleOwner
import androidx.hilt.navigation.compose.hiltViewModel
import androidx.lifecycle.HasDefaultViewModelProviderFactory
import androidx.lifecycle.Lifecycle
import androidx.lifecycle.ViewModel
import androidx.lifecycle.ViewModelProvider
import androidx.lifecycle.ViewModelStore
import androidx.lifecycle.ViewModelStoreOwner
import androidx.lifecycle.eventFlow
import androidx.lifecycle.viewModelScope
import androidx.lifecycle.viewmodel.compose.LocalViewModelStoreOwner
import androidx.lifecycle.viewmodel.compose.viewModel
import kotlinx.coroutines.CompletableJob
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.SupervisorJob
import kotlinx.coroutines.cancelChildren
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.MutableStateFlow
import kotlinx.coroutines.flow.collectLatest
import kotlinx.coroutines.flow.emptyFlow
import kotlinx.coroutines.flow.firstOrNull
import kotlinx.coroutines.flow.flatMapLatest
import kotlinx.coroutines.flow.update
import kotlinx.coroutines.launch
import kotlinx.coroutines.plus

@Composable
inline fun <reified VM : ViewModel> composableScopedHiltViewModel(
    key: Any? = null,
): VM {
    val viewModelStoreOwner = rememberViewModelStoreOwner(key)
    return hiltViewModel<VM>(viewModelStoreOwner, key = null)
}

@Composable
inline fun <reified VM : ViewModel, reified VMF> composableScopedHiltViewModel(
    key: Any? = null,
    noinline creationCallback: (VMF) -> VM,
): VM {
    val viewModelStoreOwner = rememberViewModelStoreOwner(key)
    return hiltViewModel<VM, VMF>(viewModelStoreOwner, key = null, creationCallback)
}

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

	}

}