Ensure that transitions are unregistered when the launchable is detached.

Ephemeral return animations are registered as remote transitions, and
unregistered after they run once. But if they're never run they risk
leaking, so we unregister them once the view they use is detached or the
composable is not being composed anymore.

This will help with usages of the library outside of Launcher (and maybe
Launcher itself in the future).

Bug: 339194555
Flag: com.android.systemui.shared.return_animation_framework_library
Test: manual
Change-Id: I28011978589e2dd840326539d79dfd0e09246e13

6 files changed, 114 insertions, 37 deletions

diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
index a1f0c146c507..41a00f5237f7 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/ActivityTransitionAnimator.kt
@@ -424,15 +424,16 @@ constructor(
                         newKeyguardOccludedState: Boolean?
                     ) {
                         super.onTransitionAnimationCancelled(newKeyguardOccludedState)
-                        cleanUp()
+                        onDispose()
                     }
 
                     override fun onTransitionAnimationEnd(isExpandingFullyAbove: Boolean) {
                         super.onTransitionAnimationEnd(isExpandingFullyAbove)
-                        cleanUp()
+                        onDispose()
                     }
 
-                    private fun cleanUp() {
+                    override fun onDispose() {
+                        super.onDispose()
                         cleanUpRunnable?.run()
                     }
                 }
@@ -560,6 +561,7 @@ constructor(
                 cookie: TransitionCookie? = null,
                 component: ComponentName? = null,
                 returnCujType: Int? = null,
+                isEphemeral: Boolean = true,
             ): Controller? {
                 // Make sure the View we launch from implements LaunchableView to avoid visibility
                 // issues.
@@ -587,6 +589,7 @@ constructor(
                     cookie,
                     component,
                     returnCujType,
+                    isEphemeral,
                 )
             }
         }
@@ -647,6 +650,9 @@ constructor(
          * appropriately.
          */
         fun onTransitionAnimationCancelled(newKeyguardOccludedState: Boolean? = null) {}
+
+        /** The controller will not be used again. Clean up the relevant internal state. */
+        fun onDispose() {}
     }
 
     /**
diff --git a/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt b/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
index 3ba9a2974846..b56a68cb2dd6 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/Expandable.kt
@@ -39,7 +39,8 @@ interface Expandable {
         launchCujType: Int? = null,
         cookie: ActivityTransitionAnimator.TransitionCookie? = null,
         component: ComponentName? = null,
-        returnCujType: Int? = null
+        returnCujType: Int? = null,
+        isEphemeral: Boolean = true,
     ): ActivityTransitionAnimator.Controller?
 
     /**
@@ -55,7 +56,8 @@ interface Expandable {
             launchCujType,
             cookie = null,
             component = null,
-            returnCujType = null
+            returnCujType = null,
+            isEphemeral = true,
         )
     }
 
@@ -80,14 +82,16 @@ interface Expandable {
                     launchCujType: Int?,
                     cookie: ActivityTransitionAnimator.TransitionCookie?,
                     component: ComponentName?,
-                    returnCujType: Int?
+                    returnCujType: Int?,
+                    isEphemeral: Boolean,
                 ): ActivityTransitionAnimator.Controller? {
                     return ActivityTransitionAnimator.Controller.fromView(
                         view,
                         launchCujType,
                         cookie,
                         component,
-                        returnCujType
+                        returnCujType,
+                        isEphemeral,
                     )
                 }
 
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 e626c04675e1..558c1eba2c1c 100644
--- a/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
+++ b/packages/SystemUI/animation/src/com/android/systemui/animation/GhostedViewTransitionAnimatorController.kt
@@ -67,6 +67,12 @@ constructor(
 
     /** The [CujType] associated to this return animation. */
     private val returnCujType: Int? = null,
