From db8d9c8126a694615c6fa965d230bdbec168d211 Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Wed, 9 Jan 2019 16:39:56 -0800
Subject: [PATCH 1/6] Allow creating initialState from MvRxFactory, give access
 to args in factory view model create method

---
 .../kotlin/com/airbnb/mvrx/MvRxExtensions.kt  |  82 +-----
 .../com/airbnb/mvrx/MvRxViewModelFactory.kt   |  85 +++++-
 .../com/airbnb/mvrx/MvRxViewModelProvider.kt  | 203 +++++++-------
 .../com/airbnb/mvrx/MvRxViewModelStore.kt     |  10 +-
 .../kotlin/com/airbnb/mvrx/FactoryTest.kt     | 265 +++++++++++++-----
 ...Test.kt => InitialConstructorStateTest.kt} |  14 +-
 .../com/airbnb/mvrx/MvRxExtensionsTest.kt     |  28 --
 .../features/dadjoke/DadJokeDetailFragment.kt |  10 +-
 .../features/dadjoke/DadJokeIndexViewModel.kt |  15 +-
 .../features/dadjoke/RandomDadJokeFragment.kt |  15 +-
 .../airbnb/mvrx/todomvrx/TasksViewModel.kt    |  14 +-
 11 files changed, 408 insertions(+), 333 deletions(-)
 rename mvrx/src/test/kotlin/com/airbnb/mvrx/{FragmentStateProviderTest.kt => InitialConstructorStateTest.kt} (74%)
 delete mode 100644 mvrx/src/test/kotlin/com/airbnb/mvrx/MvRxExtensionsTest.kt

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
index ed3c55340..47150e9ed 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
@@ -23,8 +23,7 @@ inline fun <T, reified VM : BaseMvRxViewModel<S>, reified S : MvRxState> T.fragm
     viewModelClass: KClass<VM> = VM::class,
     crossinline keyFactory: () -> String = { viewModelClass.java.name }
 ) where T : Fragment, T : MvRxView = lifecycleAwareLazy(this) {
-    val stateFactory: () -> S = ::_fragmentViewModelInitialStateProvider
-    MvRxViewModelProvider.get(viewModelClass.java, this, keyFactory(), stateFactory)
+    MvRxViewModelProvider.get(viewModelClass.java, S::class.java, FragmentViewModelContext(this.requireActivity(), _fragmentArgsProvider(), this), keyFactory())
         .apply { subscribe(this@fragmentViewModel, subscriber = { postInvalidate() }) }
 }
 
@@ -48,25 +47,29 @@ inline fun <T, reified VM : BaseMvRxViewModel<S>, reified S : MvRxState> T.activ
     viewModelClass: KClass<VM> = VM::class,
     noinline keyFactory: () -> String = { viewModelClass.java.name }
 ) where T : Fragment, T : MvRxView = lifecycleAwareLazy(this) {
-    val stateFactory: () -> S = { _activityViewModelInitialStateProvider(keyFactory) }
     if (requireActivity() !is MvRxViewModelStoreOwner) throw IllegalArgumentException("Your Activity must be a MvRxViewModelStoreOwner!")
-    MvRxViewModelProvider.get(viewModelClass.java, requireActivity(), keyFactory(), stateFactory)
+    MvRxViewModelProvider.get(viewModelClass.java, S::class.java, ActivityViewModelContext(requireActivity(), _activityArgsProvider(keyFactory)), keyFactory())
         .apply { subscribe(this@activityViewModel, subscriber = { postInvalidate() }) }
 }
 
 /**
  * For internal use only. Public for inline.
  *
- * Looks for [MvRx.KEY_ARG] on the arguments of the fragment receiver to create an instance of the State class.
- * The state class must have a matching single arg constructor.
+ * Looks for [MvRx.KEY_ARG] on the arguments of the fragments.
+ */
+@Suppress("FunctionName")
+fun <T : Fragment> T._fragmentArgsProvider(): Any? = arguments?.get(MvRx.KEY_ARG)
+
+/**
+ * For internal use only. Public for inline.
  *
- * If no MvRx fragment args exist it attempts to use a empty constructor. Otherwise an exception is thrown.
+ * Looks for [MvRx.KEY_ARG] on the arguments of the fragment receiver.
  *
  * Also adds the fragment's MvRx args to the host Activity's [MvRxViewModelStore] so that they can be used to recreate initial state
  * in a new process.
  */
 @Suppress("FunctionName")
