Add return animations to Toast search results.

This only applies to targets that already use the Animation library for
launch animations.

Bug: 336719605
Bug: 298089923
Flag: com.android.launcher3.enable_container_return_animations
Flag: com.android.systemui.shared.return_animation_framework_library
Test: tested manually with flag on and off
Change-Id: Ib824e78fa8b1b226b32d23d8325f06b496ba5deb
diff --git a/aconfig/launcher.aconfig b/aconfig/launcher.aconfig
index f1f9966..15ac9e3 100644
--- a/aconfig/launcher.aconfig
+++ b/aconfig/launcher.aconfig
@@ -311,6 +311,13 @@
 }
 
 flag {
+    name: "enable_container_return_animations"
+    namespace: "launcher"
+    description: "Enables the container return animation mirroring launches."
+    bug: "341017746"
+}
+
+flag {
     name: "floating_search_bar"
     namespace: "launcher"
     description: "Search bar persists at the bottom of the screen across Launcher states"
diff --git a/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java b/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
index 15180ef..d973149 100644
--- a/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
+++ b/quickstep/src/com/android/launcher3/LauncherAnimationRunner.java
@@ -238,5 +238,12 @@
         @Override
         @UiThread
         default void onAnimationCancelled() {}
+
+        /**
+         * Returns whether this animation factory supports a tightly coupled return animation.
+         */
+        default boolean supportsReturnTransition() {
+            return false;
+        }
     }
 }
diff --git a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
index fae281a..5a74f4a 100644
--- a/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
+++ b/quickstep/src/com/android/launcher3/QuickstepTransitionManager.java
@@ -43,6 +43,7 @@
 import static com.android.launcher3.BaseActivity.INVISIBLE_BY_APP_TRANSITIONS;
 import static com.android.launcher3.BaseActivity.INVISIBLE_BY_PENDING_FLAGS;
 import static com.android.launcher3.BaseActivity.PENDING_INVISIBLE_BY_WALLPAPER_ANIMATION;
+import static com.android.launcher3.Flags.enableContainerReturnAnimations;
 import static com.android.launcher3.Flags.enableScalingRevealHomeAnimation;
 import static com.android.launcher3.LauncherAnimUtils.SCALE_PROPERTY;
 import static com.android.launcher3.LauncherAnimUtils.VIEW_BACKGROUND_COLOR;
@@ -68,6 +69,7 @@
 import static com.android.quickstep.TaskViewUtils.findTaskViewToLaunch;
 import static com.android.quickstep.util.AnimUtils.clampToDuration;
 import static com.android.quickstep.util.AnimUtils.completeRunnableListCallback;
+import static com.android.systemui.shared.Flags.returnAnimationFrameworkLibrary;
 import static com.android.systemui.shared.system.QuickStepContract.getWindowCornerRadius;
 import static com.android.systemui.shared.system.QuickStepContract.supportsRoundedCornersOnWindows;
 
