[AnimLib] Reuse same unique token for same view

This change enables clients of ViewTransitionRegistry to reuse same
uniue token associated with a particular view (till it is attached to a
window). This optimization reduces the overhead of generating different
unique tokens mapped to the same view and ensures that there is a single
token, entry and listener associated with a view in the registry (this
detail is completely abstracted out from clients).

Bug: 393241010
Flag: com.android.systemui.decouple_view_controller_in_animlib
Test: GhostedViewTransitionAnimatorControllerTest, ViewTransitionRegistryTest
Change-Id: If0dbd9de5809068476cfe69268b17429f5dda11d

5 files changed, 137 insertions, 39 deletions

diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
index 444389fb26ea..fdb07bdbe7f3 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
@@ -155,7 +155,7 @@ constructor(
     /** [ViewTransitionToken] to be used for storing transitioning view in [transitionRegistry] */
     private val transitionToken =
         if (Flags.decoupleViewControllerInAnimlib()) {
-            ViewTransitionToken(transitioningView::class.java)
+            transitionRegistry?.register(transitioningView)
         } else {
             null
         }
@@ -164,7 +164,7 @@ constructor(
     private val ghostedView: View
         get() =
             if (Flags.decoupleViewControllerInAnimlib()) {
-                transitionRegistry?.getView(transitionToken!!)
+                transitionToken?.let { token -> transitionRegistry?.getView(token) }
             } else {
                 _ghostedView
             }!!
@@ -186,10 +186,6 @@ constructor(
             )
         }
 
-        if (Flags.decoupleViewControllerInAnimlib()) {
-            transitionRegistry?.register(transitionToken!!, transitioningView)
-        }
-
         /** Find the first view with a background in [view] and its children. */
         fun findBackground(view: View): Drawable? {
             if (view.background != null) {
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/IViewTransitionRegistry.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/IViewTransitionRegistry.kt
index af3ca87bf788..280d90de7ace 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/IViewTransitionRegistry.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/IViewTransitionRegistry.kt
@@ -22,12 +22,12 @@ import android.view.View
 interface IViewTransitionRegistry {
 
     /**
-     * Registers the transitioning [view] mapped to a [token]
+     * Registers the transitioning [view] mapped to returned token
      *
-     * @param token The token corresponding to the transitioning view
      * @param view The view undergoing transition
+     * @return token mapped to the transitioning view
      */
-    fun register(token: ViewTransitionToken, view: View)
+    fun register(view: View): ViewTransitionToken
 
     /**
      * Unregisters the transitioned view from its corresponding [token]
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/ViewTransitionRegistry.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/ViewTransitionRegistry.kt
index 86c7f76c6bee..882ff3b61ba9 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/ViewTransitionRegistry.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/ViewTransitionRegistry.kt
@@ -22,21 +22,21 @@ import java.lang.ref.WeakReference
 
 /**
  * A registry to temporarily store the view being transitioned into a Dialog (using
- * [DialogTransitionAnimator]) or an Activity (using [ActivityTransitionAnimator])
+ * [DialogTransitionAnimator]) or an Activity (using [ActivityTransitionAnimator]).
  */
 class ViewTransitionRegistry : IViewTransitionRegistry {
 
     /**
      * A map of a unique token to a WeakReference of the View being transitioned. WeakReference
      * ensures that Views are garbage collected whenever they become eligible and avoid any
-     * memory leaks
+     * memory leaks.
      */
-    private val registry by lazy { mutableMapOf<ViewTransitionToken, WeakReference<View>>() }
+    private val registry by lazy { mutableMapOf<ViewTransitionToken, ViewTransitionInfo>() }
 
     /**
      * A [View.OnAttachStateChangeListener] to be attached to all views stored in the registry to
      * ensure that views (and their corresponding entry) is automatically removed when the view is
-     * detached from the Window
+     * detached from the Window.
      */
     private val listener by lazy {
         object : View.OnAttachStateChangeListener {
@@ -45,74 +45,121 @@ class ViewTransitionRegistry : IViewTransitionRegistry {
             }
 
             override fun onViewDetachedFromWindow(view: View) {
-                getViewToken(view)?.let { token -> unregister(token) }
+                // if view is detached from window, remove it from registry irrespective of number
+                // of reference held by clients/user of this registry
+                getViewToken(view)?.let { token -> remove(token) }
             }
         }
     }
 
     /**
-     * Creates an entry of a unique "token" mapped to "transitioning view" in the registry
+     * Creates an entry of a unique token mapped to transitioning [view] in the registry.
      *
-     * @param token unique token associated with the transitioning view
      * @param view view undergoing transitions
+     * @return unique token mapped to the view being registered
      */
-    override fun register(token: ViewTransitionToken, view: View) {
+    override fun register(view: View): ViewTransitionToken {
+        // if view being registered is already present in the registry and has a unique token
+        // assigned to it, reuse that token
+        getViewToken(view)?.let { token ->
+            registry[token]?.let { info -> info.viewRefCount += 1 }
+            return token
+        }
+
         // token embedded as a view tag enables to use a single listener for all views
+        val token = ViewTransitionToken(view::class.java)
         view.setTag(R.id.tag_view_transition_token, token)
         view.addOnAttachStateChangeListener(listener)
-        registry[token] = WeakReference(view)
+        registry[token] = ViewTransitionInfo(WeakReference(view))
         onRegistryUpdate()
+
+        return token
     }
 
     /**
-     * Removes the entry associated with the unique "token" in the registry
+     * Unregisters a view mapped to the unique [token] in the registry. This will either remove the
+     * entry entirely from registry (if the reference count of the associated view reached zero) or
+     * will decrement the reference count of the associated view in the registry.
      *
      * @param token unique token associated with the transitioning view
      */
     override fun unregister(token: ViewTransitionToken) {
-        registry.remove(token)?.let {
-            it.get()?.let { view ->
+        registry[token]?.let { info ->
+            info.viewRefCount -= 1
+            if (info.viewRefCount == 0) {
+                remove(token)
+            }
+        }
+    }
+
+    /**
+     * Removes the entry associated with the unique [token] in the registry.
+     *
+     * @param token unique token associated with the transitioning view
+     */
+    private fun remove(token: ViewTransitionToken) {
+        registry.remove(token)?.let { removedInfo ->
+            removedInfo.viewRef.get()?.let { view ->
                 view.removeOnAttachStateChangeListener(listener)
                 view.setTag(R.id.tag_view_transition_token, null)
             }
-            it.clear()
+            removedInfo.viewRef.clear()
             onRegistryUpdate()
         }
     }
 
     /**
-     * Access a view from registry using unique "token" associated with it
+     * Access a view from registry using unique [token] associated with it.
      * WARNING - this returns a StrongReference to the View stored in the registry
      */
     override fun getView(token: ViewTransitionToken): View? {
-        return registry[token]?.get()
+        return registry[token]?.viewRef?.get()
     }
 
     /**
-     * Return token mapped to the [view], if it is present in the registry
+     * Return token mapped to the [view], if it is present in the registry.
      *
      * @param view the transitioning view whose token we are requesting
      * @return token associated with the [view] if present, else null
      */
     override fun getViewToken(view: View): ViewTransitionToken? {
-        return (view.getTag(R.id.tag_view_transition_token) as? ViewTransitionToken)?.let { token ->
-            getView(token)?.let { token }
+        // extract token from the view if it is embedded inside it as a tag
+        val token = view.getTag(R.id.tag_view_transition_token) as? ViewTransitionToken
+
+        // this should never really happen, but if token embedded inside the view as tag, doesn't
+        // point to a valid view in the registry, remove that token (tag) from the view and registry
+        if (token != null && getView(token) == null) {
+            view.setTag(R.id.tag_view_transition_token, null)
+            remove(token)
+            return null
         }
+
+        return token
     }
 
-    /** Event call to run on registry update (on both [register] and [unregister]) */
+    /** Event call to run on registry update (on both [register] and [unregister]). */
     override fun onRegistryUpdate() {
         emitCountForTrace()
     }
 
     /**
      * Utility function to emit number of non-null views in the registry whenever the registry is
-     * updated (via [register] or [unregister])
+     * updated (via [register] or [unregister]).
      */
     private fun emitCountForTrace() {
         Trace.setCounter("transition_registry_view_count", registry.count().toLong())
     }
 
+    /** Information associated with each transitioning view in the registry. */
+    private data class ViewTransitionInfo(
+
+        /** View being transitioned */
+        val viewRef: WeakReference<View>,
+
+        /** Count of clients (users of this registry) referencing same transitioning view */
+        var viewRefCount: Int = 1
+    )
+
     companion object {
         val instance by lazy(LazyThreadSafetyMode.SYNCHRONIZED) { ViewTransitionRegistry() }
     }
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
index 052d520ac92f..18b68d2fa8a3 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/GhostedViewTransitionAnimatorControllerTest.kt
@@ -153,10 +153,12 @@ class GhostedViewTransitionAnimatorControllerTest : SysuiTestCase() {
     private class FakeViewTransitionRegistry : IViewTransitionRegistry {
 
         val registry = mutableMapOf<ViewTransitionToken, View>()
+        val token = ViewTransitionToken()
 
-        override fun register(token: ViewTransitionToken, view: View) {
+        override fun register(view: View): ViewTransitionToken {
             registry[token] = view
             view.setTag(R.id.tag_view_transition_token, token)
+            return token
         }
 
         override fun unregister(token: ViewTransitionToken) {
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/ViewTransitionRegistryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/ViewTransitionRegistryTest.kt
index ef91c793a2f3..b18eafd206ca 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/ViewTransitionRegistryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/animation/ViewTransitionRegistryTest.kt
@@ -25,9 +25,9 @@ import com.google.common.truth.Truth.assertThat
 import org.junit.Before
 import org.junit.runner.RunWith
 import org.mockito.kotlin.argumentCaptor
-import org.mockito.kotlin.doReturn
 import org.mockito.kotlin.mock
 import org.mockito.kotlin.verify
+import org.mockito.kotlin.whenever
 import kotlin.test.Test
 
 @SmallTest
@@ -36,24 +36,22 @@ class ViewTransitionRegistryTest : SysuiTestCase() {
 
     private lateinit var view: View
     private lateinit var underTest: ViewTransitionRegistry
-    private var token: ViewTransitionToken = ViewTransitionToken()
 
     @Before
     fun setup() {
         view = FrameLayout(mContext)
         underTest = ViewTransitionRegistry()
-        token = ViewTransitionToken()
     }
 
     @Test
     fun testSuccessfulRegisterInViewTransitionRegistry() {
-        underTest.register(token, view)
+        val token = underTest.register(view)
         assertThat(underTest.getView(token)).isNotNull()
     }
 
     @Test
     fun testSuccessfulUnregisterInViewTransitionRegistry() {
-        underTest.register(token, view)
+        val token = underTest.register(view)
         assertThat(underTest.getView(token)).isNotNull()
 
         underTest.unregister(token)
@@ -62,13 +60,14 @@ class ViewTransitionRegistryTest : SysuiTestCase() {
 
     @Test
     fun testSuccessfulUnregisterOnViewDetachedFromWindow() {
-        val view: View = mock {
-            on { getTag(R.id.tag_view_transition_token) } doReturn token
-        }
+        val view: View = mock()
 
-        underTest.register(token, view)
+        val token = underTest.register(view)
+        assertThat(token).isEqualTo(token)
         assertThat(underTest.getView(token)).isNotNull()
 
+        whenever(view.getTag(R.id.tag_view_transition_token)).thenReturn(token)
+
         argumentCaptor<View.OnAttachStateChangeListener>()
             .apply { verify(view).addOnAttachStateChangeListener(capture()) }
             .firstValue
@@ -76,4 +75,58 @@ class ViewTransitionRegistryTest : SysuiTestCase() {
 
         assertThat(underTest.getView(token)).isNull()
     }
+
+    @Test
+    fun testMultipleRegisterOnSameView() {
+        val token = underTest.register(view)
+
+        // multiple register on same view should return same token
+        assertThat(underTest.register(view)).isEqualTo(token)
+
+        // 1st unregister doesn't remove the token from registry as refCount = 2
+        underTest.unregister(token)
+        assertThat(underTest.getView(token)).isNotNull()
+
+        // 2nd unregister removes the token from registry
+        underTest.unregister(token)
+        assertThat(underTest.getView(token)).isNull()
+    }
+
+    @Test
+    fun testMultipleRegisterOnSameViewRemovedAfterViewDetached() {
+        val view: View = mock()
+
+        val token = underTest.register(view)
+        whenever(view.getTag(R.id.tag_view_transition_token)).thenReturn(token)
+
+        assertThat(underTest.getViewToken(view)).isEqualTo(token)
+
+        // mock view's detach event
+        val caller = argumentCaptor<View.OnAttachStateChangeListener>()
+            .apply { verify(view).addOnAttachStateChangeListener(capture()) }
+            .firstValue
+
+        // register 3 times
+        underTest.register(view)
+        underTest.register(view)
+        underTest.register(view)
+
+        // unregister 1 time and verify entry should still be present in registry
+        underTest.unregister(token)
+        assertThat(underTest.getView(token)).isNotNull()
+
+        // view's associated entry should be gone from registry, after view detaches
+        caller.onViewDetachedFromWindow(view)
+        assertThat(underTest.getView(token)).isNull()
+    }
+
+    @Test
+    fun testDistinctViewsSameClassRegisterWithDifferentToken() {
+        var prev: ViewTransitionToken? = underTest.register(FrameLayout(mContext))
+        for (i in 0 until 10) {
+            val curr = underTest.register(FrameLayout(mContext))
+            assertThat(curr).isNotEqualTo(prev)
+            prev = curr
+        }
+    }
 }