-inline fun <reified S : MvRxState, T : Fragment> T._activityViewModelInitialStateProvider(keyFactory: () -> String): S {
+inline fun <T : Fragment> T._activityArgsProvider(keyFactory: () -> String): Any? {
     val args: Any? = arguments?.get(MvRx.KEY_ARG)
     val activity = requireActivity()
     if (activity is MvRxViewModelStoreOwner) {
@@ -74,65 +77,7 @@ inline fun <reified S : MvRxState, T : Fragment> T._activityViewModelInitialStat
     } else {
         throw IllegalArgumentException("Your Activity must be a MvRxViewModelStoreOwner!")
     }
-    return _initialStateProvider(S::class.java, args)
-}
-
-/**
- * For internal use only. Public for inline.
- *
- * Looks for [MvRx.KEY_ARG] on the arguments of the fragment receiver to create an instance of the State class.
- * The state class must have a matching single arg constructor.
- *
- * If no MvRx fragment args exist it attempts to use a empty constructor. Otherwise an exception is thrown.
- */
-@Suppress("FunctionName")
-inline fun <reified S : MvRxState, T : Fragment> T._fragmentViewModelInitialStateProvider(): S {
-    val args: Any? = arguments?.get(MvRx.KEY_ARG)
-    return _initialStateProvider(S::class.java, args)
-}
-
-/**
- * For internal use only. Public for inline.
- *
- * Looks for [MvRx.KEY_ARG] in intent extras on activity receiver to create an instance of the State class.
- * The state class must have a matching single arg constructor.
- *
- * If no MvRx activity args exist it attempts to use a empty constructor. Otherwise an exception is thrown.
- */
-@Suppress("FunctionName")
-inline fun <reified S : MvRxState, T : FragmentActivity> T._activityViewModelInitialStateProvider(): S {
-    val args: Any? = intent.extras?.get(MvRx.KEY_ARG)
-    return _initialStateProvider(S::class.java, args)
-}
-
-/**
- * For internal use only. Public for inline.
- *
- * Searches [stateClass] for a single argument constructor matching the type of [args]. If [args] is null, then
- * no arg constructor is invoked.
- *
- */
-@Suppress("FunctionName") // Public for inline.
-fun <S : MvRxState> _initialStateProvider(stateClass: Class<S>, args: Any?): S {
-    val argsConstructor = args?.let {
-        val argType = it::class.java
-
-        stateClass.constructors.firstOrNull { constructor ->
-            constructor.parameterTypes.size == 1 && isAssignableTo(argType, constructor.parameterTypes[0])
-            }
-        }
-
-    @Suppress("UNCHECKED_CAST")
-    return argsConstructor?.newInstance(args) as? S
-        ?: try {
-            stateClass.newInstance()
-        } catch (@Suppress("TooGenericExceptionCaught") e: Throwable) {
-            null
-        }
-        ?: throw IllegalStateException(
-            "Attempt to auto create the MvRx state class ${stateClass.simpleName} has failed. It must have default values for every property or a " +
-                "secondary constructor for ${args?.javaClass?.simpleName ?: "a fragment argument"}. "
-        )
+    return args
 }
 
 /**
@@ -143,8 +88,7 @@ inline fun <T, reified VM : BaseMvRxViewModel<S>, reified S : MvRxState> T.viewM
     crossinline keyFactory: () -> String = { viewModelClass.java.name }
 ) where T : FragmentActivity,
         T : MvRxViewModelStoreOwner = lifecycleAwareLazy(this) {
-    val stateFactory: () -> S = { _activityViewModelInitialStateProvider() }
-    MvRxViewModelProvider.get(viewModelClass.java, this, keyFactory(), stateFactory)
+    MvRxViewModelProvider.get(viewModelClass.java, S::class.java, ActivityViewModelContext(this, intent.extras?.get(MvRx.KEY_ARG)), keyFactory())
 }
 
 /**
diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
index 95aab246a..1eeb21d42 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
@@ -4,26 +4,85 @@ import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
 
 /**
- * Implement this in the companion object of a MvRxViewModel if your ViewModel needs more dependencies than just initial state.
- * If all you need is initial state, you don't need to implement this at all.
+ * Implement this on your ViewModel's companion object for hooks into state creation and ViewModel
+ * creation. See [MvRxViewModelArgFactory] for a version which automatically types your argument
+ * class.
  */
-interface MvRxViewModelFactory<S : MvRxState> {
+interface MvRxViewModelFactory<VM: BaseMvRxViewModel<S>, S : MvRxState> {
+
     /**
-     * This will be called when your ViewModel needs to created. This *needs* to be annotated with [JvmStatic].
-     * @param state: The initial state for your ViewModel. This will be populated from activity persisted state.
+     * @param viewModelContext [ViewModelContext] which contains the ViewModel owner and arguments.
+     * @param state The initial state to pass to the ViewModel. In a new process, state will have all [PersistState] annotated members restored,
+     * therefore you should never create a custom state in this method. To customize the initial state, override [initialState].
+     *
+     * @return The ViewModel. If you return `null` the ViewModel must have a single argument
+     * constructor only taking the initial state.
      */
-    fun create(activity: FragmentActivity, state: S): BaseMvRxViewModel<S>
+    fun create(viewModelContext: ViewModelContext, state: S): VM? = null
+
+    /**
+     * The initial state for the ViewModel. Override this if the initial state requires information from arguments, or the ViewModel owner.
+     * This function will take precedence over any secondary constructors defined in the state class, [S].
+     *
+     * The return value of this function will be transformed with any [PersistState] values before being used in [create].
+     *
+     * @return the initial state. If `null`, the state class constructors will be used for initial state creation.
+     */
+    fun initialState(viewModelContext: ViewModelContext): S? = null
+
 }
 
 /**
- * Implement this in the companion object of a MvRxViewModel if your ViewModel needs more dependencies than just initial state.
- * If you only need the [FragmentActivity] then you can use the [MvRxViewModelFactory].
- * If all you need is initial state, you don't need to implement this at all.
+ * Creation context for the ViewModel. Includes the ViewModel store owner (either an activity or fragment), and fragment arguments
+ * set via [MvRx.KEY_ARG].
+ *
+ * For activity scoped ViewModels see [ActivityViewModelContext].
+ * For fragment scoped ViewModels see [FragmentViewModelContext].
+ *
+ * Never store a reference to this context, or the activity/fragment.
+ *
  */
-interface MvRxFragmentViewModelFactory<S : MvRxState> {
+sealed class ViewModelContext {
+    /**
+     * The activity which is using the ViewModel.
+     */
+    abstract val activity: FragmentActivity
+    /**
+     * Fragment arguments set via [MvRx.KEY_ARG].
+     */
+    abstract val rawArgs: Any?
+
     /**
-     * This will be called when your ViewModel needs to created. This *needs* to be annotated with [JvmStatic].
-     * @param state: The initial state for your ViewModel. This will be populated from fragment args and persisted state.
+     * Convenience method to type [rawArgs].
      */
-    fun create(fragment: Fragment, state: S): BaseMvRxViewModel<S>
+    abstract fun <A> args(): A
+}
+
+/**
+ * The [ViewModelContext] for a ViewModel created with an
+ * activity scope (`val viewModel by activityViewModel<MyViewModel>`).
+ */
+class ActivityViewModelContext(
+        override val activity: FragmentActivity,
+        override val rawArgs: Any?
+) : ViewModelContext() {
+    override fun <A> args(): A {
+        @Suppress("UNCHECKED_CAST")
+        return rawArgs as A
+    }
+}
+
+/**
+ * The [ViewModelContext] for a ViewModel created with a
+ * fragment scope (`val viewModel by fragmentViewModel<MyViewModel>`).
+ */
+class FragmentViewModelContext(
+        override val activity: FragmentActivity,
+        override val rawArgs: Any?,
+        val fragment: Fragment
+) : ViewModelContext() {
+    override fun <A> args(): A {
+        @Suppress("UNCHECKED_CAST")
+        return rawArgs as A
+    }
 }
diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
index 5f14775c5..6bf761975 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
@@ -1,11 +1,9 @@
 package com.airbnb.mvrx
 
 import android.arch.lifecycle.ViewModelProviders
-import android.arch.lifecycle.ViewModelStoreOwner
 import android.support.annotation.RestrictTo
 import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
-import kotlin.reflect.full.companionObjectInstance
 import kotlin.reflect.full.primaryConstructor
 
 /**
@@ -18,130 +16,89 @@ object MvRxViewModelProvider {
      * MvRx specific ViewModelProvider used for creating a BaseMvRxViewModel scoped to either a [Fragment] or [FragmentActivity].
      * If this is in a [Fragment], it cannot be called before the Fragment has been added to an Activity or wrapped in a [Lazy] call.
      *
-     * @param viewModelClass The class of the ViewModel you would like an instance of
-     * @param storeOwner Either a [Fragment] or [FragmentActivity] to be the scope owner of the ViewModel. Activity scoped ViewModels
-     *                   can be used to share state across Fragments.
+     * @param viewModelClass The class of the ViewModel you would like an instance of.
+     * @param stateClass The class of the State used by the ViewModel.
+     * @param viewModelContext The [ViewModelContext] which contains arguments and the owner of the ViewModel.
+     *                         Either [ActivityViewModelContext] or [FragmentViewModelContext].
      * @param key An optional key for the ViewModel in the store. This is optional but should be used if you have multiple of the same
      *            ViewModel class in the same scope.
-     * @param stateFactory A factory to create the initial state if the ViewModel does not yet exist.
-     *
      */
     fun <VM : BaseMvRxViewModel<S>, S : MvRxState> get(
         viewModelClass: Class<VM>,
-        storeOwner: ViewModelStoreOwner,
-        key: String = viewModelClass.name,
-        stateFactory: () -> S
+        stateClass: Class<S>,
+        viewModelContext: ViewModelContext,
+        key: String = viewModelClass.name
     ): VM {
         // This wraps the fact that ViewModelProvider.of has individual methods for Fragment and FragmentActivity.
-        val activityOwner = storeOwner as? FragmentActivity
-        val fragmentOwner = storeOwner as? Fragment
-
-        val factory = MvRxFactory {
-            when {
-                activityOwner != null -> createViewModel(viewModelClass, activityOwner, stateFactory())
-                fragmentOwner != null -> createViewModel(viewModelClass, fragmentOwner, stateFactory())
-                else -> throw IllegalArgumentException("$storeOwner must either be an Activity or a Fragment that is attached to an Activity")
-            }
-        }
-        return when {
-            activityOwner != null -> ViewModelProviders.of(activityOwner, factory)
-            else -> ViewModelProviders.of(fragmentOwner!!, factory)
+        val factory = MvRxFactory { createViewModel(viewModelClass, stateClass, viewModelContext) }
+        return when (viewModelContext) {
+            is ActivityViewModelContext -> ViewModelProviders.of(viewModelContext.activity, factory)
+            is FragmentViewModelContext -> ViewModelProviders.of(viewModelContext.fragment, factory)
         }.get(key, viewModelClass)
     }
 
     @Suppress("UNCHECKED_CAST")
-    fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createViewModel(
-        viewModelClass: Class<VM>,
-        fragmentActivity: FragmentActivity,
-        state: S
-    ): VM = createFactoryViewModel(viewModelClass, fragmentActivity, state)
-        .let { validateViewModel(it, viewModelClass, state) }
-
-    @Suppress("UNCHECKED_CAST")
-    fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createViewModel(
-        viewModelClass: Class<VM>,
-        fragment: Fragment,
-        state: S
-    ): VM = createFactoryViewModel(viewModelClass, fragment, state)
-        .let { validateViewModel(it, viewModelClass, state) }
-
-    private fun <VM : BaseMvRxViewModel<S>, S : MvRxState> validateViewModel(
-        viewModel: VM?,
-        viewModelClass: Class<VM>,
-        state: S
+    internal fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createViewModel(
+        viewModelClass: Class<out VM>,
+        stateClass: Class<out S>,
+        viewModelContext: ViewModelContext,
+        stateRestorer: (S) -> S = { it }
     ): VM {
+        val initialState = createInitialState(viewModelClass, stateClass, viewModelContext, stateRestorer)
+        val companionClass = try { Class.forName("${viewModelClass.name}\$Companion") } catch (exception: ClassNotFoundException) { null }
+        val factoryViewModel = companionClass?.let {
+            try {
+                companionClass.getMethod("create", ViewModelContext::class.java, MvRxState::class.java)
+                    .invoke(companionClass.instance(), viewModelContext, initialState) as VM?
+            } catch (exception: NoSuchMethodException) {
+                // Check for JvmStatic method.
+                viewModelClass.getMethod("create", ViewModelContext::class.java, MvRxState::class.java)
+                    .invoke(null, viewModelContext, initialState) as VM?
+            }
+        }
+        val viewModel = factoryViewModel ?: createDefaultViewModel(viewModelClass, initialState)
         return requireNotNull(viewModel) {
-
-            // We are about to crash, so accessing Kotlin reflect is okay for a better error message.
-            when {
-                viewModelClass.kotlin.companionObjectInstance is MvRxViewModelFactory<*> -> {
-                    "${viewModelClass.simpleName} companion " +
-                        "${MvRxViewModelFactory::class.java.simpleName} is missing ${JvmStatic::class.java.name} " +
-                        "annotation on its create method."
-                }
-                viewModelClass.kotlin.companionObjectInstance != null -> {
-                    "${viewModelClass.simpleName} must have primary constructor with a single " +
-                        "parameter for initial state of ${state::class.java} or a companion object " +
-                        "implementing ${MvRxViewModelFactory::class.java} and a ${JvmStatic::class.java.simpleName} " +
-                        "annotated create method. Found a companion object which does not " +
-                        "implement ${MvRxViewModelFactory::class.java.simpleName}."
-                }
-                viewModelClass.kotlin.primaryConstructor?.parameters?.size?.let { it > 1 } == true -> {
-                    "${viewModelClass.simpleName} takes dependencies other than initialState. " +
-                        "It must have companion object implementing ${MvRxViewModelFactory::class.java.simpleName} " +
-                        "and a ${JvmStatic::class.java.simpleName} annotated create method."
-                }
-                viewModelClass.kotlin.primaryConstructor?.parameters?.size?.let { it == 0 } == true -> {
-                    "${MvRxViewModelFactory::class.java.simpleName} must have primary constructor with a " +
-                        "single parameter that takes initial state of ${state::class.java.simpleName}."
-                }
-                viewModelClass.kotlin.primaryConstructor?.parameters?.get(0)?.type != state::class -> {
-                    "${MvRxViewModelFactory::class.java.simpleName} must have primary constructor with a " +
-                        "single parameter that takes initial state of ${state::class.java.simpleName}. Found type " +
-                        "${viewModelClass.kotlin.primaryConstructor?.parameters?.get(0)?.type?.javaClass?.simpleName}"
-                }
-                viewModelClass.kotlin.primaryConstructor?.parameters?.get(0)?.isOptional == true -> {
-                    "initialState may not be an optional constructor parameter."
-                }
-                else -> {
-                    "${viewModelClass.simpleName} must have a companion object implementing " +
-                        "${MvRxViewModelFactory::class.java.simpleName} and a ${JvmStatic::class.java.simpleName} " +
-                        "annotated create method *or* have primary constructor with a single parameter for " +
-                        "initial state of ${state::class.java.simpleName}."
-                }
+            // If null, use Kotlin reflect for best error message. We will crash anyway, so performance
+            // doesn't matter.
+            if (viewModelClass.kotlin.primaryConstructor?.parameters?.size?.let { it > 1 } == true) {
+                "${viewModelClass.simpleName} takes dependencies other than initialState. " +
+                    "It must have companion object implementing ${MvRxViewModelFactory::class.java.simpleName} " +
+                    "with a create method returning a non-null ViewModel."
+            } else {
+                "${MvRxViewModelFactory::class.java.simpleName} must have primary constructor with a " +
+                    "single parameter that takes initial state of ${stateClass.simpleName}."
             }
         }
     }
 
-    @Suppress("UNCHECKED_CAST")
-    private fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createFactoryViewModel(
-        viewModelClass: Class<VM>,
-        fragment: Fragment,
-        state: S
-    ): VM? {
-        val method = try {
-            viewModelClass.getMethod("create", Fragment::class.java, state::class.java)
-        } catch (exception: NoSuchMethodException) {
-            // No fragment create found check for activity create.
-            return createFactoryViewModel(viewModelClass, fragment.requireActivity(), state)
-        }
-        val viewModel = method?.invoke(null, fragment, state) as? VM
-        return viewModel ?: createDefaultViewModel(viewModelClass, state)
+    /**
+     * Given a companion class, use Java reflection to create an instance. This is used over
+     * Kotlin reflection for performance.
+     */
+    private fun Class<*>.instance(): Any {
+        return declaredConstructors.first { it.parameterTypes.size == 1 }.newInstance(null)
     }
 
-    @Suppress("UNCHECKED_CAST")
-    private fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createFactoryViewModel(
-        viewModelClass: Class<VM>,
-        activity: FragmentActivity,
-        state: S
-    ): VM? {
-        val method = try {
-            viewModelClass.getMethod("create", FragmentActivity::class.java, state::class.java)
-        } catch (exception: NoSuchMethodException) {
-            null
+    private fun <VM : BaseMvRxViewModel<S>, S : MvRxState> createInitialState(
+        viewModelClass: Class<out VM>,
+        stateClass: Class<out S>,
+        viewModelContext: ViewModelContext,
+        stateRestorer: (S) -> S
+    ): S {
+        val companionClass = try { Class.forName("${viewModelClass.name}\$Companion") } catch (exception: ClassNotFoundException) { null }
+        @Suppress("UNCHECKED_CAST")
+        val factoryState = companionClass?.let {
+            try {
+                companionClass.getMethod("initialState", ViewModelContext::class.java)
+                    .invoke(companionClass.instance(), viewModelContext) as S?
+            } catch (exception: NoSuchMethodException) {
+                // Check for JvmStatic method.
+                viewModelClass.getMethod("initialState", ViewModelContext::class.java)
+                    .invoke(null, viewModelContext) as S?
+            }
         }
-        val viewModel = method?.invoke(null, activity, state) as? VM
-        return viewModel ?: createDefaultViewModel(viewModelClass, state)
+        return stateRestorer(factoryState
+            ?: createStateFromConstructor(stateClass, viewModelContext.rawArgs))
     }
 
     @Suppress("UNCHECKED_CAST")
@@ -156,4 +113,38 @@ object MvRxViewModelProvider {
         }
         return null
     }
+
+    /**
+     * For internal use only. Public for inline.
+     *
+     * Searches [stateClass] for a single argument constructor matching the type of [args]. If [args] is null, then
+     * no arg constructor is invoked.
+     *
+     */
+    @Suppress("FunctionName") // Public for inline.
+    internal fun <S : MvRxState> createStateFromConstructor(stateClass: Class<S>, args: Any?): S {
+        val argsConstructor = args?.let {
+            val argType = it::class.java
+
+            stateClass.constructors.firstOrNull { constructor ->
+                constructor.parameterTypes.size == 1 && isAssignableTo(argType, constructor.parameterTypes[0])
+            }
+        }
+
+        @Suppress("UNCHECKED_CAST")
+        return argsConstructor?.newInstance(args) as? S
+            ?: try {
+                stateClass.newInstance()
+            } catch (@Suppress("TooGenericExceptionCaught") e: Throwable) {
+                null
+            }
+            ?: throw IllegalStateException(
+                "Attempt to create the MvRx state class ${stateClass.simpleName} has failed. One of the following must be true:" +
+                    "\n 1) The state class has default values for every property." +
+                    "\n 2) The state class has a secondary constructor for ${args?.javaClass?.simpleName ?: "a fragment argument"}." +
+                    "\n 3) The ViewModel using the state must has a companion object implementing MvRxFactory with an initialState function " +
+                    "that does not return null. "
+            )
+    }
+
 }
diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelStore.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelStore.kt
index 439cd5db9..95750fb06 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelStore.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelStore.kt
@@ -118,18 +118,12 @@ class MvRxViewModelStore(private val viewModelStore: ViewModelStore) {
 
     private fun restoreViewModel(activity: FragmentActivity, holder: MvRxPersistedViewModelHolder, arguments: Any?): ViewModel {
         val (viewModelClass, stateClass, viewModelState) = holder
-        // If there is a key in the fragmentArgsForActivityViewModelState map, then this is an activity ViewModel. The map value will contain
-        // the Fragment args that the ViewModel was created with.
-        val state = _initialStateProvider(stateClass, arguments).let(viewModelState::restorePersistedState)
-        return createViewModel(viewModelClass, activity, state)
+        return createViewModel(viewModelClass, stateClass, ActivityViewModelContext(activity, arguments), viewModelState::restorePersistedState)
     }
 
     private fun restoreViewModel(fragment: Fragment, holder: MvRxPersistedViewModelHolder, arguments: Any?): ViewModel {
         val (viewModelClass, stateClass, viewModelState) = holder
-        // If there is a key in the fragmentArgsForActivityViewModelState map, then this is an activity ViewModel. The map value will contain
-        // the Fragment args that the ViewModel was created with.
-        val state = _initialStateProvider(stateClass, arguments).let(viewModelState::restorePersistedState)
-        return createViewModel(viewModelClass, fragment, state)
+        return createViewModel(viewModelClass, stateClass, FragmentViewModelContext(fragment.requireActivity(), arguments, fragment), viewModelState::restorePersistedState)
     }
 
     /**
diff --git a/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt b/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
index 6754cd828..8268b6973 100644
--- a/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
+++ b/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
@@ -1,31 +1,28 @@
 package com.airbnb.mvrx
 
+import android.os.Parcelable
 import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
+import kotlinx.android.parcel.Parcelize
 import org.junit.Assert.assertEquals
 import org.junit.Before
 import org.junit.Test
 import org.robolectric.Robolectric
 import java.lang.reflect.InvocationTargetException
 
-data class FactoryState(val count: Int = 0) : MvRxState
-class TestFactoryViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
-    companion object : MvRxViewModelFactory<FactoryState> {
-        @JvmStatic
-        override fun create(activity: FragmentActivity, state: FactoryState) = TestFactoryViewModel(state, 5)
-    }
+data class FactoryState(val greeting: String = "") : MvRxState {
+    constructor(args: TestArgs) : this("${args.greeting} constructor")
 }
 
-class TestFragmentFactoryViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
-    companion object : MvRxFragmentViewModelFactory<FactoryState> {
-        @JvmStatic
-        override fun create(fragment: Fragment, state: FactoryState) = TestFragmentFactoryViewModel(state, 5)
-    }
-}
+@Parcelize
+data class TestArgs(val greeting: String) : Parcelable
 
 class ViewModelFactoryTestFragment : Fragment()
 
-class FactoryTest : BaseTest() {
+/**
+ * Tests ViewModel creation when there is no factory.
+ */
+class NoFactoryTest : BaseTest() {
 
     private lateinit var activity: FragmentActivity
 
@@ -35,114 +32,246 @@ class FactoryTest : BaseTest() {
     }
 
     @Test
-    fun createDefaultViewModel() {
+    fun createFromActivityOwner() {
         class MyViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState)
-        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, activity) { FactoryState() }
+
+        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(), state)
+            assertEquals(FactoryState("hello constructor"), state)
         }
     }
 
     @Test
-    fun createDefaultViewModelFromFragment() {
+    fun createFromFragmentOwner() {
         val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
         class MyViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState)
-        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, fragment) { FactoryState() }
+
+        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(), state)
+            assertEquals(FactoryState("hello constructor"), state)
+        }
+    }
+
+    private data class PrivateState(val count1: Int = 0) : MvRxState
+
+    @Test(expected = InvocationTargetException::class)
+    fun failOnPrivateState() {
+        class MyViewModel(initialState: PrivateState) : TestMvRxViewModel<PrivateState>(initialState)
+        // Create a view model to run state validation checks.
+        @Suppress("UNUSED_VARIABLE")
+        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, PrivateState::class.java, ActivityViewModelContext(activity, null))
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun failOnDefaultState() {
+        class MyViewModel(initialState: FactoryState = FactoryState()) : TestMvRxViewModel<FactoryState>(initialState)
+        MvRxViewModelProvider.get(MyViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, null))
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun failOnWrongSingleParameterType() {
+        class ViewModel : BaseMvRxViewModel<FactoryState>(initialState = FactoryState(), debugMode = false)
+        MvRxViewModelProvider.get(ViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, null))
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun failOnMultipleParametersAndNoCompanion() {
+        class OptionalParamViewModel(initialState: FactoryState, debugMode: Boolean = false) : BaseMvRxViewModel<FactoryState>(initialState, debugMode)
+        MvRxViewModelProvider.get(OptionalParamViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, null))
+    }
+
+    @Test(expected = IllegalArgumentException::class)
+    fun failOnNoViewModelParameters() {
+        class OptionalParamViewModel : BaseMvRxViewModel<FactoryState>(initialState = FactoryState(), debugMode = false)
+        MvRxViewModelProvider.get(OptionalParamViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, null))
+    }
+}
+
+/**
+ * Test a factory which only uses a custom ViewModel create.
+ */
+class FactoryViewModelTest : BaseTest() {
+
+    private class TestFactoryViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = TestFactoryViewModel(state, 5)
+        }
+    }
+
+    private class TestFactoryJvmStaticViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryJvmStaticViewModel, FactoryState> {
+            @JvmStatic
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = TestFactoryJvmStaticViewModel(state, 5)
         }
     }
 
