Backport: Leak on backstack recreation

resaca/api/resaca.api

@@ -79,8 +79,7 @@ public final class com/sebaslogen/resaca/ScopedViewModelContainerKt {
 
 public final class com/sebaslogen/resaca/ScopedViewModelOwner {
 	public static final field $stable I
-	public fun <init> (Ljava/lang/String;Lkotlin/reflect/KClass;Landroidx/lifecycle/ViewModelProvider$Factory;Landroidx/lifecycle/viewmodel/CreationExtras;Landroidx/lifecycle/ViewModelStoreOwner;)V
-	public final fun updateViewModelProvider (Landroidx/lifecycle/ViewModelStoreOwner;)V
+	public fun <init> (Ljava/lang/String;Lkotlin/reflect/KClass;)V
 }
 
 public final class com/sebaslogen/resaca/ViewModelCreationExtrasKt {

resaca/src/main/java/com/sebaslogen/resaca/ScopedViewModelOwner.kt

@@ -21,31 +21,28 @@ import kotlin.reflect.KClass
  *
  * @param key Unique [key] required to support [SavedStateHandle] across multiple instances of the same [ViewModel] type.
  * @param modelClass Class type of the [ViewModel] to instantiate
- * @param factory [ViewModelProvider] factory to create the requested [ViewModel] when required
- * @param creationExtras [CreationExtras] with default arguments that will be provided to the [ViewModel] through the [SavedStateHandle] and creationCallbacks.
- * @param viewModelStoreOwner Used to extract possible defaultViewModelCreationExtras and defaultViewModelProviderFactory
  */
 public class ScopedViewModelOwner<T : ViewModel>(
     private val key: String,
-    private val modelClass: KClass<T>,
-    private val factory: ViewModelProvider.Factory?,
-    creationExtras: CreationExtras,
-    viewModelStoreOwner: ViewModelStoreOwner
+    private val modelClass: KClass<T>
 ) {
 
     private val viewModelStore = ViewModelStore()
-    private val scopedViewModelProvider = ScopedViewModelProvider(factory, viewModelStore, creationExtras, viewModelStoreOwner)
+    private val scopedViewModelProvider = ScopedViewModelProvider(viewModelStore)
 
-    internal val viewModel: T
+    internal fun getViewModel(factory: ViewModelProvider.Factory?, viewModelStoreOwner: ViewModelStoreOwner, creationExtras: CreationExtras): T {
+        val viewModelProvider = scopedViewModelProvider.getViewModelProvider(factory, viewModelStoreOwner, creationExtras)
         @Suppress("ReplaceGetOrSet")
-        get() {
-            val canonicalName = modelClass.qualifiedName ?: throw IllegalArgumentException("Local and anonymous classes can not be ViewModels")
-            return scopedViewModelProvider.viewModelProvider.get("$canonicalName:$key", modelClass)
-        }
+        return viewModelProvider.get(getCanonicalNameKey(), modelClass)
+    }
+
+    internal fun getCachedViewModel(): T? {
+        return scopedViewModelProvider.getCachedViewModelProvider()?.get(getCanonicalNameKey(), modelClass)
+    }
 
-    @PublishedApi
-    internal fun updateViewModelProvider(viewModelStoreOwner: ViewModelStoreOwner) {
-        scopedViewModelProvider.updateViewModelProvider(viewModelStoreOwner)
+    private fun getCanonicalNameKey(): String {
+        val canonicalName = modelClass.qualifiedName ?: throw IllegalArgumentException("Local and anonymous classes can not be ViewModels")
+        return "$canonicalName:$key"
     }
 
     internal fun clear() {
@@ -62,5 +59,4 @@ public class ScopedViewModelOwner<T : ViewModel>(
                 override fun <VM : ViewModel> create(modelClass: KClass<VM>, extras: CreationExtras): VM = builder() as VM
             }
     }
-}
-
+}

resaca/src/main/java/com/sebaslogen/resaca/ScopedViewModelProvider.kt

@@ -7,67 +7,85 @@ import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.ViewModelStore
 import androidx.lifecycle.ViewModelStoreOwner
 import androidx.lifecycle.viewmodel.CreationExtras
+import java.lang.ref.WeakReference
 
 /**
- * This class provides a [ViewModelProvider] though its public [viewModelProvider] field.
+ * This class provides a [ViewModelProvider] though the [getViewModelProvider] function.
  *
- * It creates the [ViewModelProvider] using the given [factory] and [viewModelStore] plus the
+ * It creates the [ViewModelProvider] using the given [ViewModelProvider.Factory] and [viewModelStore] plus the
  * [CreationExtras] and default [ViewModelProvider.Factory] from the [ViewModelStoreOwner].
  *
- * The created [ViewModelProvider] is cached until the [ViewModelStoreOwner] is updated and
- * contains new [CreationExtras] or [ViewModelProvider.Factory].
+ * The created [ViewModelProvider] is cached in a [WeakReference] to avoid memory leaks.
  *
- * @param factory [ViewModelProvider] factory to create the requested [ViewModel] when required
  * @param viewModelStore Used to store and clear the [ViewModel]
- * @param creationExtras [CreationExtras] with default arguments that will be provided to the [ViewModel] through the [SavedStateHandle] and creationCallbacks.
- * @param viewModelStoreOwner Used to extract possible defaultViewModelCreationExtras and defaultViewModelProviderFactory
  */
 internal class ScopedViewModelProvider(
-    private val factory: ViewModelProvider.Factory?,
     private val viewModelStore: ViewModelStore,
-    private val creationExtras: CreationExtras,
-    viewModelStoreOwner: ViewModelStoreOwner
 ) {
-    private var viewModelStoreOwnerDefaultViewModelProviderFactory: ViewModelProvider.Factory? = null
-    lateinit var viewModelProvider: ViewModelProvider
-        private set
-
-    init {
-        updateViewModelProvider(viewModelStoreOwner)
-    }
+    /**
+     * Caches the created [ViewModelProvider] in the first request for a [ViewModel], in case the same [ViewModel] is requested again.
+     * It is a [WeakReference] to avoid memory leaks because the [ViewModelProvider] has a reference to the [CreationExtras],
+     * which inside has references to Activity in Android.
+     */
+    private var cachedViewModelProvider: WeakReference<ViewModelProvider>? = null
 
-    @PublishedApi
-    internal fun updateViewModelProvider(viewModelStoreOwner: ViewModelStoreOwner) {
-        val updated = updateViewModelProviderDependencies(viewModelStoreOwner)
-        if (updated) updateViewModelProvider()
-    }
+    /**
+     * Returns a [ViewModelProvider] using the [viewModelStoreOwner] and [CreationExtras].
+     *
+     * @param factory [ViewModelProvider] factory to create the requested [ViewModel]
+     * @param viewModelStoreOwner Used to extract possible [ViewModelProvider.Factory] defaultViewModelProviderFactory
+     * @param creationExtras [CreationExtras] with default arguments that will be provided to the [ViewModel] through the [SavedStateHandle] and creationCallbacks.
+     *
+     * @return [ViewModelProvider] created with the provided [factory] and [viewModelStore]
+     */
+    internal fun getViewModelProvider(
+        factory: ViewModelProvider.Factory?,
+        viewModelStoreOwner: ViewModelStoreOwner,
+        creationExtras: CreationExtras
+    ): ViewModelProvider =
+        createViewModelProvider(
+            factory = factory,
+            defaultFactory = getDefaultFactory(viewModelStoreOwner),
+            creationExtras = creationExtras
+        )
 
-    private fun updateViewModelProviderDependencies(viewModelStoreOwner: ViewModelStoreOwner): Boolean {
-        val newViewModelStoreOwnerDefaultViewModelProviderFactory =
-            (viewModelStoreOwner as? HasDefaultViewModelProviderFactory)?.defaultViewModelProviderFactory
+    /**
+     * Returns the cached [ViewModelProvider] or null if it was not created yet.
+     * Useful to get a reference to the [ViewModelProvider] to get a [ViewModel] from it if the [ViewModel] was already created.
+     */
+    internal fun getCachedViewModelProvider(): ViewModelProvider? = cachedViewModelProvider?.get()
 
-        if (newViewModelStoreOwnerDefaultViewModelProviderFactory != viewModelStoreOwnerDefaultViewModelProviderFactory) {
-            viewModelStoreOwnerDefaultViewModelProviderFactory = newViewModelStoreOwnerDefaultViewModelProviderFactory
-            return true
-        }
-        return false
-    }
+    private fun getDefaultFactory(viewModelStoreOwner: ViewModelStoreOwner): ViewModelProvider.Factory? =
+        (viewModelStoreOwner as? HasDefaultViewModelProviderFactory)?.defaultViewModelProviderFactory
 
     /**
      * Create a [ViewModelProvider] by either:
      * - using the existing [factory], or
-     * - using the default factory provided by the [ViewModelStoreOwner] in [updateViewModelProviderDependencies], or
+     * - using the default factory provided by the [ViewModelStoreOwner], or
      * - creating a default factory (e.g. for [ViewModel]s with no parameters in the constructor) using the [viewModelStore].
+     *
+     * This function also caches the created [ViewModelProvider] in [cachedViewModelProvider].
+     *
+     * @param factory [ViewModelProvider] factory to create the requested [ViewModel]
+     * @param defaultFactory Default [ViewModelProvider.Factory] to create the requested [ViewModel] from the [ViewModelStoreOwner]
+     * @param creationExtras [CreationExtras] with default arguments that will be provided to the [ViewModel] through the [SavedStateHandle] and creationCallbacks.
+     *
+     * @return [ViewModelProvider] created with the provided [factory] and [viewModelStore]
      */
-    private fun updateViewModelProvider() {
-        val defaultFactory = viewModelStoreOwnerDefaultViewModelProviderFactory
-        viewModelProvider = when {
+    private fun createViewModelProvider(
+        factory: ViewModelProvider.Factory?,
+        defaultFactory: ViewModelProvider.Factory?,
+        creationExtras: CreationExtras
+    ): ViewModelProvider {
+        val viewModelProvider = when {
             factory != null -> ViewModelProvider.create(viewModelStore, factory, creationExtras)
             defaultFactory != null -> ViewModelProvider.create(viewModelStore, defaultFactory, creationExtras)
             else -> ViewModelProvider.create(owner = object : ViewModelStoreOwner {
                 override val viewModelStore: ViewModelStore
                     get() = this@ScopedViewModelProvider.viewModelStore
             })
         }
+        cachedViewModelProvider = WeakReference(viewModelProvider)
+        return viewModelProvider
     }
 }

resaca/src/main/java/com/sebaslogen/resaca/ScopedViewModelUtils.kt

@@ -43,31 +43,27 @@ internal object ScopedViewModelUtils {
     ): T {
         cancelDisposal(positionalMemoizationKey)
 
-        val originalScopedViewModelOwner: ScopedViewModelOwner<T>? =
-            restoreAndUpdateScopedViewModelOwner(positionalMemoizationKey, scopedObjectsContainer, viewModelStoreOwner)
+        val originalScopedViewModelOwner: ScopedViewModelOwner<T>? = restoreAndUpdateScopedViewModelOwner(positionalMemoizationKey, scopedObjectsContainer)
 
         val viewModel: T =
-            if (scopedObjectKeys.containsKey(positionalMemoizationKey)
+            if (originalScopedViewModelOwner != null
+                && scopedObjectKeys.containsKey(positionalMemoizationKey)
                 && (scopedObjectKeys[positionalMemoizationKey] == externalKey)
-                && originalScopedViewModelOwner is ScopedViewModelOwner
             ) {
                 // When the object is already present and the external key matches, then return the existing one in the ScopedViewModelOwner
-                originalScopedViewModelOwner.viewModel
+                originalScopedViewModelOwner.getViewModel(factory, viewModelStoreOwner, creationExtras)
             } else { // First time ViewModel's object creation or externalKey changed
                 scopedObjectsContainer.remove(positionalMemoizationKey) // Remove in case key changed
                     ?.also { // Old object may need to be cleared before it's forgotten
-                        clearLastDisposedObject(disposedObject = it, objectsContainer = scopedObjectsContainer.values.toList())
+                        clearLastDisposedObject(it, scopedObjectsContainer.values.toList())
                     }
                 scopedObjectKeys[positionalMemoizationKey] = externalKey // Set the new external key used to track and store the new object version
                 val newScopedViewModelOwner = ScopedViewModelOwner(
                     key = positionalMemoizationKey + externalKey, // Both keys needed to handle recreation by ViewModelProvider when any of these keys changes
-                    modelClass = modelClass,
-                    factory = factory,
-                    creationExtras = creationExtras,
-                    viewModelStoreOwner = viewModelStoreOwner
+                    modelClass = modelClass
                 )
                 scopedObjectsContainer[positionalMemoizationKey] = newScopedViewModelOwner
-                newScopedViewModelOwner.viewModel
+                newScopedViewModelOwner.getViewModel(factory, viewModelStoreOwner, creationExtras)
             }
 
         return viewModel
@@ -81,11 +77,9 @@ internal object ScopedViewModelUtils {
     @PublishedApi
     internal fun <T : ViewModel> restoreAndUpdateScopedViewModelOwner(
         positionalMemoizationKey: InternalKey,
-        scopedObjectsContainer: MutableMap<InternalKey, Any>,
-        viewModelStoreOwner: ViewModelStoreOwner
+        scopedObjectsContainer: Map<InternalKey, Any>
     ): ScopedViewModelOwner<T>? =
         (scopedObjectsContainer[positionalMemoizationKey] as? ScopedViewModelOwner<T>)
-            ?.also { it.updateViewModelProvider(viewModelStoreOwner) }
 
     /**
      * An object that is being disposed should also be cleared only if there are no more references to it in this [objectsContainer]
@@ -121,7 +115,9 @@ internal object ScopedViewModelUtils {
             objectsContainer
                 .filterIsInstance<ScopedViewModelOwner<T>>()
                 .none { storedObject ->
-                    storedObject.viewModel == scopedViewModelOwner.viewModel
+                    // Cached ViewModel will be null only when one of the ViewModels requested was never actually created before this function call
+                    val viewModel = storedObject.getCachedViewModel()
+                    viewModel != null && viewModel == scopedViewModelOwner.getCachedViewModel()
                 }
         if (viewModelMissingInContainer) scopedViewModelOwner.clear()
     }