@@ -181,6 +183,9 @@
  */
 public class QuickstepTransitionManager implements OnDeviceProfileChangeListener {
 
+    private static final String TRANSITION_COOKIE_PREFIX =
+            "com.android.launcher3.QuickstepTransitionManager_activityLaunch";
+
     private static final boolean ENABLE_SHELL_STARTING_SURFACE =
             SystemProperties.getBoolean("persist.debug.shell_starting_surface", true);
 
@@ -333,17 +338,7 @@
         restartedListener.register(onEndCallback::executeAllAndDestroy);
         onEndCallback.add(restartedListener::unregister);
 
-        mAppLaunchRunner = new AppLaunchAnimationRunner(v, onEndCallback);
-        ItemInfo tag = (ItemInfo) v.getTag();
-        if (tag != null && tag.shouldUseBackgroundAnimation()) {
-            ContainerAnimationRunner containerAnimationRunner = ContainerAnimationRunner.from(
-                    v, mLauncher, mStartingWindowListener, onEndCallback);
-            if (containerAnimationRunner != null) {
-                mAppLaunchRunner = containerAnimationRunner;
-            }
-        }
-        RemoteAnimationRunnerCompat runner = new LauncherAnimationRunner(
-                mHandler, mAppLaunchRunner, true /* startAtFrontOfQueue */);
+        RemoteAnimationRunnerCompat runner = createAppLaunchRunner(v, onEndCallback);
 
         // Note that this duration is a guess as we do not know if the animation will be a
         // recents launch or not for sure until we know the opening app targets.
@@ -360,10 +355,95 @@
         IRemoteCallback endCallback = completeRunnableListCallback(onEndCallback);
         options.setOnAnimationAbortListener(endCallback);
         options.setOnAnimationFinishedListener(endCallback);
+
+        IBinder cookie = mAppLaunchRunner.supportsReturnTransition()
+                ? ((ContainerAnimationRunner) mAppLaunchRunner).getCookie() : null;
+        addLaunchCookie(cookie, (ItemInfo) v.getTag(), options);
+
+        // Register the return animation so it can be triggered on back from the app to home.
+        maybeRegisterAppReturnTransition(v);
+
         return new ActivityOptionsWrapper(options, onEndCallback);
     }
 
     /**
+     * Selects the appropriate type of launch runner for the given view, builds it, and returns it.
+     * {@link QuickstepTransitionManager#mAppLaunchRunner} is updated as a by-product of this
+     * method.
+     */
+    private RemoteAnimationRunnerCompat createAppLaunchRunner(View v, RunnableList onEndCallback) {
+        ItemInfo tag = (ItemInfo) v.getTag();
+        ContainerAnimationRunner containerRunner = null;
+        if (tag != null && tag.shouldUseBackgroundAnimation()) {
+            // The cookie should only override the default used by launcher if container return
+            // animations are enabled.
+            ActivityTransitionAnimator.TransitionCookie cookie =
+                    checkReturnAnimationsFlags()
+                            ? new ActivityTransitionAnimator.TransitionCookie(
+                                    TRANSITION_COOKIE_PREFIX + tag.id)
+                            : null;
+            ContainerAnimationRunner launchAnimationRunner =
+                    ContainerAnimationRunner.fromView(
+                            v, cookie, true /* forLaunch */, mLauncher, mStartingWindowListener,
+                            onEndCallback);
+
+            if (launchAnimationRunner != null) {
+                containerRunner = launchAnimationRunner;
+            }
+        }
+
+        mAppLaunchRunner = containerRunner != null
+                ? containerRunner : new AppLaunchAnimationRunner(v, onEndCallback);
+        return new LauncherAnimationRunner(
+                mHandler, mAppLaunchRunner, true /* startAtFrontOfQueue */);
+    }
+
+    /**
+     * If container return animations are enabled and the current launch runner is itself a
+     * {@link ContainerAnimationRunner}, registers a matching return animation that de-registers
+     * itself after it has run once or is made obsolete by the view going away.
+     */
+    private void maybeRegisterAppReturnTransition(View v) {
+        if (!checkReturnAnimationsFlags() || !mAppLaunchRunner.supportsReturnTransition()) {
+            return;
+        }
+
+        ActivityTransitionAnimator.TransitionCookie cookie =
+                ((ContainerAnimationRunner) mAppLaunchRunner).getCookie();
+        RunnableList onEndCallback = new RunnableList();
+        ContainerAnimationRunner runner =
+                ContainerAnimationRunner.fromView(
+                        v, cookie, false /* forLaunch */, mLauncher, mStartingWindowListener,
+                        onEndCallback);
+        RemoteTransition transition =
+                new RemoteTransition(
+                        new LauncherAnimationRunner(
+                                mHandler, runner, true /* startAtFrontOfQueue */
+                        ).toRemoteTransition()
+                );
+
+        SystemUiProxy.INSTANCE.get(mLauncher).registerRemoteTransition(
+                transition, ContainerAnimationRunner.buildBackToHomeFilter(cookie, mLauncher));
+        ContainerAnimationRunner.setUpRemoteAnimationCleanup(
+                v, transition, onEndCallback, mLauncher);
+    }
+
+    /**
+     * Adds a new launch cookie for the activity launch if supported.
+     * Prioritizes the explicitly provided cookie, falling back on extracting one from the given
+     * {@link ItemInfo} if necessary.
+     */
+    private void addLaunchCookie(IBinder cookie, ItemInfo info, ActivityOptions options) {
+        if (cookie == null) {
+            cookie = mLauncher.getLaunchCookie(info);
+        }
+
+        if (cookie != null) {
+            options.setLaunchCookie(cookie);
+        }
+    }
+
+    /**
      * Whether the launch is a recents app transition and we should do a launch animation
      * from the recents view. Note that if the remote animation targets are not provided, this
      * may not always be correct as we may resolve the opening app to a task when the animation
@@ -1728,6 +1808,10 @@
         }
     }
 
+    private static boolean checkReturnAnimationsFlags() {
+        return enableContainerReturnAnimations() && returnAnimationFrameworkLibrary();
+    }
+
     /**
      * Remote animation runner for animation from the app to Launcher, including recents.
      */