+    private class TestNullFactory(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = null
+        }
+    }
+
+    private lateinit var activity: FragmentActivity
+
+    @Before
+    fun setup() {
+        activity = Robolectric.setupActivity(FragmentActivity::class.java)
+    }
+
     @Test
-    fun createDefaultViewModelWithState() {
-        class MyViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState)
-        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, activity) { FactoryState(count = 5) }
+    fun createFromActivityOwner() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(count = 5), state)
+            assertEquals(FactoryState("hello constructor"), state)
         }
+        assertEquals(5, viewModel.otherProp)
     }
 
     @Test
-    fun createDefaultViewModelWithStateFromFragment() {
+    fun createFromFragmentOwner() {
         val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
-        class MyViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState)
-        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, fragment) { FactoryState(count = 5) }
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(count = 5), state)
+            assertEquals(FactoryState("hello constructor"), state)
         }
+        assertEquals(5, viewModel.otherProp)
     }
 
     @Test
-    fun createWithFactory() {
-        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, activity) { FactoryState(count = 5) }
+    fun createWithJvmStatic() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryJvmStaticViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(count = 5), state)
+            assertEquals(FactoryState("hello constructor"), state)
         }
         assertEquals(5, viewModel.otherProp)
     }
 
     @Test
-    fun createWithActivityFactoryFromFragment() {
-        val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
-        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, fragment) { FactoryState(count = 5) }
+    fun nullInitialStateDelgatesToConstructor() {
+        val viewModel = MvRxViewModelProvider.get(TestNullFactory::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(count = 5), state)
+            assertEquals(FactoryState("hello constructor"), state)
+        }
+    }
+}
+
+/**
+ * Test a factory which only uses a custom initialState.
+ */
+class FactoryStateTest : BaseTest() {
+
+    private class TestFactoryViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
+            override fun initialState(viewModelContext: ViewModelContext): FactoryState? = FactoryState("${viewModelContext.args<TestArgs>().greeting} factory")
+        }
+    }
+
+    private class TestFactoryJvmStaticViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryJvmStaticViewModel, FactoryState> {
+            override fun initialState(viewModelContext: ViewModelContext): FactoryState? = FactoryState("${viewModelContext.args<TestArgs>().greeting} factory")
+        }
+    }
+
+    private class TestNullFactory(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestNullFactory, FactoryState> {
+            override fun initialState(viewModelContext: ViewModelContext): FactoryState? = null
+        }
+    }
+
+    private lateinit var activity: FragmentActivity
+
+    @Before
+    fun setup() {
+        activity = Robolectric.setupActivity(FragmentActivity::class.java)
+    }
+
+    @Test
+    fun createFromActivityOwner() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello factory"), state)
         }
