Fixes memory leak in ViewStateCache.

workflow-ui/container-android/api/container-android.api

@@ -33,39 +33,46 @@ public final class com/squareup/workflow1/ui/backstack/BackStackContainer$Compan
 	public fun getType ()Lkotlin/reflect/KClass;
 }
 
-public final class com/squareup/workflow1/ui/backstack/ViewStateCache : android/os/Parcelable {
-	public static final field CREATOR Lcom/squareup/workflow1/ui/backstack/ViewStateCache$CREATOR;
+public final class com/squareup/workflow1/ui/backstack/BackStackContainer$SavedState : android/view/View$BaseSavedState {
+	public static final field CREATOR Lcom/squareup/workflow1/ui/backstack/BackStackContainer$SavedState$CREATOR;
+	public fun <init> (Landroid/os/Parcel;)V
+	public fun <init> (Landroid/os/Parcelable;Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;)V
+	public final fun getSavedViewState ()Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
+	public fun writeToParcel (Landroid/os/Parcel;I)V
+}
+
+public final class com/squareup/workflow1/ui/backstack/BackStackContainer$SavedState$CREATOR : android/os/Parcelable$Creator {
+	public fun createFromParcel (Landroid/os/Parcel;)Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
+	public synthetic fun createFromParcel (Landroid/os/Parcel;)Ljava/lang/Object;
+	public fun newArray (I)[Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
+	public synthetic fun newArray (I)[Ljava/lang/Object;
+}
+
+public final class com/squareup/workflow1/ui/backstack/ViewStateCache {
 	public fun <init> ()V
 	public fun <init> (Ljava/util/Map;)V
 	public final fun attachToParentRegistryOwner (Ljava/lang/String;Landroidx/savedstate/SavedStateRegistryOwner;)V
-	public fun describeContents ()I
 	public final fun detachFromParentRegistry ()V
 	public final fun getViewStates$wf1_container_android ()Ljava/util/Map;
 	public final fun prune (Ljava/util/Collection;)V
-	public final fun restore (Lcom/squareup/workflow1/ui/backstack/ViewStateCache;)V
+	public final fun restore (Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;)V
+	public final fun save ()Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
 	public final fun update (Ljava/util/Collection;Landroid/view/View;Landroid/view/View;)V
-	public fun writeToParcel (Landroid/os/Parcel;I)V
-}
-
-public final class com/squareup/workflow1/ui/backstack/ViewStateCache$CREATOR : android/os/Parcelable$Creator {
-	public fun createFromParcel (Landroid/os/Parcel;)Lcom/squareup/workflow1/ui/backstack/ViewStateCache;
-	public synthetic fun createFromParcel (Landroid/os/Parcel;)Ljava/lang/Object;
-	public fun newArray (I)[Lcom/squareup/workflow1/ui/backstack/ViewStateCache;
-	public synthetic fun newArray (I)[Ljava/lang/Object;
 }
 
-public final class com/squareup/workflow1/ui/backstack/ViewStateCache$SavedState : android/view/View$BaseSavedState {
-	public static final field CREATOR Lcom/squareup/workflow1/ui/backstack/ViewStateCache$SavedState$CREATOR;
+public final class com/squareup/workflow1/ui/backstack/ViewStateCache$Saved : android/os/Parcelable {
+	public static final field CREATOR Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved$CREATOR;
 	public fun <init> (Landroid/os/Parcel;)V
-	public fun <init> (Landroid/os/Parcelable;Lcom/squareup/workflow1/ui/backstack/ViewStateCache;)V
-	public final fun getViewStateCache ()Lcom/squareup/workflow1/ui/backstack/ViewStateCache;
+	public fun <init> (Lcom/squareup/workflow1/ui/backstack/ViewStateCache;)V
+	public fun describeContents ()I
+	public final fun getViewStates$wf1_container_android ()Ljava/util/Map;
 	public fun writeToParcel (Landroid/os/Parcel;I)V
 }
 
-public final class com/squareup/workflow1/ui/backstack/ViewStateCache$SavedState$CREATOR : android/os/Parcelable$Creator {
-	public fun createFromParcel (Landroid/os/Parcel;)Lcom/squareup/workflow1/ui/backstack/ViewStateCache$SavedState;
+public final class com/squareup/workflow1/ui/backstack/ViewStateCache$Saved$CREATOR : android/os/Parcelable$Creator {
+	public fun createFromParcel (Landroid/os/Parcel;)Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
 	public synthetic fun createFromParcel (Landroid/os/Parcel;)Ljava/lang/Object;
-	public fun newArray (I)[Lcom/squareup/workflow1/ui/backstack/ViewStateCache$SavedState;
+	public fun newArray (I)[Lcom/squareup/workflow1/ui/backstack/ViewStateCache$Saved;
 	public synthetic fun newArray (I)[Ljava/lang/Object;
 }
 

workflow-ui/container-android/src/androidTest/java/com/squareup/workflow1/ui/backstack/test/ViewStateCacheTest.kt

@@ -43,13 +43,14 @@ internal class ViewStateCacheTest {
     )
     val parcel = Parcel.obtain()
 
-    parcel.writeParcelable(cache, 0)
+    parcel.writeParcelable(cache.save(), 0)
 
     parcel.setDataPosition(0)
-    val restoredCache =
-      parcel.readParcelable<ViewStateCache>(ViewStateCache::class.java.classLoader)!!.also {
-        ViewStateCache().restore(it)
-      }
+    val restoredCache = parcel.readParcelable<ViewStateCache.Saved>(
+      ViewStateCache.Saved::class.java.classLoader
+    )!!.let { restoredState ->
+      ViewStateCache().apply { restore(restoredState) }
+    }
 
     assertThat(restoredCache.equalsForTest(cache)).isTrue()
   }

workflow-ui/container-android/src/main/java/com/squareup/workflow1/ui/backstack/BackStackContainer.kt

@@ -1,7 +1,9 @@
 package com.squareup.workflow1.ui.backstack
 
 import android.content.Context
+import android.os.Parcel
 import android.os.Parcelable
+import android.os.Parcelable.Creator
 import android.util.AttributeSet
 import android.view.Gravity
 import android.view.View
@@ -143,14 +145,16 @@ public open class BackStackContainer @JvmOverloads constructor(
     addView(newView)
   }
 
-  override fun onSaveInstanceState(): Parcelable {
-    return ViewStateCache.SavedState(super.onSaveInstanceState(), viewStateCache)
+  override fun onSaveInstanceState(): Parcelable? {
+    return super.onSaveInstanceState()?.let {
+      SavedState(it, viewStateCache.save())
+    }
   }
 
   override fun onRestoreInstanceState(state: Parcelable) {
-    (state as? ViewStateCache.SavedState)
+    (state as? SavedState)
       ?.let {
-        viewStateCache.restore(it.viewStateCache)
+        viewStateCache.restore(it.savedViewState)
         super.onRestoreInstanceState(state.superState)
       }
       ?: super.onRestoreInstanceState(super.onSaveInstanceState())
@@ -174,6 +178,36 @@ public open class BackStackContainer @JvmOverloads constructor(
     super.onDetachedFromWindow()
   }
 
+  public class SavedState : BaseSavedState {
+    public constructor(
+      superState: Parcelable,
+      savedViewState: ViewStateCache.Saved
+    ) : super(superState) {
+      this.savedViewState = savedViewState
+    }
+
+    public constructor(source: Parcel) : super(source) {
+      this.savedViewState = source.readParcelable(ViewStateCache.Saved::class.java.classLoader)!!
+    }
+
+    public val savedViewState: ViewStateCache.Saved
+
+    override fun writeToParcel(
+      out: Parcel,
+      flags: Int
+    ) {
+      super.writeToParcel(out, flags)
+      out.writeParcelable(savedViewState, flags)
+    }
+
+    public companion object CREATOR : Creator<ViewStateCache.Saved> {
+      override fun createFromParcel(source: Parcel): ViewStateCache.Saved =
+        ViewStateCache.Saved(source)
+
+      override fun newArray(size: Int): Array<ViewStateCache.Saved?> = arrayOfNulls(size)
+    }
+  }
+
   public companion object : ViewFactory<BackStackScreen<*>>
   by BuilderViewFactory(
     type = BackStackScreen::class,

workflow-ui/container-android/src/main/java/com/squareup/workflow1/ui/backstack/ViewStateCache.kt

@@ -5,26 +5,22 @@ import android.os.Parcelable
 import android.os.Parcelable.Creator
 import android.util.SparseArray
 import android.view.View
-import android.view.View.BaseSavedState
 import androidx.annotation.VisibleForTesting
 import androidx.annotation.VisibleForTesting.PRIVATE
 import androidx.savedstate.SavedStateRegistryOwner
 import androidx.savedstate.ViewTreeSavedStateRegistryOwner
 import com.squareup.workflow1.ui.Named
 import com.squareup.workflow1.ui.WorkflowUiExperimentalApi
 import com.squareup.workflow1.ui.androidx.WorkflowSavedStateRegistryAggregator
-import com.squareup.workflow1.ui.backstack.ViewStateCache.SavedState
 import com.squareup.workflow1.ui.getRendering
 
 /**
  * Handles persistence chores for container views that manage a set of [Named] renderings,
  * showing a view for one at a time -- think back stacks or tab sets.
  *
- * - This class implements [Parcelable] so that it can be preserved from
- *  a container view's own [View.saveHierarchyState] method. A simple container can
- *  return [SavedState] from that method rather than creating its own persistence class.
- *
- * - It also handles androidx [ViewTreeSavedStateRegistryOwner] duties, via
+ * - Provides [Parcelable]-based [save] and [restore] methods for use from a
+ *   container's [View.onSaveInstanceState] and [View.onRestoreInstanceState] methods.
+ * - Also handles androidx [ViewTreeSavedStateRegistryOwner] duties, via
  *  a wrapped instance of [WorkflowSavedStateRegistryAggregator]. This means that container
  *  views using this class must call [attachToParentRegistryOwner] and
  *  [detachFromParentRegistry] when they are [attached][View.onAttachedToWindow] and
@@ -36,7 +32,7 @@ public class ViewStateCache
 internal constructor(
   @VisibleForTesting(otherwise = PRIVATE)
   internal val viewStates: MutableMap<String, ViewStateFrame>
-) : Parcelable {
+) {
   public constructor() : this(mutableMapOf())
 
   private val stateRegistryAggregator = WorkflowSavedStateRegistryAggregator()
@@ -134,77 +130,54 @@ internal constructor(
    * Replaces the state of the receiver with that of [from]. Typical usage is to call this from
    * a container view's [View.onRestoreInstanceState].
    */
-  public fun restore(from: ViewStateCache) {
+  public fun restore(from: Saved) {
     viewStates.clear()
     viewStates += from.viewStates
   }
 
   /**
-   * Convenience for use in [View.onSaveInstanceState] and [View.onRestoreInstanceState]
-   * methods of container views that have no other state of their own to save.
-   *
-   * More interesting containers should create their own subclass of [BaseSavedState]
-   * rather than trying to extend this one.
+   * Returns a [Parcelable] copy of the internal state of the receiver, for use with
+   * a container view's [View.onSaveInstanceState].
    */
-  public class SavedState : BaseSavedState {
-    public constructor(
-      superState: Parcelable?,
-      viewStateCache: ViewStateCache
-    ) : super(superState) {
-      this.viewStateCache = viewStateCache
+  public fun save(): Saved {
+    return Saved(this)
+  }
+
+  public class Saved : Parcelable {
+    internal constructor(viewStateCache: ViewStateCache) {
+      this.viewStates = viewStateCache.viewStates.toMap()
     }
 
-    public constructor(source: Parcel) : super(source) {
-      this.viewStateCache = source.readParcelable(SavedState::class.java.classLoader)!!
+    public constructor(source: Parcel) {
+      this.viewStates = mutableMapOf<String, ViewStateFrame>()
+        .apply {
+          @Suppress("UNCHECKED_CAST")
+          source.readMap(
+            this as MutableMap<Any?, Any?>,
+            ViewStateCache::class.java.classLoader
+          )
+        }
+        .toMap()
     }
 
-    public val viewStateCache: ViewStateCache
+    internal val viewStates: Map<String, ViewStateFrame>
+
+    override fun describeContents(): Int = 0
 
     override fun writeToParcel(
       out: Parcel,
       flags: Int
     ) {
-      super.writeToParcel(out, flags)
-      out.writeParcelable(viewStateCache, flags)
+      out.writeMap(viewStates)
     }
 
-    public companion object CREATOR : Creator<SavedState> {
-      override fun createFromParcel(source: Parcel): SavedState =
-        SavedState(source)
+    public companion object CREATOR : Creator<Saved> {
+      override fun createFromParcel(source: Parcel): Saved =
+        Saved(source)
 
-      override fun newArray(size: Int): Array<SavedState?> = arrayOfNulls(size)
+      override fun newArray(size: Int): Array<Saved?> = arrayOfNulls(size)
     }
   }
-
-// region Parcelable
-
-  override fun describeContents(): Int = 0
-
-  override fun writeToParcel(
-    parcel: Parcel,
-    flags: Int
-  ) {
-    @Suppress("UNCHECKED_CAST")
-    parcel.writeMap(viewStates as MutableMap<Any?, Any?>)
-  }
-
-  public companion object CREATOR : Creator<ViewStateCache> {
-    override fun createFromParcel(parcel: Parcel): ViewStateCache {
-      @Suppress("UNCHECKED_CAST")
-      return mutableMapOf<String, ViewStateFrame>()
-        .apply {
-          parcel.readMap(
-            this as MutableMap<Any?, Any?>,
-            ViewStateCache::class.java.classLoader
-          )
-        }
-        .let { ViewStateCache(it) }
-    }
-
-    override fun newArray(size: Int): Array<ViewStateCache?> = arrayOfNulls(size)
-  }
-
-// endregion
 }
 
 @WorkflowUiExperimentalApi