Allow async restoring of view state

leakcanary/leakcanary-android-core/src/main/java/leakcanary/internal/activity/screen/HeapDumpsScreen.kt

@@ -21,6 +21,7 @@ import leakcanary.internal.navigation.goTo
 import leakcanary.internal.navigation.inflate
 import leakcanary.internal.navigation.onCreateOptionsMenu
 import leakcanary.internal.navigation.onScreenExiting
+import leakcanary.internal.navigation.restoreViewStateFromTag
 
 internal class HeapDumpsScreen : Screen() {
   override fun createView(container: ViewGroup) =
@@ -111,5 +112,6 @@ internal class HeapDumpsScreen : Screen() {
         )
         countView.text = count
       }
+    restoreViewStateFromTag()
   }
 }

leakcanary/leakcanary-android-core/src/main/java/leakcanary/internal/activity/screen/LeaksScreen.kt

@@ -19,6 +19,7 @@ import leakcanary.internal.navigation.activity
 import leakcanary.internal.navigation.goTo
 import leakcanary.internal.navigation.inflate
 import leakcanary.internal.navigation.onScreenExiting
+import leakcanary.internal.navigation.restoreViewStateFromTag
 
 internal class LeaksScreen : Screen() {
   override fun createView(container: ViewGroup) =
@@ -70,5 +71,6 @@ internal class LeaksScreen : Screen() {
     listView.setOnItemClickListener { _, _, position, _ ->
       goTo(LeakScreen(projections[position].signature))
     }
+    restoreViewStateFromTag()
   }
-}
+}

leakcanary/leakcanary-android-core/src/main/java/leakcanary/internal/navigation/BackstackFrame.kt

@@ -4,6 +4,7 @@ import android.os.Parcel
 import android.os.Parcelable
 import android.util.SparseArray
 import android.view.View
+import com.squareup.leakcanary.core.R
 
 internal class BackstackFrame : Parcelable {
 
@@ -32,11 +33,15 @@ internal class BackstackFrame : Parcelable {
     this.screen = screen
     viewState = SparseArray()
     view.saveHierarchyState(viewState)
+    view.setTag(R.id.leak_canary_restored_view_state, null)
   }
 
   fun restore(view: View) {
     if (viewState != null) {
       view.restoreHierarchyState(viewState)
+      view.setTag(R.id.leak_canary_restored_view_state, viewState)
+    } else {
+      view.setTag(R.id.leak_canary_restored_view_state, null)
     }
   }
 

leakcanary/leakcanary-android-core/src/main/java/leakcanary/internal/navigation/Views.kt

@@ -3,6 +3,8 @@ package leakcanary.internal.navigation
 import android.app.Activity
 import android.content.Context
 import android.os.Build.VERSION
+import android.os.Parcelable
+import android.util.SparseArray
 import android.view.LayoutInflater
 import android.view.Menu
 import android.view.View
@@ -12,6 +14,13 @@ import com.squareup.leakcanary.core.R
 internal fun ViewGroup.inflate(layoutResId: Int) = LayoutInflater.from(context)
   .inflate(layoutResId, this, false)!!
 
+internal fun View.restoreViewStateFromTag() {
+  val viewState = getTag(R.id.leak_canary_restored_view_state) as SparseArray<Parcelable>?
+  if (viewState != null) {
+    restoreHierarchyState(viewState)
+  }
+}
+
 internal val View.activity
   get() = context as Activity
 
@@ -54,4 +63,4 @@ internal fun View.notifyScreenExiting() {
   val callbacks = getTag(R.id.leak_canary_notification_on_screen_exit)
     as MutableList<() -> Unit>?
   callbacks?.forEach { it.invoke() }
-}
+}

leakcanary/leakcanary-android-core/src/main/res/values/leak_canary_ids.xml

@@ -22,4 +22,5 @@
   <item type="id" name="leak_canary_notification_retained_objects" />
   <item type="id" name="leak_canary_notification_no_retained_object_on_tap" />
   <item type="id" name="leak_canary_notification_on_screen_exit" />
+  <item type="id" name="leak_canary_restored_view_state" />
 </resources>