-        assertEquals(5, viewModel.otherProp)
     }
 
     @Test
-    fun createWithFragmentFactoryFromFragment() {
+    fun createFromFragmentOwner() {
         val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
-        val viewModel = MvRxViewModelProvider.get(TestFragmentFactoryViewModel::class.java, fragment) { FactoryState(count = 5) }
+
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
         withState(viewModel) { state ->
-            assertEquals(FactoryState(count = 5), state)
+            assertEquals(FactoryState("hello factory"), state)
         }
-        assertEquals(5, viewModel.otherProp)
     }
 
-    private data class PrivateState(val count1: Int = 0) : MvRxState
+    @Test
+    fun createWithJvmStatic() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryJvmStaticViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello factory"), state)
+        }
+    }
 
-    @Test(expected = InvocationTargetException::class)
-    fun failOnPrivateState() {
-        class MyViewModel(initialState: PrivateState) : TestMvRxViewModel<PrivateState>(initialState)
-        // Create a view model to run state validation checks.
-        @Suppress("UNUSED_VARIABLE")
-        val viewModel = MvRxViewModelProvider.get(MyViewModel::class.java, activity) { PrivateState() }
+    @Test
+    fun nullInitialStateDelgatesToConstructor() {
+        val viewModel = MvRxViewModelProvider.get(TestNullFactory::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello constructor"), state)
+        }
     }
+}
 
