Properly animate canceled deep shortcut drag back to original icon

- Added getFirstMatch() instead of using mapOverItems() (was a bit cleaner using ItemInfoMatcher)
- Match based on package name / UserHandle for deep shortcuts case

Test: drag deep shortcut from taskbar icon, inside folder, inside all apps; drag regular icons as well
Fixes: 222574524
Change-Id: Id5fdee29110f143c1125edc6945af09ab0a8d8ce
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
index 24c5d0e..af98b7f 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarDragController.java
@@ -16,6 +16,7 @@
 package com.android.launcher3.taskbar;
 
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_ALL_APPS;
+import static com.android.launcher3.LauncherSettings.Favorites.ITEM_TYPE_DEEP_SHORTCUT;
 
 import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
@@ -55,20 +56,20 @@
 import com.android.launcher3.dragndrop.DragOptions;
 import com.android.launcher3.dragndrop.DragView;
 import com.android.launcher3.dragndrop.DraggableView;
-import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.graphics.DragPreviewProvider;
 import com.android.launcher3.logging.StatsLogManager;
-import com.android.launcher3.model.data.FolderInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.popup.PopupContainerWithArrow;
 import com.android.launcher3.shortcuts.DeepShortcutView;
 import com.android.launcher3.shortcuts.ShortcutDragPreviewProvider;
-import com.android.launcher3.util.LauncherBindableItemsContainer;
+import com.android.launcher3.util.IntSet;
+import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.systemui.shared.recents.model.Task;
 
 import java.io.PrintWriter;
 import java.util.Arrays;
+import java.util.Collections;
 
 /**
  * Handles long click on Taskbar items to start a system drag and drop operation.
@@ -422,23 +423,18 @@
             ItemInfo item = (ItemInfo) tag;
             TaskbarViewController taskbarViewController = mControllers.taskbarViewController;
             if (item.container == CONTAINER_ALL_APPS) {
-                // Since all apps closes when the drag starts, target the all apps button instead
+                // Since all apps closes when the drag starts, target the all apps button instead.
                 target = taskbarViewController.getAllAppsButtonView();
             } else if (item.container >= 0) {
-                // Since folders close when the drag starts, target the folder icon instead
-                LauncherBindableItemsContainer.ItemOperator op = (info, v) -> {
-                    if (info instanceof FolderInfo && v instanceof FolderIcon) {
-                        FolderInfo fi = (FolderInfo) info;
-                        for (WorkspaceItemInfo si : fi.contents) {
-                            if (si.id == item.id) {
-                                // Found the parent
-                                return true;
-                            }
-                        }
-                    }
-                    return false;
-                };
-                target = taskbarViewController.mapOverItems(op);
+                // Since folders close when the drag starts, target the folder icon instead.
+                ItemInfoMatcher matcher = ItemInfoMatcher.forFolderMatch(
+                        ItemInfoMatcher.ofItemIds(IntSet.wrap(item.id)));
+                target = taskbarViewController.getFirstIconMatch(matcher);
+            } else if (item.itemType == ITEM_TYPE_DEEP_SHORTCUT) {
+                // Find first icon with same package/user as the deep shortcut.
+                ItemInfoMatcher packageUserMatcher = ItemInfoMatcher.ofPackages(
+                        Collections.singleton(item.getTargetPackage()), item.user);
+                target = taskbarViewController.getFirstIconMatch(packageUserMatcher);
             }
         }
 
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
index ade58a9..0b537e2 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarView.java
@@ -39,6 +39,7 @@
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.WorkspaceItemInfo;
 import com.android.launcher3.uioverrides.ApiWrapper;
+import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.LauncherBindableItemsContainer;
 import com.android.launcher3.views.ActivityContext;
 import com.android.launcher3.views.AllAppsButton;
@@ -367,17 +368,36 @@
     }
 
     /**
-     * Maps {@code op} over all the child views, returning the view that {@code op} evaluates
-     * {@code true} for, or {@code null} if none satisfy {@code op}.
+     * Maps {@code op} over all the child views.
      */
-    protected View mapOverItems(LauncherBindableItemsContainer.ItemOperator op) {
+    public void mapOverItems(LauncherBindableItemsContainer.ItemOperator op) {
         // map over all the shortcuts on the taskbar
         for (int i = 0; i < getChildCount(); i++) {
             View item = getChildAt(i);
             if (op.evaluate((ItemInfo) item.getTag(), item)) {
-                return item;
+                return;
             }
         }
-        return null;
+    }
+
+    /**
+     * Finds the first icon to match one of the given matchers, from highest to lowest priority.
+     * @return The first match, or All Apps button if no match was found.
+     */
+    public View getFirstMatch(ItemInfoMatcher... matchers) {
+        for (ItemInfoMatcher matcher : matchers) {
+            for (int i = 0; i < getChildCount(); i++) {
+                View item = getChildAt(i);
+                if (!(item.getTag() instanceof ItemInfo)) {
+                    // Should only happen for All Apps button.
+                    continue;
+                }
+                ItemInfo info = (ItemInfo) item.getTag();
+                if (matcher.matchesInfo(info)) {
+                    return item;
+                }
+            }
+        }
+        return mAllAppsButton;
     }
 }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
index 153ed14..6e34ee0 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarViewController.java
@@ -37,6 +37,7 @@
 import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.folder.FolderIcon;
 import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.ItemInfoMatcher;
 import com.android.launcher3.util.LauncherBindableItemsContainer;
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.quickstep.AnimatedFloat;
@@ -274,8 +275,22 @@
         mTaskbarNavButtonTranslationY.updateValue(-deviceProfile.getTaskbarOffsetY());
     }
 
-    public View mapOverItems(LauncherBindableItemsContainer.ItemOperator op) {
-        return mTaskbarView.mapOverItems(op);
+    /**
+     * Maps the given operator to all the top-level children of TaskbarView.
+     */
+    public void mapOverItems(LauncherBindableItemsContainer.ItemOperator op) {
+        mTaskbarView.mapOverItems(op);
+    }
+
+    /**
+     * Returns the first icon to match the given parameter, in priority from:
+     * 1) Icons directly on Taskbar
+     * 2) FolderIcon of the Folder containing the given icon
+     * 3) All Apps button
+     */
+    public View getFirstIconMatch(ItemInfoMatcher matcher) {
+        ItemInfoMatcher folderMatcher = ItemInfoMatcher.forFolderMatch(matcher);
+        return mTaskbarView.getFirstMatch(matcher, folderMatcher);
     }
 
     /**