Add support for swiping back to the shortcut that launched the activity

Bug: 129067201
Test: Open a shortcut on the workspace, go home

Change-Id: If5d3c3e8e93f09af50aa4994094657347890ef45
Signed-off-by: Winson Chung <winsonc@google.com>
diff --git a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
index fb9765c..caf52e3 100644
--- a/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
+++ b/quickstep/src/com/android/launcher3/BaseQuickstepLauncher.java
@@ -26,10 +26,13 @@
 
 import android.animation.AnimatorSet;
 import android.animation.ValueAnimator;
+import android.app.ActivityOptions;
 import android.content.Intent;
 import android.content.IntentSender;
+import android.os.Binder;
 import android.os.Bundle;
 import android.os.CancellationSignal;
+import android.os.IBinder;
 import android.view.View;
 
 import androidx.annotation.Nullable;
@@ -37,6 +40,8 @@
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.model.WellbeingModel;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.popup.SystemShortcut;
 import com.android.launcher3.proxy.ProxyActivityStarter;
 import com.android.launcher3.proxy.StartActivityParams;
@@ -50,6 +55,7 @@
 import com.android.launcher3.uioverrides.RecentsViewStateController;
 import com.android.launcher3.util.ActivityOptionsWrapper;
 import com.android.launcher3.util.DisplayController;
+import com.android.launcher3.util.ObjectWrapper;
 import com.android.launcher3.util.UiThreadHelper;
 import com.android.quickstep.RecentsModel;
 import com.android.quickstep.SysUINavigationMode;
@@ -67,6 +73,7 @@
 import com.android.systemui.shared.system.ActivityOptionsCompat;
 import com.android.systemui.shared.system.RemoteAnimationTargetCompat;
 
+import java.util.HashMap;
 import java.util.List;
 import java.util.stream.Stream;
 
@@ -418,18 +425,38 @@
     }
 
     @Override
-    public ActivityOptionsWrapper getActivityLaunchOptions(View v) {
+    public ActivityOptionsWrapper getActivityLaunchOptions(View v, @Nullable ItemInfo item) {
         ActivityOptionsWrapper activityOptions =
                 mAppTransitionManager.hasControlRemoteAppTransitionPermission()
                         ? mAppTransitionManager.getActivityLaunchOptions(this, v)
-                        : super.getActivityLaunchOptions(v);
+                        : super.getActivityLaunchOptions(v, item);
         if (mLastTouchUpTime > 0) {
             ActivityOptionsCompat.setLauncherSourceInfo(
                     activityOptions.options, mLastTouchUpTime);
         }
+        addLaunchCookie(item, activityOptions.options);
         return activityOptions;
     }
 
+    /**
+     * Adds a new launch cookie for the activity launch of the given {@param info} if supported.
+     */
+    public void addLaunchCookie(ItemInfo info, ActivityOptions opts) {
+        if (info == null) {
+            return;
+        }
+        switch (info.itemType) {
+            case LauncherSettings.Favorites.ITEM_TYPE_APPLICATION:
+            case LauncherSettings.Favorites.ITEM_TYPE_SHORTCUT:
+            case LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT:
+                // Fall through and continue if it's an app or shortcut
+                break;
+            default:
+                return;
+        }
+        opts.setLaunchCookie(ObjectWrapper.wrap(new Integer(info.id)));
+    }
+
     public void setHintUserWillBeActive() {
         addActivityFlags(ACTIVITY_STATE_USER_WILL_BE_ACTIVE);
     }
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 1b59c49..80b9cf3 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -66,6 +66,7 @@
 import android.graphics.Rect;
 import android.os.Build;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.SystemClock;
 import android.view.MotionEvent;
 import android.view.View;
@@ -1043,7 +1044,8 @@
                 interpolator, target, velocityPxPerMs));
     }
 
-    protected abstract HomeAnimationFactory createHomeAnimationFactory(long duration);
+    protected abstract HomeAnimationFactory createHomeAnimationFactory(
+            ArrayList<IBinder> launchCookies, long duration);
 
     private final TaskStackChangeListener mActivityRestartListener = new TaskStackChangeListener() {
         @Override
@@ -1084,7 +1086,8 @@
             final RemoteAnimationTargetCompat runningTaskTarget = mRecentsAnimationTargets != null
                     ? mRecentsAnimationTargets.findTask(mGestureState.getRunningTaskId())
                     : null;
-            HomeAnimationFactory homeAnimFactory = createHomeAnimationFactory(duration);
+            HomeAnimationFactory homeAnimFactory = createHomeAnimationFactory(
+                    runningTaskTarget.taskInfo.launchCookies, duration);
             mIsSwipingPipToHome = homeAnimFactory.supportSwipePipToHome()
                     && runningTaskTarget != null
                     && runningTaskTarget.taskInfo.pictureInPictureParams != null
diff --git a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
index f5698f7..9846ee7 100644
--- a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
@@ -38,6 +38,7 @@
 import android.os.Build;
 import android.os.Bundle;
 import android.os.Handler;
+import android.os.IBinder;
 import android.os.Looper;
 import android.os.Message;
 import android.os.Messenger;
@@ -68,6 +69,7 @@
 import com.android.systemui.shared.system.SyncRtSurfaceTransactionApplierCompat.SurfaceParams;
 
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
 import java.util.UUID;
 import java.util.function.Consumer;
 
@@ -126,7 +128,8 @@
     }
 
     @Override