-    @Test(expected = IllegalArgumentException::class)
-    fun failOnDefaultState() {
-        class MyViewModel(initialState: FactoryState = FactoryState()) : TestMvRxViewModel<FactoryState>(initialState)
-        MvRxViewModelProvider.get(MyViewModel::class.java, activity) { FactoryState(count = 5) }
+/**
+ * Test a factory which uses both a custom State and ViewModel create.
+ */
+class FactoryViewModelAndStateTest : BaseTest() {
+
+    private class TestFactoryViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = TestFactoryViewModel(FactoryState("${viewModelContext.args<TestArgs>().greeting} factory"), 5)
+        }
     }
 
-    @Test(expected = IllegalArgumentException::class)
-    fun failOnWrongSingleParameterType() {
-        class ViewModel : BaseMvRxViewModel<FactoryState>(initialState = FactoryState(), debugMode = false)
-        MvRxViewModelProvider.get(ViewModel::class.java, activity) { FactoryState(count = 5) }
+    private class TestFactoryJvmStaticViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
+        companion object : MvRxViewModelFactory<TestFactoryJvmStaticViewModel, FactoryState> {
+            @JvmStatic
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = TestFactoryJvmStaticViewModel(FactoryState("${viewModelContext.args<TestArgs>().greeting} factory"), 5)
+        }
     }
 
-    @Test(expected = IllegalArgumentException::class)
-    fun failOnMultipleParametersAndNoCompanion() {
-        class OptionalParamViewModel(initialState: FactoryState, debugMode: Boolean = false) : BaseMvRxViewModel<FactoryState>(initialState, debugMode)
-        MvRxViewModelProvider.get(OptionalParamViewModel::class.java, activity) { FactoryState(count = 5) }
+    private lateinit var activity: FragmentActivity
+
+    @Before
+    fun setup() {
+        activity = Robolectric.setupActivity(FragmentActivity::class.java)
     }
 
-    @Test(expected = IllegalArgumentException::class)
-    fun failOnNoViewModelParameters() {
-        class OptionalParamViewModel : BaseMvRxViewModel<FactoryState>(initialState = FactoryState(), debugMode = false)
-        MvRxViewModelProvider.get(OptionalParamViewModel::class.java, activity) { FactoryState(count = 5) }
+    @Test
+    fun createFromActivityOwner() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello factory"), state)
+        }
+        assertEquals(5, viewModel.otherProp)
     }
 
