Make settings cog animate back to original button

Provide the appropriate itemInfo so that launcherBackAnimation would play correctly.

bug:324569690
Test manual video: https://drive.google.com/file/d/1GfjL4PhdGOxs2KfDpieTNiFiKQ61aCgh/view?usp=sharing
Flag: ACONFIG com.android.launcher3.Flags.enable_private_space trunkfood

Change-Id: I1a9458643e932b279bf4ba4ea33e3dc0a2c121c2
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index c91d4d0..dc7c349 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -2508,7 +2508,13 @@
         final int itemCount = container.getChildCount();
         for (int itemIdx = 0; itemIdx < itemCount; itemIdx++) {
             View item = container.getChildAt(itemIdx);
-            if (op.test((ItemInfo) item.getTag())) {
+            if (item instanceof ViewGroup viewGroup) {
+                View view = mapOverViewGroup(viewGroup, op);
+                if (view != null) {
+                    return view;
+                }
+            }
+            if (item.getTag() instanceof ItemInfo itemInfo && op.test(itemInfo)) {
                 return item;
             }
         }
diff --git a/src/com/android/launcher3/LauncherSettings.java b/src/com/android/launcher3/LauncherSettings.java
index 84b8ba1..87ac193 100644
--- a/src/com/android/launcher3/LauncherSettings.java
+++ b/src/com/android/launcher3/LauncherSettings.java
@@ -183,6 +183,7 @@
         public static final int CONTAINER_SHORTCUTS = -107;
         public static final int CONTAINER_SETTINGS = -108;
         public static final int CONTAINER_TASKSWITCHER = -109;
+        public static final int CONTAINER_PRIVATESPACE = -110;
 
         // Represents any of the extended containers implemented in non-AOSP variants.
         public static final int EXTENDED_CONTAINERS = -200;
diff --git a/src/com/android/launcher3/allapps/PrivateProfileManager.java b/src/com/android/launcher3/allapps/PrivateProfileManager.java
index 90ed3eb..f67071b 100644
--- a/src/com/android/launcher3/allapps/PrivateProfileManager.java
+++ b/src/com/android/launcher3/allapps/PrivateProfileManager.java
@@ -21,6 +21,7 @@
 import static android.view.View.VISIBLE;
 
 import static com.android.launcher3.LauncherAnimUtils.VIEW_ALPHA;
+import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_PRIVATESPACE;
 import static com.android.launcher3.allapps.ActivityAllAppsContainerView.AdapterHolder.MAIN;
 import static com.android.launcher3.allapps.BaseAllAppsAdapter.VIEW_TYPE_ICON;
 import static com.android.launcher3.allapps.BaseAllAppsAdapter.VIEW_TYPE_PRIVATE_SPACE_HEADER;
@@ -216,11 +217,27 @@
         resetPrivateSpaceDecorator(updatedState);
     }
 
-    /** Opens the Private Space Settings Page. */
-    public void openPrivateSpaceSettings() {
+    /**
+     * Opens the Private Space Settings Page.
+     *
+     * @param view the view that was clicked to open the settings page and which will be the same
+     *             view to animate back. Otherwise if there is no view, simply start the activity.
+     */
+    public void openPrivateSpaceSettings(View view) {
         if (mPrivateSpaceSettingsAvailable) {
-            mAllApps.getContext().startActivity(
-                    ApiWrapper.INSTANCE.get(mAllApps.getContext()).getPrivateSpaceSettingsIntent());
+            Context context = mAllApps.getContext();
+            Intent intent = ApiWrapper.INSTANCE.get(context).getPrivateSpaceSettingsIntent();
+            if (view == null) {
+                context.startActivity(intent);
+                return;
+            }
+            ActivityContext activityContext = ActivityContext.lookupContext(context);
+            AppInfo itemInfo = new AppInfo();
+            itemInfo.id = CONTAINER_PRIVATESPACE;
+            itemInfo.componentName = intent.getComponent();
+            itemInfo.container = CONTAINER_PRIVATESPACE;
+            view.setTag(itemInfo);
+            activityContext.startActivitySafely(view, intent, itemInfo);
         }
     }
 
@@ -423,7 +440,7 @@
             settingsButton.setOnClickListener(
                     view -> {
                         logEvents(LAUNCHER_PRIVATE_SPACE_SETTINGS_TAP);
-                        openPrivateSpaceSettings();
+                        openPrivateSpaceSettings(view);
                     });
         } else {
             settingsButton.setVisibility(GONE);
diff --git a/src/com/android/launcher3/views/FloatingIconView.java b/src/com/android/launcher3/views/FloatingIconView.java
index f76b53b..f560311 100644
--- a/src/com/android/launcher3/views/FloatingIconView.java
+++ b/src/com/android/launcher3/views/FloatingIconView.java
@@ -55,6 +55,7 @@
 import com.android.launcher3.graphics.PreloadIconDrawable;
 import com.android.launcher3.icons.FastBitmapDrawable;
 import com.android.launcher3.icons.LauncherIcons;
+import com.android.launcher3.model.data.AppInfo;
 import com.android.launcher3.model.data.ItemInfo;
 import com.android.launcher3.model.data.ItemInfoWithIcon;
 import com.android.launcher3.popup.SystemShortcut;
@@ -282,6 +283,8 @@
         } else if (btvIcon instanceof PreloadIconDrawable) {
             // Force the progress bar to display.
             drawable = btvIcon;
+        } else if (originalView instanceof ImageView) {
+            drawable = ((ImageView) originalView).getDrawable();
         } else {
             int width = (int) pos.width();
             int height = (int) pos.height();