-    protected HomeAnimationFactory createHomeAnimationFactory(long duration) {
+    protected HomeAnimationFactory createHomeAnimationFactory(ArrayList<IBinder> launchCookies,
+            long duration) {
         mActiveAnimationFactory = new FallbackHomeAnimationFactory(duration);
         ActivityOptions options = ActivityOptions.makeCustomAnimation(mContext, 0, 0);
         Intent intent = new Intent(mGestureState.getHomeIntent());
diff --git a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
index 5ecd385..dd35d68 100644
--- a/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
+++ b/quickstep/src/com/android/quickstep/LauncherSwipeHandlerV2.java
@@ -31,6 +31,7 @@
 import android.animation.ValueAnimator;
 import android.content.Context;
 import android.graphics.RectF;
+import android.os.IBinder;
 import android.os.UserHandle;
 import android.view.View;
 
@@ -47,6 +48,7 @@
 import com.android.launcher3.dragndrop.DragLayer;
 import com.android.launcher3.states.StateAnimationConfig;
 import com.android.launcher3.util.DynamicResource;
+import com.android.launcher3.util.ObjectWrapper;
 import com.android.launcher3.views.FloatingIconView;
 import com.android.quickstep.util.AppCloseConfig;
 import com.android.quickstep.util.RectFSpringAnim;
@@ -56,6 +58,8 @@
 import com.android.systemui.plugins.ResourceProvider;
 import com.android.systemui.shared.system.InputConsumerController;
 
+import java.util.ArrayList;
+
 /**
  * Temporary class to allow easier refactoring
  */
@@ -71,20 +75,12 @@
 
 
     @Override
-    protected HomeAnimationFactory createHomeAnimationFactory(long duration) {
+    protected HomeAnimationFactory createHomeAnimationFactory(ArrayList<IBinder> launchCookies,
+            long duration) {
         HomeAnimationFactory homeAnimFactory;
         if (mActivity != null) {
-            final TaskView runningTaskView = mRecentsView.getRunningTaskView();
-            final View workspaceView;
-            if (runningTaskView != null
-                    && !mIsSwipingPipToHome
-                    && runningTaskView.getTask().key.getComponent() != null) {
-                workspaceView = mActivity.getWorkspace().getFirstMatchForAppClose(
-                        runningTaskView.getTask().key.getComponent().getPackageName(),
-                        UserHandle.of(runningTaskView.getTask().key.userId));
-            } else {
-                workspaceView = null;
-            }
+            final View workspaceView = findWorkspaceView(launchCookies,
+                    mRecentsView.getRunningTaskView());
             boolean canUseWorkspaceView =
                     workspaceView != null && workspaceView.isAttachedToWindow();
 
@@ -232,6 +228,37 @@
         return homeAnimFactory;
     }
 
+    /**
+     * Returns the associated view on the workspace matching one of the launch cookies, or the app
+     * associated with the running task.
+     */
+    @Nullable
+    private View findWorkspaceView(ArrayList<IBinder> launchCookies, TaskView runningTaskView) {
+        if (mIsSwipingPipToHome) {
+            // Disable if swiping to PIP
+            return null;
+        }
+        if (runningTaskView == null || runningTaskView.getTask() == null
+                || runningTaskView.getTask().key.getComponent() == null) {
+            // Disable if it's an invalid task
+            return null;
+        }
+
+        // Find the associated item info for the launch cookie (if available)
+        int launchCookieItemId = -1;
+        for (IBinder cookie : launchCookies) {
+            Integer itemId = ObjectWrapper.unwrap(cookie);
+            if (itemId != null) {
+                launchCookieItemId = itemId;
+                break;
+            }
+        }
+
+        return mActivity.getWorkspace().getFirstMatchForAppClose(launchCookieItemId,
+                runningTaskView.getTask().key.getComponent().getPackageName(),
+                UserHandle.of(runningTaskView.getTask().key.userId));
+    }
+
     @Override
     protected void finishRecentsControllerToHome(Runnable callback) {
         mRecentsAnimationController.finish(
diff --git a/quickstep/src/com/android/quickstep/RecentsActivity.java b/quickstep/src/com/android/quickstep/RecentsActivity.java
index 7c453e7..ec224ed 100644
--- a/quickstep/src/com/android/quickstep/RecentsActivity.java
+++ b/quickstep/src/com/android/quickstep/RecentsActivity.java
@@ -37,6 +37,8 @@
 import android.os.Looper;
 import android.view.View;
 
+import androidx.annotation.Nullable;
+
 import com.android.launcher3.AbstractFloatingView;
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.InvariantDeviceProfile;
@@ -48,6 +50,7 @@
 import com.android.launcher3.anim.Interpolators;
 import com.android.launcher3.anim.PendingAnimation;
 import com.android.launcher3.compat.AccessibilityManagerCompat;
+import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.statemanager.StateManager;
 import com.android.launcher3.statemanager.StateManager.AtomicAnimationFactory;
 import com.android.launcher3.statemanager.StateManager.StateHandler;
@@ -177,9 +180,9 @@
     }
 
     @Override
-    public ActivityOptionsWrapper getActivityLaunchOptions(final View v) {
+    public ActivityOptionsWrapper getActivityLaunchOptions(final View v, @Nullable ItemInfo item) {
         if (!(v instanceof TaskView)) {
-            return super.getActivityLaunchOptions(v);
+            return super.getActivityLaunchOptions(v, item);
         }
 
         final TaskView taskView = (TaskView) v;
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 2bb14d2..35acdd1 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -523,7 +523,7 @@
         if (mTask != null) {
             TestLogging.recordEvent(
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
-            ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this);
+            ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this, null);
             if (ActivityManagerWrapper.getInstance()
                     .startActivityFromRecents(mTask.key, opts.options)) {
                 return opts.onEndCallback;
diff --git a/src/com/android/launcher3/BaseDraggingActivity.java b/src/com/android/launcher3/BaseDraggingActivity.java
index e9412d9..cc9f594 100644
--- a/src/com/android/launcher3/BaseDraggingActivity.java
+++ b/src/com/android/launcher3/BaseDraggingActivity.java
@@ -165,7 +165,7 @@
     }
 
     @NonNull
-    public ActivityOptionsWrapper getActivityLaunchOptions(View v) {
+    public ActivityOptionsWrapper getActivityLaunchOptions(View v, @Nullable ItemInfo item) {
         int left = 0, top = 0;
         int width = v.getMeasuredWidth(), height = v.getMeasuredHeight();
         if (v instanceof BubbleTextView) {
@@ -192,7 +192,7 @@
             return false;
         }
 
-        Bundle optsBundle = (v != null) ? getActivityLaunchOptions(v).toBundle() : null;
+        Bundle optsBundle = (v != null) ? getActivityLaunchOptions(v, item).toBundle() : null;
         UserHandle user = item == null ? null : item.user;
 
         // Prepare intent
diff --git a/src/com/android/launcher3/Workspace.java b/src/com/android/launcher3/Workspace.java
index 5d941a7..5e37c7a 100644
--- a/src/com/android/launcher3/Workspace.java
+++ b/src/com/android/launcher3/Workspace.java
@@ -2999,20 +2999,30 @@
      * Similar to {@link #getFirstMatch} but optimized to finding a suitable view for the app close
      * animation.
      *
+     * @param preferredItemId The id of the preferred item to match to if it exists.
      * @param packageName The package name of the app to match.
      * @param user The user of the app to match.
      */
-    public View getFirstMatchForAppClose(String packageName, UserHandle user) {
-        List<CellLayout> cellLayouts = new ArrayList<>(getPanelCount() + 1);
-        cellLayouts.add(getHotseat());
-        getVisiblePages().forEach(page -> cellLayouts.add((CellLayout) page));
-
-        final Workspace.ItemOperator packageAndUser = (ItemInfo info, View view) -> info != null
-                && info.getTargetComponent() != null
-                && TextUtils.equals(info.getTargetComponent().getPackageName(), packageName)
-                && info.user.equals(user);
+    public View getFirstMatchForAppClose(int preferredItemId, String packageName, UserHandle user) {
+        final Workspace.ItemOperator preferredItem = (ItemInfo info, View view) ->
+                info != null && info.id == preferredItemId;
+        final Workspace.ItemOperator preferredItemInFolder = (info, view) -> {
+            if (info instanceof FolderInfo) {
+                FolderInfo folderInfo = (FolderInfo) info;
+                for (WorkspaceItemInfo shortcutInfo : folderInfo.contents) {
+                    if (preferredItem.evaluate(shortcutInfo, view)) {
+                        return true;
+                    }
+                }
+            }
+            return false;
+        };
         final Workspace.ItemOperator packageAndUserAndApp = (ItemInfo info, View view) ->
-                packageAndUser.evaluate(info, view) && info.itemType == ITEM_TYPE_APPLICATION;
+                info != null
+                        && info.getTargetComponent() != null
+                        && TextUtils.equals(info.getTargetComponent().getPackageName(), packageName)
+                        && info.user.equals(user)
+                        && info.itemType == ITEM_TYPE_APPLICATION;
         final Workspace.ItemOperator packageAndUserAndAppInFolder = (info, view) -> {
             if (info instanceof FolderInfo) {
                 FolderInfo folderInfo = (FolderInfo) info;
@@ -3025,13 +3035,18 @@
             return false;
         };
 
+        List<CellLayout> cellLayouts = new ArrayList<>(getPanelCount() + 1);
+        cellLayouts.add(getHotseat());
+        getVisiblePages().forEach(page -> cellLayouts.add((CellLayout) page));
+
         // Order: App icons, app in folder. Items in hotseat get returned first.
         if (ADAPTIVE_ICON_WINDOW_ANIM.get()) {
-            return getFirstMatch(cellLayouts, packageAndUserAndApp, packageAndUserAndAppInFolder);
+            return getFirstMatch(cellLayouts, preferredItem, preferredItemInFolder,
+                    packageAndUserAndApp, packageAndUserAndAppInFolder);
         } else {
             // Do not use Folder as a criteria, since it'll cause a crash when trying to draw
             // FolderAdaptiveIcon as the background.
-            return getFirstMatch(cellLayouts, packageAndUserAndApp);
+            return getFirstMatch(cellLayouts, preferredItem, packageAndUserAndApp);
         }
     }
 
diff --git a/src/com/android/launcher3/touch/ItemClickHandler.java b/src/com/android/launcher3/touch/ItemClickHandler.java
index 2e54904..ce7dc07 100644
--- a/src/com/android/launcher3/touch/ItemClickHandler.java
+++ b/src/com/android/launcher3/touch/ItemClickHandler.java
@@ -180,7 +180,7 @@
                 LauncherApps launcherApps = launcher.getSystemService(LauncherApps.class);
                 try {
                     launcherApps.startPackageInstallerSessionDetailsActivity(sessionInfo, null,
-                            launcher.getActivityLaunchOptions(v).toBundle());
+                            launcher.getActivityLaunchOptions(v, item).toBundle());
                     return;
                 } catch (Exception e) {
                     Log.e(TAG, "Unable to launch market intent for package=" + packageName, e);
diff --git a/src/com/android/launcher3/util/ActivityTracker.java b/src/com/android/launcher3/util/ActivityTracker.java
index 59266b4..b5b9c2f 100644
--- a/src/com/android/launcher3/util/ActivityTracker.java
+++ b/src/com/android/launcher3/util/ActivityTracker.java
@@ -75,9 +75,8 @@
     private boolean handleIntent(T activity, Intent intent, boolean alreadyOnHome) {
         if (intent != null && intent.getExtras() != null) {
             IBinder stateBinder = intent.getExtras().getBinder(EXTRA_SCHEDULER_CALLBACK);
-            if (stateBinder instanceof ObjectWrapper) {
-                SchedulerCallback<T> handler =
-                        ((ObjectWrapper<SchedulerCallback>) stateBinder).get();
+            SchedulerCallback<T> handler = ObjectWrapper.unwrap(stateBinder);
+            if (handler != null) {
                 if (!handler.init(activity, alreadyOnHome)) {
                     intent.getExtras().remove(EXTRA_SCHEDULER_CALLBACK);
                 }
diff --git a/src/com/android/launcher3/util/ObjectWrapper.java b/src/com/android/launcher3/util/ObjectWrapper.java
index e5b4707..a715821 100644
--- a/src/com/android/launcher3/util/ObjectWrapper.java
+++ b/src/com/android/launcher3/util/ObjectWrapper.java
@@ -42,4 +42,11 @@
     public static IBinder wrap(Object obj) {
         return new ObjectWrapper<>(obj);
     }
+
+    public static <T> T unwrap(IBinder binder) {
+        if (binder instanceof ObjectWrapper) {
+            return ((ObjectWrapper<T>) binder).get();
+        }
+        return null;
+    }
 }
diff --git a/src/com/android/launcher3/views/FloatingSurfaceView.java b/src/com/android/launcher3/views/FloatingSurfaceView.java
index f13484f..e2e3be7 100644
--- a/src/com/android/launcher3/views/FloatingSurfaceView.java
+++ b/src/com/android/launcher3/views/FloatingSurfaceView.java
@@ -158,7 +158,7 @@
         if (mContract == null) {
             return;
         }
-        View icon = mLauncher.getWorkspace().getFirstMatchForAppClose(
+        View icon = mLauncher.getWorkspace().getFirstMatchForAppClose(-1,
                 mContract.componentName.getPackageName(), mContract.user);
 
         boolean iconChanged = mIcon != icon;