SavedStateViewModelHelper.kt

package androidx.lifecycle

import android.os.Bundle
import androidx.fragment.app.Fragment
import androidx.fragment.app.FragmentActivity
import androidx.savedstate.SavedStateRegistryOwner

object ViewModelHelper {

  /**
   * [androidx.lifecycle.ViewModelProvider.DEFAULT_KEY]
   **/
  private const val DEFAULT_KEY = "androidx.lifecycle.ViewModelProvider.DefaultKey"


  private fun getProvider(
    activity: FragmentActivity,
    viewmodelFactory: ViewModelProvider.Factory
  ) = ViewModelProviders.of(activity, viewmodelFactory)

  private fun getProvider(
    activity: Fragment,
    viewmodelFactory: ViewModelProvider.Factory
  ) = ViewModelProviders.of(activity, viewmodelFactory)

  private fun <VIEWMODEL : ViewModel> getHandle(
    provider: ViewModelProvider,
    owner: SavedStateRegistryOwner,
    type: Class<VIEWMODEL>,
    defaultBundle: Bundle? = null
  ) = requireNotNull(provider.get(type)).also {
    if (it is SavedStateViewModel) {
      it.savedStateHandle = SavedStateHandleController.create(
        owner.savedStateRegistry,
        owner.lifecycle,
        "${DEFAULT_KEY}:${type.canonicalName}",
        defaultBundle
      ).handle
    }
  }

  fun <VIEWMODEL : ViewModel> getViewModel(
    activity: FragmentActivity,
    viewmodelFactory: ViewModelProvider.Factory,
    type: Class<VIEWMODEL>,
    defaultBundle: Bundle? = null
  ): VIEWMODEL = getHandle(getProvider(activity, viewmodelFactory), activity, type, defaultBundle)

  fun <VIEWMODEL : ViewModel> getViewModel(
    fragment: Fragment,
    viewmodelFactory: ViewModelProvider.Factory,
    type: Class<VIEWMODEL>,
    defaultBundle: Bundle? = null
  ): VIEWMODEL = getHandle(getProvider(fragment, viewmodelFactory), fragment, type, defaultBundle)
}


abstract class SavedStateViewModel : ViewModel() {
  lateinit var savedStateHandle: SavedStateHandle
    internal set
}