A variation of FragmentViewModelLazy with added [SavedStateHandle] support

FragmentSavedStateViewModelLazy.kt

/**
 * A variation of [Fragment.viewModels] with added [SavedStateHandle] support
 */
@MainThread
inline fun <reified VM : ViewModel> Fragment.savedStateViewModels(
        noinline viewModelProducer: (SavedStateHandle) -> VM
) = createViewModelLazy(
        viewModelClass = VM::class,
        storeProducer = { this.viewModelStore },
        savedStateRegistryOwnerProducer = { this },
        viewModelProducer = viewModelProducer
)

/**
 * A variation of [Fragment.activityViewModels] with added [SavedStateHandle] support
 */
@MainThread
inline fun <reified VM : ViewModel> Fragment.savedStateActivityViewModels(
        noinline viewModelProducer: (SavedStateHandle) -> VM
) = createViewModelLazy(
        viewModelClass = VM::class,
        storeProducer = { requireActivity().viewModelStore },
        savedStateRegistryOwnerProducer = { requireActivity() },
        viewModelProducer = viewModelProducer
)

/**
 * Helper method for creation of [ViewModelLazy] with [SavedStateHandle] support.
 */
@MainThread
fun <VM : ViewModel> createViewModelLazy(
        viewModelClass: KClass<VM>,
        storeProducer: () -> ViewModelStore,
        savedStateRegistryOwnerProducer: () -> SavedStateRegistryOwner,
        viewModelProducer: (SavedStateHandle) -> VM
): ViewModelLazy<VM> {
    val factoryProducer = {
        object : AbstractSavedStateViewModelFactory(savedStateRegistryOwnerProducer(), null) {
            @Suppress("UNCHECKED_CAST")
            override fun <T : ViewModel?> create(key: String, modelClass: Class<T>, handle: SavedStateHandle) = viewModelProducer(handle) as T
        }
    }
    return ViewModelLazy(viewModelClass, storeProducer, factoryProducer)
}