+
+    /**
+     * Whether this controller should be invalidated after its first use, and whenever [ghostedView]
+     * is detached.
+     */
+    private val isEphemeral: Boolean = false,
     private var interactionJankMonitor: InteractionJankMonitor =
         InteractionJankMonitor.getInstance(),
 ) : ActivityTransitionAnimator.Controller {
@@ -119,6 +125,19 @@ constructor(
                 returnCujType
             }
 
+    /**
+     * Used to automatically clean up the internal state once [ghostedView] is detached from the
+     * hierarchy.
+     */
+    private val detachListener =
+        object : View.OnAttachStateChangeListener {
+            override fun onViewAttachedToWindow(v: View) {}
+
+            override fun onViewDetachedFromWindow(v: View) {
+                onDispose()
+            }
+        }
+
     init {
         // Make sure the View we launch from implements LaunchableView to avoid visibility issues.
         if (ghostedView !is LaunchableView) {
@@ -155,6 +174,16 @@ constructor(
         }
 
         background = findBackground(ghostedView)
+
+        if (TransitionAnimator.returnAnimationsEnabled() && isEphemeral) {
+            ghostedView.addOnAttachStateChangeListener(detachListener)
+        }
+    }
+
+    override fun onDispose() {
+        if (TransitionAnimator.returnAnimationsEnabled()) {
+            ghostedView.removeOnAttachStateChangeListener(detachListener)
+        }
     }
 
     /**
@@ -164,7 +193,7 @@ constructor(
     protected open fun setBackgroundCornerRadius(
         background: Drawable,
         topCornerRadius: Float,
-        bottomCornerRadius: Float
+        bottomCornerRadius: Float,
     ) {
         // By default, we rely on WrappedDrawable to set/restore the background radii before/after
         // each draw.
@@ -195,7 +224,7 @@ constructor(
         val state =
             TransitionAnimator.State(
                 topCornerRadius = getCurrentTopCornerRadius(),
-                bottomCornerRadius = getCurrentBottomCornerRadius()
+                bottomCornerRadius = getCurrentBottomCornerRadius(),
             )
         fillGhostedViewState(state)
         return state
@@ -269,7 +298,7 @@ constructor(
     override fun onTransitionAnimationProgress(
         state: TransitionAnimator.State,
         progress: Float,
-        linearProgress: Float
+        linearProgress: Float,
     ) {
         val ghostView = this.ghostView ?: return
         val backgroundView = this.backgroundView!!
@@ -317,11 +346,11 @@ constructor(
             scale,
             scale,
             ghostedViewState.centerX - transitionContainerLocation[0],
-            ghostedViewState.centerY - transitionContainerLocation[1]
+            ghostedViewState.centerY - transitionContainerLocation[1],
         )
         ghostViewMatrix.postTranslate(
             (leftChange + rightChange) / 2f,
-            (topChange + bottomChange) / 2f
+            (topChange + bottomChange) / 2f,
         )
         ghostView.animationMatrix = ghostViewMatrix
 
@@ -462,7 +491,7 @@ constructor(
         private fun updateRadii(
             radii: FloatArray,
             topCornerRadius: Float,
-            bottomCornerRadius: Float
+            bottomCornerRadius: Float,
         ) {
             radii[0] = topCornerRadius
             radii[1] = topCornerRadius
diff --git a/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt b/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
index a55df2b36a80..103a9b5cf5f4 100644
--- a/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
+++ b/packages/SystemUI/compose/core/src/com/android/compose/animation/ExpandableController.kt
@@ -52,6 +52,9 @@ import kotlin.math.roundToInt
 interface ExpandableController {
     /** The [Expandable] controlled by this controller. */
     val expandable: Expandable
+
+    /** Called when the [Expandable] stop being included in the composition. */
+    fun onDispose()
 }
 
 /**
@@ -88,33 +91,44 @@ fun rememberExpandableController(
     // Whether this composable is still composed. We only do the dialog exit animation if this is
     // true.
     val isComposed = remember { mutableStateOf(true) }
-    DisposableEffect(Unit) { onDispose { isComposed.value = false } }
-
-    return remember(
-        color,
-        contentColor,
-        shape,
-        borderStroke,
-        composeViewRoot,
-        density,
-        layoutDirection,
-    ) {
-        ExpandableControllerImpl(
+
+    val controller =
+        remember(
             color,
             contentColor,
             shape,
             borderStroke,
             composeViewRoot,
             density,
-            animatorState,
-            isDialogShowing,
-            overlay,
-            currentComposeViewInOverlay,
-            boundsInComposeViewRoot,
             layoutDirection,
-            isComposed,
-        )
+        ) {
+            ExpandableControllerImpl(
+                color,
+                contentColor,
+                shape,
+                borderStroke,
+                composeViewRoot,
+                density,
+                animatorState,
+                isDialogShowing,
+                overlay,
+                currentComposeViewInOverlay,
+                boundsInComposeViewRoot,
+                layoutDirection,
+                isComposed,
+            )
+        }
+
+    DisposableEffect(Unit) {
+        onDispose {
+            isComposed.value = false
+            if (TransitionAnimator.returnAnimationsEnabled()) {
+                controller.onDispose()
+            }
+        }
     }
+
+    return controller
 }
 
 internal class ExpandableControllerImpl(
@@ -132,19 +146,29 @@ internal class ExpandableControllerImpl(
     private val layoutDirection: LayoutDirection,
     private val isComposed: State<Boolean>,
 ) : ExpandableController {
+    /** The [ActivityTransitionAnimator.Controller] to be cleaned up [onDispose]. */
+    private var activityControllerForDisposal: ActivityTransitionAnimator.Controller? = null
+
     override val expandable: Expandable =
         object : Expandable {
             override fun activityTransitionController(
                 launchCujType: Int?,
                 cookie: ActivityTransitionAnimator.TransitionCookie?,
                 component: ComponentName?,
-                returnCujType: Int?
+                returnCujType: Int?,
+                isEphemeral: Boolean,
             ): ActivityTransitionAnimator.Controller? {
                 if (!isComposed.value) {
                     return null
                 }
 
-                return activityController(launchCujType, cookie, component, returnCujType)
+                val controller = activityController(launchCujType, cookie, component, returnCujType)
+                if (TransitionAnimator.returnAnimationsEnabled() && isEphemeral) {
+                    activityControllerForDisposal?.onDispose()
+                    activityControllerForDisposal = controller
+                }
+
+                return controller
             }
 
             override fun dialogTransitionController(
@@ -158,6 +182,11 @@ internal class ExpandableControllerImpl(
             }
         }
 
+    override fun onDispose() {
+        activityControllerForDisposal?.onDispose()
+        activityControllerForDisposal = null
+    }
+
     /**
      * Create a [TransitionAnimator.Controller] that is going to be used to drive an activity or
      * dialog animation. This controller will:
@@ -181,7 +210,7 @@ internal class ExpandableControllerImpl(
             override fun onTransitionAnimationProgress(
                 state: TransitionAnimator.State,
                 progress: Float,
-                linearProgress: Float
+                linearProgress: Float,
             ) {
                 // We copy state given that it's always the same object that is mutated by
                 // ActivityTransitionAnimator.
@@ -269,7 +298,7 @@ internal class ExpandableControllerImpl(
         launchCujType: Int?,
         cookie: ActivityTransitionAnimator.TransitionCookie?,
         component: ComponentName?,
-        returnCujType: Int?
+        returnCujType: Int?,
     ): ActivityTransitionAnimator.Controller {
         val delegate = transitionController()
         return object :
diff --git a/packages/SystemUI/src/com/android/systemui/haptics/msdl/qs/StateAwareExpandable.kt b/packages/SystemUI/src/com/android/systemui/haptics/msdl/qs/StateAwareExpandable.kt
index 215ceacaef14..0ed4007d68ef 100644
--- a/packages/SystemUI/src/com/android/systemui/haptics/msdl/qs/StateAwareExpandable.kt
+++ b/packages/SystemUI/src/com/android/systemui/haptics/msdl/qs/StateAwareExpandable.kt
@@ -78,9 +78,16 @@ fun Expandable.withStateAwareness(
             cookie: ActivityTransitionAnimator.TransitionCookie?,
             component: ComponentName?,
             returnCujType: Int?,
+            isEphemeral: Boolean,
         ): ActivityTransitionAnimator.Controller? =
             delegate
-                .activityTransitionController(launchCujType, cookie, component, returnCujType)
+                .activityTransitionController(
+                    launchCujType,
+                    cookie,
+                    component,
+                    returnCujType,
+                    isEphemeral,
+                )
                 ?.withStateAwareness(onActivityLaunchTransitionStart, onActivityLaunchTransitionEnd)
 
         override fun dialogTransitionController(
diff --git a/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt b/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
index b82aa817afd8..1504402279b4 100644
--- a/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
+++ b/packages/SystemUI/src/com/android/systemui/haptics/qs/QSLongPressEffect.kt
@@ -284,6 +284,7 @@ constructor(
                     cookie: ActivityTransitionAnimator.TransitionCookie?,
                     component: ComponentName?,
                     returnCujType: Int?,
+                    isEphemeral: Boolean,
                 ): ActivityTransitionAnimator.Controller? {
                     val delegatedController =
                         ActivityTransitionAnimator.Controller.fromView(
@@ -292,6 +293,7 @@ constructor(
                             cookie,
                             component,
                             returnCujType,
+                            isEphemeral,
                         )
                     return delegatedController?.let { createTransitionControllerDelegate(it) }
                 }