-    class TestFactoryViewModelNoJvmStatic(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
-        companion object : MvRxViewModelFactory<FactoryState> {
-            override fun create(activity: FragmentActivity, state: FactoryState) = TestFactoryViewModelNoJvmStatic(state, 5)
+    @Test
+    fun createFromFragmentOwner() {
+        val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
+
+        val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello factory"), state)
         }
+        assertEquals(5, viewModel.otherProp)
     }
 
-    @Test(expected = IllegalArgumentException::class)
-    fun failOnNoJvmStaticInCompanion() {
-        MvRxViewModelProvider.get(TestFactoryViewModelNoJvmStatic::class.java, activity) { FactoryState(count = 5) }
+    @Test
+    fun createWithJvmStatic() {
+        val viewModel = MvRxViewModelProvider.get(TestFactoryJvmStaticViewModel::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
+        withState(viewModel) { state ->
+            assertEquals(FactoryState("hello factory"), state)
+        }
+        assertEquals(5, viewModel.otherProp)
     }
+
 }
diff --git a/mvrx/src/test/kotlin/com/airbnb/mvrx/FragmentStateProviderTest.kt b/mvrx/src/test/kotlin/com/airbnb/mvrx/InitialConstructorStateTest.kt
similarity index 74%
rename from mvrx/src/test/kotlin/com/airbnb/mvrx/FragmentStateProviderTest.kt
rename to mvrx/src/test/kotlin/com/airbnb/mvrx/InitialConstructorStateTest.kt
index 3199a1e91..b422f79ba 100644
--- a/mvrx/src/test/kotlin/com/airbnb/mvrx/FragmentStateProviderTest.kt
+++ b/mvrx/src/test/kotlin/com/airbnb/mvrx/InitialConstructorStateTest.kt
@@ -10,13 +10,13 @@ import org.junit.Test
 import java.io.Serializable
 
 /** Test auto creating state from fragment arguments. */
-class FragmentStateProviderTest : BaseTest() {
+class InitialConstructorStateTest : BaseTest() {
 
     @Test
     fun testParcelableArgs() {
         val args = ParcelableArgs("hello")
         val frag = Frag(args)
-        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+        val state = MvRxViewModelProvider.createStateFromConstructor(TestState::class.java, frag._fragmentArgsProvider())
         assertEquals(args.str, state.str)
         assertEquals(null, state.num)
     }
@@ -25,7 +25,7 @@ class FragmentStateProviderTest : BaseTest() {
     fun testSerializableArgs() {
         val args = SerializableArgs(7)
         val frag = Frag(args)
-        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+        val state = MvRxViewModelProvider.createStateFromConstructor(TestState::class.java, frag._fragmentArgsProvider())
         assertEquals(null, state.str)
         assertEquals(args.num, state.num)
     }
@@ -33,7 +33,7 @@ class FragmentStateProviderTest : BaseTest() {
     @Test
     fun testLongArgs() {
         val frag = Frag(8L)
-        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+        val state = MvRxViewModelProvider.createStateFromConstructor(TestState::class.java, frag._fragmentArgsProvider())
         assertEquals("id", state.str)
         assertEquals(8, state.num)
     }
@@ -41,7 +41,7 @@ class FragmentStateProviderTest : BaseTest() {
     @Test
     fun testNoArgs() {
         val frag = Frag(null)
-        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+        val state = MvRxViewModelProvider.createStateFromConstructor(TestState::class.java, frag._fragmentArgsProvider())
         assertEquals("empty", state.str)
         assertEquals(2, state.num)
     }
@@ -49,7 +49,7 @@ class FragmentStateProviderTest : BaseTest() {
     @Test
     fun testNoMatchingArgsFallsbackToEmptyConstructor() {
         val frag = Frag("unknown arg type")
-        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+        val state = MvRxViewModelProvider.createStateFromConstructor(TestState::class.java, frag._fragmentArgsProvider())
         assertEquals("empty", state.str)
         assertEquals(2, state.num)
     }
@@ -58,7 +58,7 @@ class FragmentStateProviderTest : BaseTest() {
     //    @Test(expected = IllegalStateException::class)
     //    fun testFailsWithNoMatchingConstructor() {
     //        val frag = Frag("string")
-    //        val state: TestState = frag._fragmentViewModelInitialStateProvider()
+    //        val state: TestState = frag._fragmentViewModelConstructorInitialState()
     //    }
 }
 
diff --git a/mvrx/src/test/kotlin/com/airbnb/mvrx/MvRxExtensionsTest.kt b/mvrx/src/test/kotlin/com/airbnb/mvrx/MvRxExtensionsTest.kt
deleted file mode 100644
index 1a24a41a4..000000000
--- a/mvrx/src/test/kotlin/com/airbnb/mvrx/MvRxExtensionsTest.kt
+++ /dev/null
@@ -1,28 +0,0 @@
-package com.airbnb.mvrx
-
-import android.os.Parcelable
-import kotlinx.android.parcel.Parcelize
-import org.junit.*
-import org.junit.Assert.*
-
-class MvRxExtensionsTest {
-
-  @Parcelize
-  data class Arg(val value: String) : Parcelable
-
-  class DemoState1(val arg: Parcelable) : MvRxState
-
-  class DemoState2(val arg: Arg) : MvRxState
-
-  @Test
-  fun testInitialDemoState1() {
-    val state = _initialStateProvider(DemoState1::class.java, Arg("Hello"))
-    assertEquals("Hello", (state.arg as Arg).value)
-  }
-
-  @Test
-  fun testInitialDemoState2() {
-    val state = _initialStateProvider(DemoState2::class.java, Arg("Hello"))
-    assertEquals("Hello", state.arg.value)
-  }
-}
diff --git a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeDetailFragment.kt b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeDetailFragment.kt
index c7542b3d6..264b8e9a4 100644
--- a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeDetailFragment.kt
+++ b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeDetailFragment.kt
@@ -41,13 +41,9 @@ class DadJokeDetailViewModel(
         dadJokeService.fetch(state.id).execute { copy(joke = it) }
     }
 
-    companion object : MvRxViewModelFactory<DadJokeDetailState> {
-        @JvmStatic
-        override fun create(
-            activity: FragmentActivity,
-            state: DadJokeDetailState
-        ): BaseMvRxViewModel<DadJokeDetailState> {
-            val service: DadJokeService by activity.inject()
+    companion object : MvRxViewModelFactory<DadJokeDetailViewModel, DadJokeDetailState> {
+        override fun create(viewModelContext: ViewModelContext, state: DadJokeDetailState): DadJokeDetailViewModel {
+            val service: DadJokeService by viewModelContext.activity.inject()
             return DadJokeDetailViewModel(state, service)
         }
     }
diff --git a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeIndexViewModel.kt b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeIndexViewModel.kt
index 6a341d126..ccfcb310e 100644
--- a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeIndexViewModel.kt
+++ b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/DadJokeIndexViewModel.kt
@@ -1,12 +1,6 @@
 package com.airbnb.mvrx.sample.features.dadjoke
 
-import android.support.v4.app.FragmentActivity
-import com.airbnb.mvrx.Async
-import com.airbnb.mvrx.BaseMvRxViewModel
-import com.airbnb.mvrx.Loading
-import com.airbnb.mvrx.MvRxState
-import com.airbnb.mvrx.MvRxViewModelFactory
-import com.airbnb.mvrx.Uninitialized
+import com.airbnb.mvrx.*
 import com.airbnb.mvrx.sample.core.MvRxViewModel
 import com.airbnb.mvrx.sample.models.Joke
 import com.airbnb.mvrx.sample.models.JokesResponse
@@ -50,9 +44,10 @@ class DadJokeIndexViewModel(
      *
      * @see MvRxViewModelFactory
      */
-    companion object : MvRxViewModelFactory<DadJokeIndexState> {
-        @JvmStatic override fun create(activity: FragmentActivity, state: DadJokeIndexState): BaseMvRxViewModel<DadJokeIndexState> {
-            val service: DadJokeService by activity.inject()
+    companion object : MvRxViewModelFactory<DadJokeIndexViewModel, DadJokeIndexState> {
+
+        override fun create(viewModelContext: ViewModelContext, state: DadJokeIndexState): DadJokeIndexViewModel {
+            val service: DadJokeService by viewModelContext.activity.inject()
             return DadJokeIndexViewModel(state, service)
         }
     }
diff --git a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/RandomDadJokeFragment.kt b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/RandomDadJokeFragment.kt
index 7533bb6d5..8a206b85e 100644
--- a/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/RandomDadJokeFragment.kt
+++ b/sample/src/main/java/com/airbnb/mvrx/sample/features/dadjoke/RandomDadJokeFragment.kt
@@ -1,6 +1,5 @@
 package com.airbnb.mvrx.sample.features.dadjoke
 
-import android.support.v4.app.FragmentActivity
 import com.airbnb.mvrx.*
 import com.airbnb.mvrx.sample.core.BaseFragment
 import com.airbnb.mvrx.sample.core.MvRxViewModel
@@ -10,6 +9,7 @@ import com.airbnb.mvrx.sample.network.DadJokeService
 import com.airbnb.mvrx.sample.views.basicRow
 import com.airbnb.mvrx.sample.views.loadingRow
 import com.airbnb.mvrx.sample.views.marquee
+import io.reactivex.schedulers.Schedulers
 import org.koin.android.ext.android.inject
 
 data class RandomDadJokeState(val joke: Async<Joke> = Uninitialized) : MvRxState
@@ -23,16 +23,13 @@ class RandomDadJokeViewModel(
     }
 
     fun fetchRandomJoke() {
-        dadJokeService.random().execute { copy(joke = it) }
+        dadJokeService.random().subscribeOn(Schedulers.io()).execute { copy(joke = it) }
     }
 
-    companion object : MvRxViewModelFactory<RandomDadJokeState> {
-        @JvmStatic
-        override fun create(
-            activity: FragmentActivity,
-            state: RandomDadJokeState
-        ): BaseMvRxViewModel<RandomDadJokeState> {
-            val service: DadJokeService by activity.inject()
+    companion object : MvRxViewModelFactory<RandomDadJokeViewModel, RandomDadJokeState> {
+
+        override fun create(viewModelContext: ViewModelContext, state: RandomDadJokeState): RandomDadJokeViewModel {
+            val service: DadJokeService by viewModelContext.activity.inject()
             return RandomDadJokeViewModel(state, service)
         }
     }
diff --git a/todomvrx/src/main/java/com/airbnb/mvrx/todomvrx/TasksViewModel.kt b/todomvrx/src/main/java/com/airbnb/mvrx/todomvrx/TasksViewModel.kt
index 16d75ac6b..27399a1c0 100644
--- a/todomvrx/src/main/java/com/airbnb/mvrx/todomvrx/TasksViewModel.kt
+++ b/todomvrx/src/main/java/com/airbnb/mvrx/todomvrx/TasksViewModel.kt
@@ -1,11 +1,7 @@
 package com.airbnb.mvrx.todomvrx
 
 import android.support.v4.app.FragmentActivity
-import com.airbnb.mvrx.Async
-import com.airbnb.mvrx.BaseMvRxViewModel
-import com.airbnb.mvrx.MvRxState
-import com.airbnb.mvrx.MvRxViewModelFactory
-import com.airbnb.mvrx.Uninitialized
+import com.airbnb.mvrx.*
 import com.airbnb.mvrx.todomvrx.core.MvRxViewModel
 import com.airbnb.mvrx.todomvrx.data.Task
 import com.airbnb.mvrx.todomvrx.data.Tasks
@@ -64,15 +60,17 @@ class TasksViewModel(initialState: TasksState, private val sources: List<TasksDa
         sources.forEach { it.deleteTask(id) }
     }
 
-    companion object : MvRxViewModelFactory<TasksState> {
-        @JvmStatic override fun create(activity: FragmentActivity, state: TasksState): BaseMvRxViewModel<TasksState> {
-            val database = ToDoDatabase.getInstance(activity)
+    companion object : MvRxViewModelFactory<TasksViewModel, TasksState> {
+
+        override fun create(viewModelContext: ViewModelContext, state: TasksState): TasksViewModel {
+            val database = ToDoDatabase.getInstance(viewModelContext.activity)
             // Simulate data sources of different speeds.
             // The slower one can be thought of as the network data source.
             val dataSource1 = DatabaseDataSource(database.taskDao(), 2000)
             val dataSource2 = DatabaseDataSource(database.taskDao(), 3500)
             return TasksViewModel(state, listOf(dataSource1, dataSource2))
         }
+
     }
 }
 

From 5e8191582d553ea0327e158cb8da0623b931dfba Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Wed, 9 Jan 2019 17:07:00 -0800
Subject: [PATCH 2/6] Better comment

---
 mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
index 1eeb21d42..0d77fceaa 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
@@ -4,9 +4,8 @@ import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
 
 /**
- * Implement this on your ViewModel's companion object for hooks into state creation and ViewModel
- * creation. See [MvRxViewModelArgFactory] for a version which automatically types your argument
- * class.
+ * Implement this on your ViewModel's companion object for hooks into state creation and ViewModel creation. For example, if you need access
+ * to the fragment or activity owner for dependency injection.
  */
 interface MvRxViewModelFactory<VM: BaseMvRxViewModel<S>, S : MvRxState> {
 

From ec5e67460ba6c505683b14b4d5a25a5b711d370f Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Wed, 9 Jan 2019 21:49:03 -0800
Subject: [PATCH 3/6] comments, restrict to library

---
 mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt       | 4 ++++
 .../src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt | 5 ++---
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
index 47150e9ed..c83345f62 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
@@ -1,6 +1,8 @@
 package com.airbnb.mvrx
 
 import android.arch.lifecycle.ViewModelProviders
+import android.support.annotation.RestrictTo
+import android.support.annotation.RestrictTo.Scope.LIBRARY
 import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
 import kotlin.properties.ReadOnlyProperty
@@ -58,6 +60,7 @@ inline fun <T, reified VM : BaseMvRxViewModel<S>, reified S : MvRxState> T.activ
  * Looks for [MvRx.KEY_ARG] on the arguments of the fragments.
  */
 @Suppress("FunctionName")
+@RestrictTo(LIBRARY)
 fun <T : Fragment> T._fragmentArgsProvider(): Any? = arguments?.get(MvRx.KEY_ARG)
 
 /**
@@ -69,6 +72,7 @@ fun <T : Fragment> T._fragmentArgsProvider(): Any? = arguments?.get(MvRx.KEY_ARG
  * in a new process.
  */
 @Suppress("FunctionName")
+@RestrictTo(LIBRARY)
 inline fun <T : Fragment> T._activityArgsProvider(keyFactory: () -> String): Any? {
     val args: Any? = arguments?.get(MvRx.KEY_ARG)
     val activity = requireActivity()
diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
index 6bf761975..bf491a662 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
@@ -115,7 +115,6 @@ object MvRxViewModelProvider {
     }
 
     /**
-     * For internal use only. Public for inline.
      *
      * Searches [stateClass] for a single argument constructor matching the type of [args]. If [args] is null, then
      * no arg constructor is invoked.
@@ -140,9 +139,9 @@ object MvRxViewModelProvider {
             }
             ?: throw IllegalStateException(
                 "Attempt to create the MvRx state class ${stateClass.simpleName} has failed. One of the following must be true:" +
-                    "\n 1) The state class has default values for every property." +
+                    "\n 1) The state class has default values for every constructor property." +
                     "\n 2) The state class has a secondary constructor for ${args?.javaClass?.simpleName ?: "a fragment argument"}." +
-                    "\n 3) The ViewModel using the state must has a companion object implementing MvRxFactory with an initialState function " +
+                    "\n 3) The ViewModel using the state must have a companion object implementing MvRxFactory with an initialState function " +
                     "that does not return null. "
             )
     }

From 92f7de430665a4eebc2a384f55477b6ff31146e9 Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Wed, 9 Jan 2019 21:50:44 -0800
Subject: [PATCH 4/6] Chain arg providers

---
 mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
index c83345f62..d648672fb 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxExtensions.kt
@@ -74,7 +74,7 @@ fun <T : Fragment> T._fragmentArgsProvider(): Any? = arguments?.get(MvRx.KEY_ARG
 @Suppress("FunctionName")
 @RestrictTo(LIBRARY)
 inline fun <T : Fragment> T._activityArgsProvider(keyFactory: () -> String): Any? {
-    val args: Any? = arguments?.get(MvRx.KEY_ARG)
+    val args: Any? = _fragmentArgsProvider()
     val activity = requireActivity()
     if (activity is MvRxViewModelStoreOwner) {
         activity.mvrxViewModelStore._saveActivityViewModelArgs(keyFactory(), args)

From edbb23ef2c7a3aafd383725dc5bcc54e2aeffd7d Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Wed, 9 Jan 2019 22:24:44 -0800
Subject: [PATCH 5/6] Better explanation for why activity view model can only
 get activity reference, more robust tests

---
 .../com/airbnb/mvrx/MvRxViewModelFactory.kt   | 34 +++++++------------
 .../kotlin/com/airbnb/mvrx/FactoryTest.kt     | 29 ++++++++++++----
 2 files changed, 35 insertions(+), 28 deletions(-)

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
index 0d77fceaa..8d625a816 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
@@ -7,7 +7,7 @@ import android.support.v4.app.FragmentActivity
  * Implement this on your ViewModel's companion object for hooks into state creation and ViewModel creation. For example, if you need access
  * to the fragment or activity owner for dependency injection.
  */
-interface MvRxViewModelFactory<VM: BaseMvRxViewModel<S>, S : MvRxState> {
+interface MvRxViewModelFactory<VM : BaseMvRxViewModel<S>, S : MvRxState> {
 
     /**
      * @param viewModelContext [ViewModelContext] which contains the ViewModel owner and arguments.
@@ -54,34 +54,26 @@ sealed class ViewModelContext {
     /**
      * Convenience method to type [rawArgs].
      */
-    abstract fun <A> args(): A
+    @Suppress("UNCHECKED_CAST")
+    fun <A> args(): A = rawArgs as A
 }
 
 /**
- * The [ViewModelContext] for a ViewModel created with an
- * activity scope (`val viewModel by activityViewModel<MyViewModel>`).
+ * The [ViewModelContext] for a ViewModel created with an activity scope (`val viewModel by activityViewModel<MyViewModel>`). Although a fragment
+ * reference is available when an activity scoped ViewModel is first created, during process restoration, activity scoped ViewModels will be created
+ * _without_ a fragment reference, so it is only safe to reference the activity.
  */
 class ActivityViewModelContext(
-        override val activity: FragmentActivity,
-        override val rawArgs: Any?
-) : ViewModelContext() {
-    override fun <A> args(): A {
-        @Suppress("UNCHECKED_CAST")
-        return rawArgs as A
-    }
-}
+    override val activity: FragmentActivity,
+    override val rawArgs: Any?
+) : ViewModelContext()
 
 /**
  * The [ViewModelContext] for a ViewModel created with a
  * fragment scope (`val viewModel by fragmentViewModel<MyViewModel>`).
  */
 class FragmentViewModelContext(
-        override val activity: FragmentActivity,
-        override val rawArgs: Any?,
-        val fragment: Fragment
-) : ViewModelContext() {
-    override fun <A> args(): A {
-        @Suppress("UNCHECKED_CAST")
-        return rawArgs as A
-    }
-}
+    override val activity: FragmentActivity,
+    override val rawArgs: Any?,
+    val fragment: Fragment
+) : ViewModelContext()
diff --git a/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt b/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
index 8268b6973..276e49d41 100644
--- a/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
+++ b/mvrx/src/test/kotlin/com/airbnb/mvrx/FactoryTest.kt
@@ -1,5 +1,6 @@
 package com.airbnb.mvrx
 
+import android.os.Bundle
 import android.os.Parcelable
 import android.support.v4.app.Fragment
 import android.support.v4.app.FragmentActivity
@@ -94,7 +95,14 @@ class FactoryViewModelTest : BaseTest() {
 
     private class TestFactoryViewModel(initialState: FactoryState, val otherProp: Long) : TestMvRxViewModel<FactoryState>(initialState) {
         companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
-            override fun create(viewModelContext: ViewModelContext, state: FactoryState) = TestFactoryViewModel(state, 5)
+            override fun create(viewModelContext: ViewModelContext, state: FactoryState): TestFactoryViewModel {
+                return when (viewModelContext) {
+                    // Use Fragment args to test that there is a valid fragment reference.
+                    is FragmentViewModelContext -> TestFactoryViewModel(state, viewModelContext.fragment.arguments?.getLong("otherProp")!!)
+                    else -> TestFactoryViewModel(state, 5L)
+                }
+            }
+
         }
     }
 
@@ -130,11 +138,12 @@ class FactoryViewModelTest : BaseTest() {
     @Test
     fun createFromFragmentOwner() {
         val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
+        fragment.arguments = Bundle().apply { putLong("otherProp", 6L) }
         val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
         withState(viewModel) { state ->
             assertEquals(FactoryState("hello constructor"), state)
         }
-        assertEquals(5, viewModel.otherProp)
+        assertEquals(6, viewModel.otherProp)
     }
 
     @Test
@@ -147,7 +156,7 @@ class FactoryViewModelTest : BaseTest() {
     }
 
     @Test
-    fun nullInitialStateDelgatesToConstructor() {
+    fun nullInitialStateDelegatesToConstructor() {
         val viewModel = MvRxViewModelProvider.get(TestNullFactory::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
             assertEquals(FactoryState("hello constructor"), state)
@@ -162,7 +171,13 @@ class FactoryStateTest : BaseTest() {
 
     private class TestFactoryViewModel(initialState: FactoryState) : TestMvRxViewModel<FactoryState>(initialState) {
         companion object : MvRxViewModelFactory<TestFactoryViewModel, FactoryState> {
-            override fun initialState(viewModelContext: ViewModelContext): FactoryState? = FactoryState("${viewModelContext.args<TestArgs>().greeting} factory")
+            override fun initialState(viewModelContext: ViewModelContext): FactoryState? {
+                return when (viewModelContext) {
+                    is FragmentViewModelContext -> FactoryState("${viewModelContext.fragment.arguments?.getString("greeting")!!} and ${viewModelContext.args<TestArgs>().greeting} factory")
+                    else -> FactoryState("${viewModelContext.args<TestArgs>().greeting} factory")
+                }
+
+            }
         }
     }
 
@@ -196,10 +211,10 @@ class FactoryStateTest : BaseTest() {
     @Test
     fun createFromFragmentOwner() {
         val (_, fragment) = createFragment<ViewModelFactoryTestFragment, TestActivity>()
-
+        fragment.arguments = Bundle().apply { putString("greeting", "howdy") }
         val viewModel = MvRxViewModelProvider.get(TestFactoryViewModel::class.java, FactoryState::class.java, FragmentViewModelContext(activity, TestArgs("hello"), fragment))
         withState(viewModel) { state ->
-            assertEquals(FactoryState("hello factory"), state)
+            assertEquals(FactoryState("howdy and hello factory"), state)
         }
     }
 
@@ -212,7 +227,7 @@ class FactoryStateTest : BaseTest() {
     }
 
     @Test
-    fun nullInitialStateDelgatesToConstructor() {
+    fun nullInitialStateDelegatesToConstructor() {
         val viewModel = MvRxViewModelProvider.get(TestNullFactory::class.java, FactoryState::class.java, ActivityViewModelContext(activity, TestArgs("hello")))
         withState(viewModel) { state ->
             assertEquals(FactoryState("hello constructor"), state)

From 3df423c1dbd87e4edbb52dad0a86db9ba3cb5811 Mon Sep 17 00:00:00 2001
From: Ben Schwab <ben.schwab@airbnb.com>
Date: Thu, 10 Jan 2019 13:31:16 -0800
Subject: [PATCH 6/6] Add convience methods to type fragment and activity

---
 .../com/airbnb/mvrx/MvRxViewModelFactory.kt   | 28 +++++++++++++++----
 .../com/airbnb/mvrx/MvRxViewModelProvider.kt  |  2 +-
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
index 8d625a816..65a93b16d 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelFactory.kt
@@ -46,16 +46,23 @@ sealed class ViewModelContext {
      * The activity which is using the ViewModel.
      */
     abstract val activity: FragmentActivity
+
+    /**
+     * Convenience method to type [activity].
+     */
+    @Suppress("UNCHECKED_CAST")
+    fun <A : FragmentActivity> activity() : A = activity as A
+
     /**
      * Fragment arguments set via [MvRx.KEY_ARG].
      */
-    abstract val rawArgs: Any?
+    abstract val args: Any?
 
     /**
-     * Convenience method to type [rawArgs].
+     * Convenience method to type [args].
      */
     @Suppress("UNCHECKED_CAST")
-    fun <A> args(): A = rawArgs as A
+    fun <A> args(): A = args as A
 }
 
 /**
@@ -65,7 +72,7 @@ sealed class ViewModelContext {
  */
 class ActivityViewModelContext(
     override val activity: FragmentActivity,
-    override val rawArgs: Any?
+    override val args: Any?
 ) : ViewModelContext()
 
 /**
@@ -74,6 +81,15 @@ class ActivityViewModelContext(
  */
 class FragmentViewModelContext(
     override val activity: FragmentActivity,
-    override val rawArgs: Any?,
+    override val args: Any?,
+    /**
+     * The fragment owner of the ViewModel.
+     */
     val fragment: Fragment
-) : ViewModelContext()
+) : ViewModelContext() {
+    /**
+     * Convenience method to type [fragment].
+     */
+    @Suppress("UNCHECKED_CAST")
+    fun <F : Fragment> fragment() : F = fragment as F
+}
diff --git a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
index bf491a662..35f32d4e4 100644
--- a/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
+++ b/mvrx/src/main/kotlin/com/airbnb/mvrx/MvRxViewModelProvider.kt
@@ -98,7 +98,7 @@ object MvRxViewModelProvider {
             }
         }
         return stateRestorer(factoryState
-            ?: createStateFromConstructor(stateClass, viewModelContext.rawArgs))
+            ?: createStateFromConstructor(stateClass, viewModelContext.args))
     }
 
     @Suppress("UNCHECKED_CAST")