@@ -1844,38 +1928,45 @@
         /** The delegate runner that handles the actual animation. */
         private final RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> mDelegate;
 
+        @Nullable
+        private final ActivityTransitionAnimator.TransitionCookie mCookie;
+
         private ContainerAnimationRunner(
-                RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> delegate) {
+                RemoteAnimationDelegate<IRemoteAnimationFinishedCallback> delegate,
+                ActivityTransitionAnimator.TransitionCookie cookie) {
             mDelegate = delegate;
+            mCookie = cookie;
         }
 
         @Nullable
-        private static ContainerAnimationRunner from(View v, Launcher launcher,
-                StartingWindowListener startingWindowListener, RunnableList onEndCallback) {
-            View viewToUse = findLaunchableViewWithBackground(v);
-            if (viewToUse == null) {
-                return null;
+        ActivityTransitionAnimator.TransitionCookie getCookie() {
+            return mCookie;
+        }
+
+        @Nullable
+        static ContainerAnimationRunner fromView(
+                View v,
+                ActivityTransitionAnimator.TransitionCookie cookie,
+                boolean forLaunch,
+                Launcher launcher,
+                StartingWindowListener startingWindowListener,
+                RunnableList onEndCallback) {
+            if (!forLaunch && !checkReturnAnimationsFlags()) {
+                throw new IllegalStateException(
+                        "forLaunch cannot be false when the enableContainerReturnAnimations or "
+                                + "returnAnimationFrameworkLibrary flag is disabled");
             }
 
-            // The CUJ is logged by the click handler, so we don't log it inside the animation
-            // library.
-            ActivityTransitionAnimator.Controller controllerDelegate =
-                    ActivityTransitionAnimator.Controller.fromView(viewToUse, null /* cujType */);
-
-            if (controllerDelegate == null) {
-                return null;
-            }
-
-            // This wrapper allows us to override the default value, telling the controller that the
-            // current window is below the animating window.
+            // First the controller is created. This is used by the runner to animate the
+            // origin/target view.
             ActivityTransitionAnimator.Controller controller =
-                    new DelegateTransitionAnimatorController(controllerDelegate) {
-                        @Override
-                        public boolean isBelowAnimatingWindow() {
-                            return true;
-                        }
-                    };
+                    buildController(v, cookie, forLaunch);
+            if (controller == null) {
+                return null;
+            }
 
+            // The callback is used to make sure that we use the right color to fade between view
+            // and the window.
             ActivityTransitionAnimator.Callback callback = task -> {
                 final int backgroundColor =
                         startingWindowListener.mBackgroundColor == Color.TRANSPARENT
@@ -1894,7 +1985,52 @@
 
             return new ContainerAnimationRunner(
                     new ActivityTransitionAnimator.AnimationDelegate(
-                            MAIN_EXECUTOR, controller, callback, listener));
+                            MAIN_EXECUTOR, controller, callback, listener),
+                    cookie);
+        }
+
+        /**
+         * Constructs a {@link ActivityTransitionAnimator.Controller} that can be used by a
+         * {@link ContainerAnimationRunner} to animate a view into an opening window or from a
+         * closing one.
+         */
+        @Nullable
+        private static ActivityTransitionAnimator.Controller buildController(
+                View v, ActivityTransitionAnimator.TransitionCookie cookie, boolean isLaunching) {
+            View viewToUse = findLaunchableViewWithBackground(v);
+            if (viewToUse == null) {
+                return null;
+            }
+
+            // The CUJ is logged by the click handler, so we don't log it inside the animation
+            // library. TODO: figure out return CUJ.
+            ActivityTransitionAnimator.Controller controllerDelegate =
+                    ActivityTransitionAnimator.Controller.fromView(viewToUse, null /* cujType */);
+
+            if (controllerDelegate == null) {
+                return null;
+            }
+
+            // This wrapper allows us to override the default value, telling the controller that the
+            // current window is below the animating window as well as information about the return
+            // animation.
+            return new DelegateTransitionAnimatorController(controllerDelegate) {
+                @Override
+                public boolean isLaunching() {
+                    return isLaunching;
+                }
+
+                @Override
+                public boolean isBelowAnimatingWindow() {
+                    return true;
+                }
+
+                @Nullable
+                @Override
+                public ActivityTransitionAnimator.TransitionCookie getTransitionCookie() {
+                    return cookie;
+                }
+            };
         }
 
         /**
@@ -1916,6 +2052,67 @@
             return (T) current;
         }
 
+        /**
+         * Builds the filter used by WM Shell to match app closing transitions (only back, no home
+         * button/gesture) to the given launch cookie.
+         */
+        static TransitionFilter buildBackToHomeFilter(
+                ActivityTransitionAnimator.TransitionCookie cookie, Launcher launcher) {
+            // Closing activity must include the cookie in its list of launch cookies.
+            TransitionFilter.Requirement appRequirement = new TransitionFilter.Requirement();
+            appRequirement.mActivityType = ACTIVITY_TYPE_STANDARD;
+            appRequirement.mLaunchCookie = cookie;
+            appRequirement.mModes = new int[]{TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            // Opening activity must be Launcher.
+            TransitionFilter.Requirement launcherRequirement = new TransitionFilter.Requirement();
+            launcherRequirement.mActivityType = ACTIVITY_TYPE_HOME;
+            launcherRequirement.mModes = new int[]{TRANSIT_OPEN, TRANSIT_TO_FRONT};
+            launcherRequirement.mTopActivity = launcher.getComponentName();
+            // Transition types CLOSE and TO_BACK match the back button/gesture but not the  home
+            // button/gesture.
+            TransitionFilter filter = new TransitionFilter();
+            filter.mTypeSet = new int[]{TRANSIT_CLOSE, TRANSIT_TO_BACK};
+            filter.mRequirements =
+                    new TransitionFilter.Requirement[]{appRequirement, launcherRequirement};
+            return filter;
+        }
+
+        /**
+         * Creates various conditions to ensure that the given transition is cleaned up correctly
+         * when necessary:
+         * - if the transition has run, it is the callback that unregisters it;
+         * - if the associated view is detached before the transition has had an opportunity to run,
+         *   a {@link View.OnAttachStateChangeListener} allows us to do the same (and removes
+         *   itself).
+         */
+        static void setUpRemoteAnimationCleanup(
+                View v, RemoteTransition transition, RunnableList callback, Launcher launcher) {
+            View.OnAttachStateChangeListener listener = new View.OnAttachStateChangeListener() {
+                @Override
+                public void onViewAttachedToWindow(@NonNull View v) {}
+
+                @Override
+                public void onViewDetachedFromWindow(@NonNull View v) {
+                    SystemUiProxy.INSTANCE.get(launcher)
+                            .unregisterRemoteTransition(transition);
+                    v.removeOnAttachStateChangeListener(this);
+                }
+            };
+
+            // Remove the animation as soon as it has run once.
+            callback.add(() -> {
+                SystemUiProxy.INSTANCE.get(launcher).unregisterRemoteTransition(transition);
+                if (v != null) {
+                    v.removeOnAttachStateChangeListener(listener);
+                }
+            });
+
+            // Remove the animation when the view is detached from the hierarchy.
+            // This is so that if back is not invoked (e.g. if we go back home through the home
+            // gesture) we don't have obsolete transitions staying registered.
+            v.addOnAttachStateChangeListener(listener);
+        }
+
         @Override
         public void onAnimationStart(int transit, RemoteAnimationTarget[] appTargets,
                 RemoteAnimationTarget[] wallpaperTargets, RemoteAnimationTarget[] nonAppTargets,
@@ -1928,6 +2125,11 @@
         public void onAnimationCancelled() {
             mDelegate.onAnimationCancelled();
         }
+
+        @Override
+        public boolean supportsReturnTransition() {
+            return true;
+        }
     }
 
     /**
diff --git a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
index 2168f7a..6c7f052 100644
--- a/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/uioverrides/QuickstepLauncher.java
@@ -1182,7 +1182,6 @@
                         : Display.DEFAULT_DISPLAY);
         activityOptions.options.setPendingIntentBackgroundActivityStartMode(
                 ActivityOptions.MODE_BACKGROUND_ACTIVITY_START_ALLOWED);
-        addLaunchCookie(item, activityOptions.options);
         return activityOptions;
     }
 
@@ -1207,19 +1206,6 @@
     }
 
     /**
-     * Adds a new launch cookie for the activity launch if supported.
-     *
-     * @param info the item info for the launch
-     * @param opts the options to set the launchCookie on.
-     */
-    public void addLaunchCookie(ItemInfo info, ActivityOptions opts) {
-        IBinder launchCookie = getLaunchCookie(info);
-        if (launchCookie != null) {
-            opts.setLaunchCookie(launchCookie);
-        }
-    }
-
-    /**
      * Return a new launch cookie for the activity launch if supported.
      *
      * @param info the item info for the launch