ViewModel.kt

@MainThread
inline fun <reified VM : ViewModel> ComponentActivity.viewModel(
    noinline provider: (SavedStateHandle) -> VM
) = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { this },
    viewModelStoreOwnerProducer = { this },
    viewModelProvider = provider
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.viewModel(
    noinline provider: (SavedStateHandle) -> VM
) = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { this },
    viewModelStoreOwnerProducer = { this },
    viewModelProvider = provider
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.parentViewModel() = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { requireParentFragment() },
    viewModelStoreOwnerProducer = { requireParentFragment() },
    viewModelProvider = {
        throw IllegalStateException("${VM::class.java.name} should already exist scoped to the parent fragment.")
    }
)

@MainThread
inline fun <reified VM : ViewModel> Fragment.activityViewModel() = createLazyViewModel(
    viewModelClass = VM::class,
    savedStateRegistryOwnerProducer = { requireActivity() },
    viewModelStoreOwnerProducer = { requireActivity() },
    viewModelProvider = {
        throw IllegalStateException("${VM::class.java.name} should already exist scoped to the activity.")
    }
)

fun <VM : ViewModel> createLazyViewModel(
    viewModelClass: KClass<VM>,
    savedStateRegistryOwnerProducer: () -> SavedStateRegistryOwner,
    viewModelStoreOwnerProducer: () -> ViewModelStoreOwner,
    viewModelProvider: (SavedStateHandle) -> VM
) = ViewModelLazy(viewModelClass, { viewModelStoreOwnerProducer().viewModelStore }) {
    object : AbstractSavedStateViewModelFactory(savedStateRegistryOwnerProducer(), Bundle()) {
        @Suppress("UNCHECKED_CAST")
        override fun <T : ViewModel?> create(
            key: String,
            modelClass: Class<T>,
            handle: SavedStateHandle
        ): T {
            return viewModelProvider(handle) as T
        }
    }
}