diff --git a/Android.bp b/Android.bp
index ac7e136..35731b5 100644
--- a/Android.bp
+++ b/Android.bp
@@ -280,16 +280,15 @@
     srcs: [
         ":launcher-src-no-build-config",
     ],
-    resource_dirs: [
-        "quickstep/res",
-    ],
+    resource_dirs: [],
     libs: [
         "framework-statsd",
     ],
     static_libs: [
+        "QuickstepResLib",
         "SystemUI-statsd",
         "SystemUISharedLib",
-        "Launcher3CommonDepsLib"
+        "Launcher3CommonDepsLib",
     ],
     manifest: "quickstep/AndroidManifest.xml",
     platform_apis: true,
diff --git a/quickstep/res/drawable/task_menu_item_bg.xml b/quickstep/res/drawable/task_menu_item_bg.xml
index b6a8b90..16c13eb 100644
--- a/quickstep/res/drawable/task_menu_item_bg.xml
+++ b/quickstep/res/drawable/task_menu_item_bg.xml
@@ -15,7 +15,8 @@
      limitations under the License.
 -->
 
-<shape xmlns:android="http://schemas.android.com/apk/res/android">
-    <solid android:color="?android:attr/colorPrimary"/>
-    <corners android:radius="@dimen/task_menu_item_corner_radius"/>
+<shape xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:androidprv="http://schemas.android.com/apk/prv/res/android">
+    <solid android:color="?androidprv:attr/colorSurface" />
+    <corners android:radius="@dimen/task_menu_item_corner_radius" />
 </shape>
diff --git a/quickstep/res/layout/task_menu_with_arrow.xml b/quickstep/res/layout/task_menu_with_arrow.xml
new file mode 100644
index 0000000..38573fd
--- /dev/null
+++ b/quickstep/res/layout/task_menu_with_arrow.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2021 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<com.android.quickstep.views.TaskMenuViewWithArrow
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:layout_width="wrap_content"
+    android:layout_height="wrap_content"
+    android:animateLayoutChanges="true"
+    android:background="@drawable/task_menu_bg"
+    android:orientation="vertical"
+    android:visibility="invisible">
+
+    <LinearLayout
+        android:id="@+id/menu_option_layout"
+        android:layout_width="match_parent"
+        android:layout_height="wrap_content"
+        android:orientation="vertical"
+        android:showDividers="middle" />
+
+</com.android.quickstep.views.TaskMenuViewWithArrow>
\ No newline at end of file
diff --git a/quickstep/res/layout/task_view_menu_option.xml b/quickstep/res/layout/task_view_menu_option.xml
index 5978b97..8a8fc36 100644
--- a/quickstep/res/layout/task_view_menu_option.xml
+++ b/quickstep/res/layout/task_view_menu_option.xml
@@ -18,7 +18,7 @@
     xmlns:android="http://schemas.android.com/apk/res/android"
     android:layout_width="wrap_content"
     android:layout_height="wrap_content"
-    android:orientation="vertical"
+    android:orientation="horizontal"
     android:paddingTop="@dimen/task_card_menu_option_vertical_padding"
     android:paddingBottom="@dimen/task_card_menu_option_vertical_padding"
     android:background="@drawable/task_menu_item_bg"
diff --git a/quickstep/res/values/dimens.xml b/quickstep/res/values/dimens.xml
index 0b534e3..e48c9e8 100644
--- a/quickstep/res/values/dimens.xml
+++ b/quickstep/res/values/dimens.xml
@@ -27,7 +27,7 @@
     <dimen name="task_menu_corner_radius">22dp</dimen>
     <dimen name="task_menu_item_corner_radius">4dp</dimen>
     <dimen name="task_menu_spacing">2dp</dimen>
-    <dimen name="task_menu_width_grid">200dp</dimen>
+    <dimen name="task_menu_width_grid">234dp</dimen>
     <dimen name="overview_proactive_row_height">48dp</dimen>
     <dimen name="overview_proactive_row_bottom_margin">16dp</dimen>
 
@@ -91,7 +91,7 @@
     <dimen name="task_menu_vertical_padding">8dp</dimen>
     <dimen name="task_card_margin">8dp</dimen>
     <dimen name="task_card_menu_shadow_height">3dp</dimen>
-    <dimen name="task_menu_option_start_margin">12dp</dimen>
+    <dimen name="task_menu_option_start_margin">16dp</dimen>
     <!-- Copied from framework resource:
        docked_stack_divider_thickness - 2 * docked_stack_divider_insets -->
     <dimen name="multi_window_task_divider_size">10dp</dimen>
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index a0a940b..45b2081 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -122,6 +122,7 @@
 
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.Arrays;
 import java.util.function.Consumer;
 
 /**
@@ -976,12 +977,13 @@
     }
 
     /** @return Whether this was the task we were waiting to appear, and thus handled it. */
-    protected boolean handleTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {
+    protected boolean handleTaskAppeared(RemoteAnimationTargetCompat[] appearedTaskTarget) {
         if (mStateCallback.hasStates(STATE_HANDLER_INVALIDATED)) {
             return false;
         }
-        if (mStateCallback.hasStates(STATE_START_NEW_TASK)
-                && appearedTaskTarget.taskId == mGestureState.getLastStartedTaskId()) {
+        boolean hasStartedTaskBefore = Arrays.stream(appearedTaskTarget).anyMatch(
+                targetCompat -> targetCompat.taskId == mGestureState.getLastStartedTaskId());
+        if (mStateCallback.hasStates(STATE_START_NEW_TASK) && hasStartedTaskBefore) {
             reset();
             return true;
         }
@@ -1809,13 +1811,8 @@
                 mGestureState.updateLastStartedTaskId(taskId);
                 boolean hasTaskPreviouslyAppeared = mGestureState.getPreviouslyAppearedTaskIds()
                         .contains(taskId);
-                boolean isOldTaskSplit = LauncherSplitScreenListener.INSTANCE.getNoCreate()
-                        .getRunningSplitTaskIds().length > 0;
                 nextTask.launchTask(success -> {
                     resultCallback.accept(success);
-                    if (isOldTaskSplit) {
-                        SystemUiProxy.INSTANCE.getNoCreate().exitSplitScreen(taskId);
-                    }
                     if (success) {
                         if (hasTaskPreviouslyAppeared) {
                             onRestartPreviouslyAppearedTask();
@@ -1871,9 +1868,9 @@
     }
 
     @Override
-    public void onTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {
+    public void onTasksAppeared(RemoteAnimationTargetCompat[] appearedTaskTargets) {
         if (mRecentsAnimationController != null) {
-            if (handleTaskAppeared(appearedTaskTarget)) {
+            if (handleTaskAppeared(appearedTaskTargets)) {
                 mRecentsAnimationController.finish(false /* toRecents */,
                         null /* onFinishComplete */);
                 mActivityInterface.onLaunchTaskSuccess();
diff --git a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
index fed5ae5..a82137e 100644
--- a/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
+++ b/quickstep/src/com/android/quickstep/FallbackSwipeHandler.java
@@ -147,7 +147,7 @@
     }
 
     @Override
-    protected boolean handleTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {
+    protected boolean handleTaskAppeared(RemoteAnimationTargetCompat[] appearedTaskTarget) {
         if (mActiveAnimationFactory != null
                 && mActiveAnimationFactory.handleHomeTaskAppeared(appearedTaskTarget)) {
             mActiveAnimationFactory = null;
@@ -260,7 +260,8 @@
             }
         }
 
-        public boolean handleHomeTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {
+        public boolean handleHomeTaskAppeared(RemoteAnimationTargetCompat[] appearedTaskTargets) {
+            RemoteAnimationTargetCompat appearedTaskTarget = appearedTaskTargets[0];
             if (appearedTaskTarget.activityType == ACTIVITY_TYPE_HOME) {
                 RemoteAnimationTargets targets = new RemoteAnimationTargets(
                         new RemoteAnimationTargetCompat[] {appearedTaskTarget},
diff --git a/quickstep/src/com/android/quickstep/KtR.java b/quickstep/src/com/android/quickstep/KtR.java
new file mode 100644
index 0000000..57dad08
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/KtR.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep;
+
+import com.android.launcher3.R;
+
+/**
+ * Bridge class to allow using resources in Kotlin.
+ * <br/>
+ * TODO(b/204069723) Can't use resources directly in Kotlin
+ */
+public class KtR {
+    public static final class id {
+        public static int menu_option_layout = R.id.menu_option_layout;
+    }
+
+    public static final class dimen {
+        public static int task_menu_spacing = R.dimen.task_menu_spacing;
+    }
+
+    public static final class layout {
+        public static int task_menu_with_arrow = R.layout.task_menu_with_arrow;
+        public static int task_view_menu_option = R.layout.task_view_menu_option;
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/RecentTasksList.java b/quickstep/src/com/android/quickstep/RecentTasksList.java
index 79dd633..c5f4a53 100644
--- a/quickstep/src/com/android/quickstep/RecentTasksList.java
+++ b/quickstep/src/com/android/quickstep/RecentTasksList.java
@@ -27,12 +27,14 @@
 
 import androidx.annotation.VisibleForTesting;
 
+import com.android.quickstep.util.GroupTask;
 import com.android.launcher3.util.LooperExecutor;
-import com.android.systemui.shared.recents.model.GroupTask;
+import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.system.KeyguardManagerCompat;
 import com.android.wm.shell.recents.IRecentTasksListener;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
+import com.android.wm.shell.util.StagedSplitBounds;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -192,12 +194,23 @@
                                 tmpLockedUsers.get(task2Key.userId) /* isLocked */);
                 task2.setLastSnapshotData(taskInfo2);
             }
-            allTasks.add(new GroupTask(task1, task2));
+            final SplitConfigurationOptions.StagedSplitBounds launcherSplitBounds =
+                    convertSplitBounds(rawTask.mStagedSplitBounds);
+            allTasks.add(new GroupTask(task1, task2, launcherSplitBounds));
         }
 
         return allTasks;
     }
 
+    private SplitConfigurationOptions.StagedSplitBounds convertSplitBounds(
+            StagedSplitBounds shellSplitBounds) {
+        return shellSplitBounds == null ?
+                null :
+                new SplitConfigurationOptions.StagedSplitBounds(
+                        shellSplitBounds.leftTopBounds, shellSplitBounds.rightBottomBounds,
+                        shellSplitBounds.leftTopTaskId, shellSplitBounds.rightBottomTaskId);
+    }
+
     private ArrayList<GroupTask> copyOf(ArrayList<GroupTask> tasks) {
         ArrayList<GroupTask> newTasks = new ArrayList<>();
         for (int i = 0; i < tasks.size(); i++) {
diff --git a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
index 750985a..dd6392c 100644
--- a/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
+++ b/quickstep/src/com/android/quickstep/RecentsAnimationCallbacks.java
@@ -136,10 +136,10 @@
 
     @BinderThread
     @Override
-    public void onTaskAppeared(RemoteAnimationTargetCompat app) {
+    public void onTasksAppeared(RemoteAnimationTargetCompat[] apps) {
         Utilities.postAsyncCallback(MAIN_EXECUTOR.getHandler(), () -> {
             for (RecentsAnimationListener listener : getListeners()) {
-                listener.onTaskAppeared(app);
+                listener.onTasksAppeared(apps);
             }
         });
     }
@@ -177,6 +177,6 @@
         /**
          * Callback made when a task started from the recents is ready for an app transition.
          */
-        default void onTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {}
+        default void onTasksAppeared(RemoteAnimationTargetCompat[] appearedTaskTarget) {}
     }
 }
diff --git a/quickstep/src/com/android/quickstep/RecentsModel.java b/quickstep/src/com/android/quickstep/RecentsModel.java
index 71153c6..ac97dd6 100644
--- a/quickstep/src/com/android/quickstep/RecentsModel.java
+++ b/quickstep/src/com/android/quickstep/RecentsModel.java
@@ -34,7 +34,7 @@
 import com.android.launcher3.icons.IconProvider.IconChangeListener;
 import com.android.launcher3.util.Executors.SimpleThreadFactory;
 import com.android.launcher3.util.MainThreadInitializedObject;
-import com.android.systemui.shared.recents.model.GroupTask;
+import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.ThumbnailData;
 import com.android.systemui.shared.system.ActivityManagerWrapper;
diff --git a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
index 5f2b49d..b031c47 100644
--- a/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
+++ b/quickstep/src/com/android/quickstep/RemoteTargetGluer.java
@@ -20,7 +20,7 @@
 
 import androidx.annotation.Nullable;
 
-import com.android.launcher3.util.SplitConfigurationOptions;
+import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import com.android.quickstep.util.AnimatorControllerWithResistance;
 import com.android.quickstep.util.LauncherSplitScreenListener;
 import com.android.quickstep.util.TaskViewSimulator;
@@ -33,7 +33,7 @@
  */
 public class RemoteTargetGluer {
     private RemoteTargetHandle[] mRemoteTargetHandles;
-    private SplitConfigurationOptions.StagedSplitBounds mStagedSplitBounds;
+    private StagedSplitBounds mStagedSplitBounds;
 
     /**
      * Use this constructor if remote targets are split-screen independent
@@ -85,13 +85,21 @@
 
     /**
      * Similar to {@link #assignTargets(RemoteAnimationTargets)}, except this matches the
-     * apps in targets.apps to that of the split screened tasks. If split screen is active, then
-     * {@link #mRemoteTargetHandles} index 0 will be the left/top task, index one right/bottom
+     * apps in targets.apps to that of the _active_ split screened tasks.
+     * See {@link #assignTargetsForSplitScreen(RemoteAnimationTargets, int[])}
      */
     public RemoteTargetHandle[] assignTargetsForSplitScreen(RemoteAnimationTargets targets) {
         int[] splitIds = LauncherSplitScreenListener.INSTANCE.getNoCreate()
                 .getRunningSplitTaskIds();
+        return assignTargetsForSplitScreen(targets, splitIds);
+    }
 
+    /**
+     * Assigns the provided splitIDs to the {@link #mRemoteTargetHandles}, with index 0 will beint
+     * the left/top task, index 1 right/bottom
+     */
+    public RemoteTargetHandle[] assignTargetsForSplitScreen(RemoteAnimationTargets targets,
+            int[] splitIds) {
         RemoteAnimationTargetCompat primaryTaskTarget;
         RemoteAnimationTargetCompat secondaryTaskTarget;
         if (mRemoteTargetHandles.length == 1) {
@@ -108,9 +116,9 @@
             primaryTaskTarget = targets.findTask(splitIds[0]);
             secondaryTaskTarget = targets.findTask(splitIds[1]);
 
-            mStagedSplitBounds = new SplitConfigurationOptions.StagedSplitBounds(
+            mStagedSplitBounds = new StagedSplitBounds(
                     primaryTaskTarget.screenSpaceBounds,
-                    secondaryTaskTarget.screenSpaceBounds);
+                    secondaryTaskTarget.screenSpaceBounds, splitIds[0], splitIds[1]);
             mRemoteTargetHandles[0].mTransformParams.setTargetSet(
                     createRemoteAnimationTargetsForTarget(primaryTaskTarget, targets));
             mRemoteTargetHandles[0].mTaskViewSimulator.setPreview(primaryTaskTarget,
@@ -135,7 +143,7 @@
         return mRemoteTargetHandles;
     }
 
-    public SplitConfigurationOptions.StagedSplitBounds getStagedSplitBounds() {
+    public StagedSplitBounds getStagedSplitBounds() {
         return mStagedSplitBounds;
     }
 
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.java b/quickstep/src/com/android/quickstep/SystemUiProxy.java
index a875b69..4239739 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.java
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.java
@@ -572,53 +572,6 @@
         mPendingSplitScreenListener = null;
     }
 
-    public void setSideStageVisibility(boolean visible) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.setSideStageVisibility(visible);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call setSideStageVisibility");
-            }
-        }
-    }
-
-    /**
-     * To be called whenever the user exits out of split screen apps (either by launching another
-     * app or by swiping home)
-     * @param topTaskId The taskId of the new app that was launched. System will then move this task
-     *                  to the front of what the user sees while removing all other split stages.
-     *                  If swiping to home (or there is no task to put at the top), can pass in -1.
-     */
-    public void exitSplitScreen(int topTaskId) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.exitSplitScreen(topTaskId);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call exitSplitScreen");
-            }
-        }
-    }
-
-    public void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.exitSplitScreenOnHide(exitSplitScreenOnHide);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call exitSplitScreen");
-            }
-        }
-    }
-
-    public void startTask(int taskId, int stage, int position, Bundle options) {
-        if (mSplitScreen != null) {
-            try {
-                mSplitScreen.startTask(taskId, stage, position, options);
-            } catch (RemoteException e) {
-                Log.w(TAG, "Failed call startTask");
-            }
-        }
-    }
-
     /** Start multiple tasks in split-screen simultaneously. */
     public void startTasks(int mainTaskId, Bundle mainOptions, int sideTaskId, Bundle sideOptions,
             @SplitConfigurationOptions.StagePosition int sidePosition,
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index ed7f0b3..e69330a 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -147,16 +147,16 @@
             }
 
             @Override
-            public void onTaskAppeared(RemoteAnimationTargetCompat appearedTaskTarget) {
+            public void onTasksAppeared(RemoteAnimationTargetCompat[] appearedTaskTargets) {
+                RemoteAnimationTargetCompat appearedTaskTarget = appearedTaskTargets[0];
                 BaseActivityInterface activityInterface = mLastGestureState.getActivityInterface();
                 if (ENABLE_QUICKSTEP_LIVE_TILE.get() && activityInterface.isInLiveTileMode()
                         && activityInterface.getCreatedActivity() != null) {
                     RecentsView recentsView =
                             activityInterface.getCreatedActivity().getOverviewPanel();
                     if (recentsView != null) {
-                        RemoteAnimationTargetCompat[] apps = new RemoteAnimationTargetCompat[1];
-                        apps[0] = appearedTaskTarget;
-                        recentsView.launchSideTaskInLiveTileMode(appearedTaskTarget.taskId, apps,
+                        recentsView.launchSideTaskInLiveTileMode(appearedTaskTarget.taskId,
+                                appearedTaskTargets,
                                 new RemoteAnimationTargetCompat[0] /* wallpaper */,
                                 new RemoteAnimationTargetCompat[0] /* nonApps */);
                         return;
diff --git a/quickstep/src/com/android/quickstep/TaskViewUtils.java b/quickstep/src/com/android/quickstep/TaskViewUtils.java
index 15729e1..67a94cc 100644
--- a/quickstep/src/com/android/quickstep/TaskViewUtils.java
+++ b/quickstep/src/com/android/quickstep/TaskViewUtils.java
@@ -177,8 +177,8 @@
         } else {
             RemoteTargetGluer gluer = new RemoteTargetGluer(v.getContext(),
                     recentsView.getSizeStrategy(), targets);
-            if (recentsViewHandles != null && recentsViewHandles.length > 1) {
-                remoteTargetHandles = gluer.assignTargetsForSplitScreen(targets);
+            if (v.containsMultipleTasks()) {
+                remoteTargetHandles = gluer.assignTargetsForSplitScreen(targets, v.getTaskIds());
             } else {
                 remoteTargetHandles = gluer.assignTargets(targets);
             }
diff --git a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
index 4044f7b..4f4edab 100644
--- a/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
+++ b/quickstep/src/com/android/quickstep/fallback/FallbackRecentsView.java
@@ -42,7 +42,7 @@
 import com.android.quickstep.views.OverviewActionsView;
 import com.android.quickstep.views.RecentsView;
 import com.android.quickstep.views.TaskView;
-import com.android.systemui.shared.recents.model.GroupTask;
+import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 
@@ -176,7 +176,7 @@
                 newList.addAll(taskGroups);
                 newList.add(new GroupTask(
                         Task.from(new TaskKey(mHomeTaskInfo), mHomeTaskInfo, false),
-                        null));
+                        null, null));
                 taskGroups = newList;
             }
         }
diff --git a/quickstep/src/com/android/quickstep/util/GroupTask.java b/quickstep/src/com/android/quickstep/util/GroupTask.java
new file mode 100644
index 0000000..e2563e3
--- /dev/null
+++ b/quickstep/src/com/android/quickstep/util/GroupTask.java
@@ -0,0 +1,56 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.quickstep.util;
+
+import androidx.annotation.NonNull;
+import androidx.annotation.Nullable;
+
+import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
+import com.android.systemui.shared.recents.model.Task;
+
+/**
+ * A {@link Task} container that can contain one or two tasks, depending on if the two tasks
+ * are represented as an app-pair in the recents task list.
+ */
+public class GroupTask {
+    public @NonNull Task task1;
+    public @Nullable Task task2;
+    public @Nullable StagedSplitBounds mStagedSplitBounds;
+
+    public GroupTask(@NonNull Task t1, @Nullable Task t2,
+            @Nullable StagedSplitBounds stagedSplitBounds) {
+        task1 = t1;
+        task2 = t2;
+        mStagedSplitBounds = stagedSplitBounds;
+    }
+
+    public GroupTask(@NonNull GroupTask group) {
+        task1 = new Task(group.task1);
+        task2 = group.task2 != null
+                ? new Task(group.task2)
+                : null;
+        mStagedSplitBounds = group.mStagedSplitBounds;
+    }
+
+    public boolean containsTask(int taskId) {
+        return task1.key.id == taskId || (task2 != null && task2.key.id == taskId);
+    }
+
+    public boolean hasMultipleTasks() {
+        return task2 != null;
+    }
+}
diff --git a/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java b/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
index fa4cddc..99efb39 100644
--- a/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
+++ b/quickstep/src/com/android/quickstep/util/LauncherSplitScreenListener.java
@@ -149,7 +149,6 @@
             return;
         }
 
-        SystemUiProxy.INSTANCE.getNoCreate().exitSplitScreen(-1);
         mPersistentGroupedIds = EMPTY_ARRAY;
     }
 
diff --git a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
index 4c300ec..e1afa97 100644
--- a/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
+++ b/quickstep/src/com/android/quickstep/util/SplitSelectStateController.java
@@ -56,6 +56,7 @@
     private Task mInitialTask;
     private Task mSecondTask;
     private Rect mInitialBounds;
+    private boolean mRecentsAnimationRunning;
 
     public SplitSelectStateController(Handler handler, SystemUiProxy systemUiProxy) {
         mHandler = handler;
@@ -113,6 +114,10 @@
         return mStagePosition;
     }
 
+    public void setRecentsAnimationRunning(boolean running) {
+        this.mRecentsAnimationRunning = running;
+    }
+
     /**
      * Requires Shell Transitions
      */
@@ -172,7 +177,9 @@
         public void onAnimationCancelled() {
             postAsyncCallback(mHandler, () -> {
                 if (mSuccessCallback != null) {
-                    mSuccessCallback.accept(false);
+                    // Launching legacy tasks while recents animation is running will always cause
+                    // onAnimationCancelled to be called (should be fixed w/ shell transitions?)
+                    mSuccessCallback.accept(mRecentsAnimationRunning);
                 }
                 resetState();
             });
@@ -187,6 +194,7 @@
         mSecondTask = null;
         mStagePosition = SplitConfigurationOptions.STAGE_POSITION_UNDEFINED;
         mInitialBounds = null;
+        mRecentsAnimationRunning = false;
     }
 
     /**
diff --git a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
index 977c696..280b5d7 100644
--- a/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
+++ b/quickstep/src/com/android/quickstep/util/TaskViewSimulator.java
@@ -22,7 +22,6 @@
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_TOP_OR_LEFT;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_UNDEFINED;
 import static com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
-import static com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import static com.android.quickstep.util.RecentsOrientedState.postDisplayRotation;
 import static com.android.quickstep.util.RecentsOrientedState.preDisplayRotation;
 import static com.android.systemui.shared.system.WindowManagerWrapper.WINDOWING_MODE_FULLSCREEN;
@@ -41,6 +40,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.Utilities;
 import com.android.launcher3.anim.PendingAnimation;
+import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import com.android.launcher3.util.TraceHelper;
 import com.android.quickstep.AnimatedFloat;
 import com.android.quickstep.BaseActivityInterface;
diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index 357fb67..eace227 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -14,6 +14,7 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.util.RunnableList;
+import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import com.android.launcher3.util.TransformingTouchDelegate;
 import com.android.quickstep.RecentsModel;
@@ -118,7 +119,11 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
-        return TaskMenuView.showForTask(mTaskIdAttributeContainer[iconView == mIconView ? 0 : 1]);
+        if (mActivity.getDeviceProfile().overviewShowAsGrid) {
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+        } else {
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        }
     }
 
     public void updateSplitBoundsConfig(StagedSplitBounds stagedSplitBounds) {
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index 9a76f2b..72e01bb 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -137,8 +137,8 @@
 import com.android.launcher3.util.MultiValueAlpha;
 import com.android.launcher3.util.ResourceBasedOverride.Overrides;
 import com.android.launcher3.util.RunnableList;
-import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.StagePosition;
+import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import com.android.launcher3.util.Themes;
 import com.android.launcher3.util.TranslateEdgeEffect;
 import com.android.launcher3.util.ViewPool;
@@ -157,7 +157,6 @@
 import com.android.quickstep.TaskThumbnailCache;
 import com.android.quickstep.TaskViewUtils;
 import com.android.quickstep.ViewUtils;
-import com.android.quickstep.util.LauncherSplitScreenListener;
 import com.android.quickstep.util.LayoutUtils;
 import com.android.quickstep.util.RecentsOrientedState;
 import com.android.quickstep.util.SplitScreenBounds;
@@ -167,7 +166,7 @@
 import com.android.quickstep.util.TransformParams;
 import com.android.quickstep.util.VibratorWrapper;
 import com.android.systemui.plugins.ResourceProvider;
-import com.android.systemui.shared.recents.model.GroupTask;
+import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.recents.model.Task;
 import com.android.systemui.shared.recents.model.Task.TaskKey;
 import com.android.systemui.shared.recents.model.ThumbnailData;
@@ -599,7 +598,7 @@
      */
     private TaskView mSplitHiddenTaskView;
     private TaskView mSecondSplitHiddenTaskView;
-    private SplitConfigurationOptions.StagedSplitBounds mSplitBoundsConfig;
+    private StagedSplitBounds mSplitBoundsConfig;
     private final Toast mSplitToast = Toast.makeText(getContext(),
             R.string.toast_split_select_app, Toast.LENGTH_SHORT);
     private final Toast mSplitUnsupportedToast = Toast.makeText(getContext(),
@@ -665,9 +664,8 @@
         mClearAllButton.setOnClickListener(this::dismissAllTasks);
         mTaskViewPool = new ViewPool<>(context, this, R.layout.task, 20 /* max size */,
                 10 /* initial size */);
-        // There's only one pair of grouped tasks we can envision at the moment
         mGroupedTaskViewPool = new ViewPool<>(context, this,
-                R.layout.task_grouped, 2 /* max size */, 1 /* initial size */);
+                R.layout.task_grouped, 20 /* max size */, 10 /* initial size */);
 
         mIsRtl = mOrientationHandler.getRecentsRtlSetting(getResources());
         setLayoutDirection(mIsRtl ? View.LAYOUT_DIRECTION_RTL : View.LAYOUT_DIRECTION_LTR);
@@ -1071,6 +1069,10 @@
     }
 
     public TaskView getTaskViewByTaskId(int taskId) {
+        if (taskId == -1) {
+            return null;
+        }
+
         for (int i = 0; i < getTaskViewCount(); i++) {
             TaskView taskView = getTaskViewAt(i);
             int[] taskIds = taskView.getTaskIds();
@@ -1323,72 +1325,38 @@
         TaskView ignoreResetTaskView =
                 mIgnoreResetTaskId == -1 ? null : getTaskViewByTaskId(mIgnoreResetTaskId);
 
-        int[] splitTaskIds =
-                LauncherSplitScreenListener.INSTANCE.getNoCreate().getPersistentSplitIds();
-        int requiredGroupTaskViews = splitTaskIds.length / 2;
-
-        // TODO(b/202740477): Update once grouped tasks are returned
-        // Subtract half the number of split tasks and not total number because we've already
-        // added a GroupedTaskView when swipe up gesture happens.
-        // This will need to change if we start showing GroupedTaskViews during swipe up from home
-        int requiredTaskViewCount = taskGroups.size() - requiredGroupTaskViews;
-
-        if (getTaskViewCount() != requiredTaskViewCount) {
-            if (indexOfChild(mClearAllButton) != -1) {
-                removeView(mClearAllButton);
-            }
-
-            for (int i = getTaskViewCount(); i < requiredTaskViewCount; i++) {
-                addView(getTaskViewFromPool(false));
-            }
-            while (getTaskViewCount() > requiredTaskViewCount) {
-                removeView(getChildAt(getChildCount() - 1));
-            }
-            int groupedTaskViewCount = getGroupedTaskViewCount();
-            while (requiredGroupTaskViews > groupedTaskViewCount) {
-                // Add to front of list
-                addView(getTaskViewFromPool(true), 0);
-                requiredGroupTaskViews--;
-            }
-            if (requiredTaskViewCount > 0) {
-                addView(mClearAllButton);
-            }
-        }
-
         // Save running task ID if it exists before rebinding all taskViews, otherwise the task from
         // the runningTaskView currently bound could get assigned to another TaskView
-        // TODO set these type to array and check all taskIDs? Maybe we can get away w/ only one
         int runningTaskId = getTaskIdsForTaskViewId(mRunningTaskViewId)[0];
         int focusedTaskId = getTaskIdsForTaskViewId(mFocusedTaskViewId)[0];
-        Log.d(TASK_VIEW_ID_CRASH, "runningTaskId beforeBind: " + runningTaskId
-                + " runningTaskViewId: " + mRunningTaskViewId
-                + " forTaskView: " + getTaskViewFromTaskViewId(mRunningTaskViewId));
 
-        for (int taskViewIndex = requiredTaskViewCount - 1, taskDataIndex = taskGroups.size() - 1;
-                taskViewIndex >= 0;
-                taskViewIndex--, taskDataIndex--) {
-            final int pageIndex = requiredTaskViewCount - taskViewIndex - 1;
-            // TODO(b/202740477): Temporary assumption, to be updated once groups are actually used
-            final Task task = taskGroups.get(taskDataIndex).task1;
-            final TaskView taskView = (TaskView) getChildAt(pageIndex);
-            if (taskView instanceof GroupedTaskView) {
-                Task leftTop;
-                Task rightBottom;
-                if (task.key.id == splitTaskIds[0]) {
-                    leftTop = task;
-                    taskDataIndex--;
-                    rightBottom = taskGroups.get(taskDataIndex).task1;
-                } else {
-                    rightBottom = task;
-                    taskDataIndex--;
-                    leftTop = taskGroups.get(taskDataIndex).task1;
-                }
-                ((GroupedTaskView) taskView).bind(leftTop, rightBottom, mOrientationState,
-                        mSplitBoundsConfig);
+        // Removing views sets the currentPage to 0, so we save this and restore it after
+        // the new set of views are added
+        int previousPage = mCurrentPage;
+        removeAllViews();
+
+        // Add views as children based on whether it's grouped or single task
+        for (int i = taskGroups.size() - 1; i >= 0; i--) {
+            GroupTask groupTask = taskGroups.get(i);
+            boolean hasMultipleTasks = groupTask.hasMultipleTasks();
+            TaskView taskView = getTaskViewFromPool(hasMultipleTasks);
+            addView(taskView);
+
+            if (hasMultipleTasks) {
+                boolean firstTaskIsLeftTopTask =
+                        groupTask.mStagedSplitBounds.leftTopTaskId == groupTask.task1.key.id;
+                Task leftTopTask = firstTaskIsLeftTopTask ? groupTask.task1 : groupTask.task2;
+                Task rightBottomTask = firstTaskIsLeftTopTask ? groupTask.task2 : groupTask.task1;
+                ((GroupedTaskView) taskView).bind(leftTopTask, rightBottomTask, mOrientationState,
+                        groupTask.mStagedSplitBounds);
             } else {
-                taskView.bind(task, mOrientationState);
+                taskView.bind(groupTask.task1, mOrientationState);
             }
         }
+        if (!taskGroups.isEmpty()) {
+            addView(mClearAllButton);
+        }
+        setCurrentPage(previousPage);
 
         // Keep same previous focused task
         TaskView newFocusedTaskView = getTaskViewByTaskId(focusedTaskId);
@@ -1406,21 +1374,10 @@
             // Update mRunningTaskViewId to be the new TaskView that was assigned by binding
             // the full list of tasks to taskViews
             newRunningTaskView = getTaskViewByTaskId(runningTaskId);
-            if (newRunningTaskView == null) {
-                StringBuilder sb = new StringBuilder();
-                for (int i = requiredTaskViewCount - 1; i >= 0; i--) {
-                    final int pageIndex = requiredTaskViewCount - i - 1;
-                    final TaskView taskView = (TaskView) getChildAt(pageIndex);
-                    int taskViewId = taskView.getTaskViewId();
-                    sb.append(" taskViewId: " + taskViewId
-                            + " taskId: " + getTaskIdsForTaskViewId(taskViewId)[0]
-                            + " for taskView: " + taskView + "\n");
-                }
-                Log.d(TASK_VIEW_ID_CRASH, "taskViewCount: " + getTaskViewCount()
-                        + " " + sb.toString());
-                mRunningTaskViewId = -1;
-            } else {
+            if (newRunningTaskView != null) {
                 mRunningTaskViewId = newRunningTaskView.getTaskViewId();
+            } else {
+                mRunningTaskViewId = -1;
             }
         }
 
@@ -1921,6 +1878,7 @@
             remoteTargetHandle.getTransformParams().setTargetSet(null);
             remoteTargetHandle.getTaskViewSimulator().setDrawsBelowRecents(true);
         });
+        mSplitSelectStateController.resetState();
 
         // These are relatively expensive and don't need to be done this frame (RecentsView isn't
         // visible anyway), so defer by a frame to get off the critical path, e.g. app to home.
@@ -4380,6 +4338,7 @@
     public void setRecentsAnimationTargets(RecentsAnimationController recentsAnimationController,
             RecentsAnimationTargets recentsAnimationTargets) {
         mRecentsAnimationController = recentsAnimationController;
+        mSplitSelectStateController.setRecentsAnimationRunning(true);
         if (recentsAnimationTargets == null || recentsAnimationTargets.apps.length == 0) {
             return;
         }
@@ -4479,6 +4438,7 @@
         // taps on QSB (3) user goes back to Overview and launch the most recent task.
         setCurrentTask(-1);
         mRecentsAnimationController = null;
+        mSplitSelectStateController.setRecentsAnimationRunning(false);
         executeSideTaskLaunchCallback();
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuView.java b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
index 77ac373..d55b89c 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuView.java
@@ -34,7 +34,6 @@
 import android.view.Gravity;
 import android.view.MotionEvent;
 import android.view.View;
-import android.view.ViewGroup;
 import android.view.ViewOutlineProvider;
 import android.view.ViewTreeObserver.OnScrollChangedListener;
 import android.widget.LinearLayout;
@@ -270,15 +269,9 @@
         BaseDragLayer.LayoutParams params = (BaseDragLayer.LayoutParams) getLayoutParams();
         int padding = getResources()
                 .getDimensionPixelSize(R.dimen.task_menu_vertical_padding);
-        if (deviceProfile.overviewShowAsGrid) {
-            // TODO(b/193432925) temporary so it doesn't look terrible on large screen
-            params.width =
-                    getContext().getResources().getDimensionPixelSize(R.dimen.task_menu_width_grid);
-        } else {
-            params.width = orientationHandler
-                    .getTaskMenuWidth(taskContainer.getThumbnailView(),
-                            deviceProfile) - (2 * padding);
-        }
+        params.width = orientationHandler
+                .getTaskMenuWidth(taskContainer.getThumbnailView(),
+                        deviceProfile) - (2 * padding);
         // Gravity set to Left instead of Start as sTempRect.left measures Left distance not Start
         params.gravity = Gravity.LEFT;
         setLayoutParams(params);
@@ -349,13 +342,4 @@
         return new RoundedRectRevealOutlineProvider(radius, radius, fromRect, toRect);
     }
 
-    public View findMenuItemByText(String text) {
-        for (int i = mOptionLayout.getChildCount() - 1; i >= 0; --i) {
-            final ViewGroup menuOptionView = (ViewGroup) mOptionLayout.getChildAt(i);
-            if (text.equals(menuOptionView.<TextView>findViewById(R.id.text).getText())) {
-                return menuOptionView;
-            }
-        }
-        return null;
-    }
 }
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
index 9b86c73..5059f8b 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -1,14 +1,167 @@
+/*
+ * Copyright (C) 2021 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
 package com.android.quickstep.views
 
-import android.util.Log
+import android.animation.AnimatorSet
+import android.animation.ObjectAnimator
+import android.content.Context
+import android.graphics.Rect
+import android.graphics.drawable.ShapeDrawable
+import android.graphics.drawable.shapes.RectShape
+import android.util.AttributeSet
+import android.view.MotionEvent
+import android.view.View
+import android.view.ViewGroup
+import android.widget.LinearLayout
+import com.android.launcher3.BaseDraggingActivity
+import com.android.launcher3.DeviceProfile
+import com.android.launcher3.R
+import com.android.launcher3.popup.ArrowPopup
+import com.android.launcher3.popup.SystemShortcut
+import com.android.launcher3.util.Themes
+import com.android.quickstep.KtR
+import com.android.quickstep.TaskOverlayFactory
+import com.android.quickstep.views.TaskView.TaskIdAttributeContainer
 
-// TODO(http://b/193432925)
-class TaskMenuViewWithArrow {
+class TaskMenuViewWithArrow<T : BaseDraggingActivity> : ArrowPopup<T> {
     companion object {
         const val TAG = "TaskMenuViewWithArrow"
 
-        fun logSomething() {
-            Log.d(TAG, "It worked!")
+        fun showForTask(taskContainer: TaskIdAttributeContainer): Boolean {
+            val activity = BaseDraggingActivity
+                    .fromContext<BaseDraggingActivity>(taskContainer.taskView.context)
+            val taskMenuViewWithArrow = activity.layoutInflater
+                    .inflate(KtR.layout.task_menu_with_arrow, activity.dragLayer, false) as TaskMenuViewWithArrow<*>
+
+            return taskMenuViewWithArrow.populateAndShowForTask(taskContainer)
         }
     }
+
+    constructor(context: Context) : super(context)
+    constructor(context: Context, attrs: AttributeSet) : super(context, attrs)
+    constructor(context: Context, attrs: AttributeSet, defStyleAttr: Int) : super(context, attrs, defStyleAttr)
+
+    init {
+        clipToOutline = true
+    }
+
+    private val menuWidth = context.resources.getDimensionPixelSize(R.dimen.task_menu_width_grid)
+
+    private lateinit var taskView: TaskView
+    private lateinit var optionLayout: LinearLayout
+    private lateinit var taskContainer: TaskIdAttributeContainer
+
+    override fun isOfType(type: Int): Boolean = type and TYPE_TASK_MENU != 0
+
+    override fun getTargetObjectLocation(outPos: Rect?) {
+        popupContainer.getDescendantRectRelativeToSelf(taskView.iconView, outPos)
+    }
+
+    override fun onControllerInterceptTouchEvent(ev: MotionEvent?): Boolean {
+        if (ev?.action == MotionEvent.ACTION_DOWN) {
+            if (!popupContainer.isEventOverView(this, ev)) {
+                close(true)
+                return true
+            }
+        }
+        return false
+    }
+
+    override fun onFinishInflate() {
+        super.onFinishInflate()
+        optionLayout = findViewById(KtR.id.menu_option_layout)
+    }
+
+    private fun populateAndShowForTask(taskContainer: TaskIdAttributeContainer): Boolean {
+        if (isAttachedToWindow) {
+            return false
+        }
+
+        taskView = taskContainer.taskView
+        this.taskContainer = taskContainer
+        if (!populateMenu()) return false
+        show()
+        return true
+    }
+
+    /** @return true if successfully able to populate task view menu, false otherwise
+     */
+    private fun populateMenu(): Boolean {
+        // Icon may not be loaded
+        if (taskContainer.task.icon == null) return false
+
+        addMenuOptions()
+        return true
+    }
+
+    private fun addMenuOptions() {
+        // Add the options
+        TaskOverlayFactory
+            .getEnabledShortcuts(taskView, mActivityContext.deviceProfile, taskContainer)
+            .forEach { this.addMenuOption(it) }
+
+        // Add the spaces between items
+        val divider = ShapeDrawable(RectShape())
+        divider.paint.color = resources.getColor(android.R.color.transparent)
+        val dividerSpacing = resources.getDimension(KtR.dimen.task_menu_spacing).toInt()
+        optionLayout.showDividers = SHOW_DIVIDER_MIDDLE
+
+        // Set the orientation, which makes the menu show
+        val recentsView: RecentsView<*, *> = mActivityContext.getOverviewPanel()
+        val orientationHandler = recentsView.pagedOrientationHandler
+        val deviceProfile: DeviceProfile = mActivityContext.deviceProfile
+        orientationHandler.setTaskOptionsMenuLayoutOrientation(
+            deviceProfile,
+            optionLayout,
+            dividerSpacing,
+            divider
+        )
+    }
+
+    private fun addMenuOption(menuOption: SystemShortcut<*>) {
+        val menuOptionView = mActivityContext.layoutInflater.inflate(
+            KtR.layout.task_view_menu_option, this, false
+        ) as LinearLayout
+        menuOption.setIconAndLabelFor(
+            menuOptionView.findViewById(R.id.icon),
+            menuOptionView.findViewById(R.id.text)
+        )
+        val lp = menuOptionView.layoutParams as LayoutParams
+        lp.width = menuWidth
+        menuOptionView.setOnClickListener { view: View? -> menuOption.onClick(view) }
+        optionLayout.addView(menuOptionView)
+    }
+
+    override fun assignMarginsAndBackgrounds(viewGroup: ViewGroup) {
+        assignMarginsAndBackgrounds(this, Themes.getAttrColor(context, com.android.internal.R.attr.colorSurface))
+    }
+
+    override fun onCreateOpenAnimation(anim: AnimatorSet) {
+        anim.play(
+            ObjectAnimator.ofFloat(
+                taskContainer.thumbnailView, TaskThumbnailView.DIM_ALPHA,
+                TaskView.MAX_PAGE_SCRIM_ALPHA
+            )
+        )
+    }
+
+    override fun onCreateCloseAnimation(anim: AnimatorSet) {
+        anim.play(
+            ObjectAnimator.ofFloat(taskContainer.thumbnailView, TaskThumbnailView.DIM_ALPHA, 0f)
+        )
+    }
 }
\ No newline at end of file
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index b1c733f..4132006 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -557,6 +557,10 @@
         return mTaskIdContainer;
     }
 
+    public boolean containsMultipleTasks() {
+        return mTaskIdContainer[1] != -1;
+    }
+
     public TaskThumbnailView getThumbnail() {
         return mSnapshotView;
     }
@@ -689,13 +693,8 @@
                     TestProtocol.SEQUENCE_MAIN, "startActivityFromRecentsAsync", mTask);
             ActivityOptionsWrapper opts =  mActivity.getActivityLaunchOptions(this, null);
             opts.options.setLaunchDisplayId(getRootViewDisplayId());
-            boolean isOldTaskSplit = LauncherSplitScreenListener.INSTANCE.getNoCreate()
-                    .getPersistentSplitIds().length > 0;
             if (ActivityManagerWrapper.getInstance()
                     .startActivityFromRecents(mTask.key, opts.options)) {
-                if (isOldTaskSplit) {
-                    SystemUiProxy.INSTANCE.getNoCreate().exitSplitScreen(mTask.key.id);
-                }
                 RecentsView recentsView = getRecentsView();
                 if (ENABLE_QUICKSTEP_LIVE_TILE.get() && recentsView.getRunningTaskViewId() != -1) {
                     recentsView.onTaskLaunchedInLiveTileMode();
@@ -838,9 +837,11 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
-        // TODO(http://b/193432925)
-        if (DEBUG) TaskMenuViewWithArrow.Companion.logSomething();
-        return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        if (mActivity.getDeviceProfile().overviewShowAsGrid) {
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+        } else {
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        }
     }
 
     protected void setIcon(IconView iconView, Drawable icon) {
@@ -850,7 +851,15 @@
                 if (confirmSecondSplitSelectApp()) {
                     return;
                 }
-                showTaskMenu(iconView);
+                if (ENABLE_QUICKSTEP_LIVE_TILE.get() && isRunningTask()) {
+                    RecentsView recentsView = getRecentsView();
+                    recentsView.switchToScreenshot(
+                            () -> recentsView.finishRecentsAnimation(true /* toRecents */,
+                                    false /* shouldPip */,
+                                    () -> showTaskMenu(iconView)));
+                } else {
+                    showTaskMenu(iconView);
+                }
             });
             iconView.setOnLongClickListener(v -> {
                 requestDisallowInterceptTouchEvent(true);
diff --git a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
index 6b2d5ed..4e49716 100644
--- a/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
+++ b/quickstep/tests/src/com/android/quickstep/RecentTasksListTest.java
@@ -30,9 +30,7 @@
 import androidx.test.filters.SmallTest;
 
 import com.android.launcher3.util.LooperExecutor;
-import com.android.systemui.shared.recents.model.GroupTask;
-import com.android.systemui.shared.recents.model.Task;
-import com.android.systemui.shared.system.ActivityManagerWrapper;
+import com.android.quickstep.util.GroupTask;
 import com.android.systemui.shared.system.KeyguardManagerCompat;
 import com.android.wm.shell.util.GroupedRecentTaskInfo;
 
@@ -73,7 +71,7 @@
     @Test
     public void loadTasksInBackground_onlyKeys_noValidTaskDescription() {
         GroupedRecentTaskInfo recentTaskInfos = new GroupedRecentTaskInfo(
-                new ActivityManager.RecentTaskInfo(), new ActivityManager.RecentTaskInfo());
+                new ActivityManager.RecentTaskInfo(), new ActivityManager.RecentTaskInfo(), null);
         when(mockSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
                 .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
@@ -93,7 +91,7 @@
         ActivityManager.RecentTaskInfo task2 = new ActivityManager.RecentTaskInfo();
         task2.taskDescription = new ActivityManager.TaskDescription();
         GroupedRecentTaskInfo recentTaskInfos = new GroupedRecentTaskInfo(
-                task1, task2);
+                task1, task2, null);
         when(mockSystemUiProxy.getRecentTasks(anyInt(), anyInt()))
                 .thenReturn(new ArrayList<>(Collections.singletonList(recentTaskInfos)));
 
diff --git a/src/com/android/launcher3/Launcher.java b/src/com/android/launcher3/Launcher.java
index 8d92bf2..8154168 100644
--- a/src/com/android/launcher3/Launcher.java
+++ b/src/com/android/launcher3/Launcher.java
@@ -154,6 +154,7 @@
 import com.android.launcher3.notification.NotificationListener;
 import com.android.launcher3.pm.PinRequestHelper;
 import com.android.launcher3.pm.UserCache;
+import com.android.launcher3.popup.ArrowPopup;
 import com.android.launcher3.popup.PopupContainerWithArrow;
 import com.android.launcher3.popup.PopupDataProvider;
 import com.android.launcher3.popup.SystemShortcut;
@@ -3032,4 +3033,13 @@
     public StatsLogManager getStatsLogManager() {
         return super.getStatsLogManager().withDefaultInstanceId(mAllAppsSessionLogId);
     }
+
+    /**
+     * Returns the current popup for testing, if any.
+     */
+    @VisibleForTesting
+    @Nullable
+    public ArrowPopup<?> getOptionsPopup() {
+        return findViewById(R.id.popup_container);
+    }
 }
diff --git a/src/com/android/launcher3/logging/StatsLogManager.java b/src/com/android/launcher3/logging/StatsLogManager.java
index d987212..8b7bebc 100644
--- a/src/com/android/launcher3/logging/StatsLogManager.java
+++ b/src/com/android/launcher3/logging/StatsLogManager.java
@@ -518,7 +518,10 @@
         LAUNCHER_TASKBAR_LONGPRESS_HIDE(896),
 
         @UiEvent(doc = "User long pressed on the taskbar gesture handle to show the taskbar")
-        LAUNCHER_TASKBAR_LONGPRESS_SHOW(897);
+        LAUNCHER_TASKBAR_LONGPRESS_SHOW(897),
+
+        @UiEvent(doc = "User clicks on the search icon on header to launch search in app.")
+        LAUNCHER_ALLAPPS_SEARCHINAPP_LAUNCH(913);
 
         // ADD MORE
 
diff --git a/src/com/android/launcher3/model/data/SearchActionItemInfo.java b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
index 293c095..c6e5e8a 100644
--- a/src/com/android/launcher3/model/data/SearchActionItemInfo.java
+++ b/src/com/android/launcher3/model/data/SearchActionItemInfo.java
@@ -46,6 +46,7 @@
     public static final int FLAG_PRIMARY_ICON_FROM_TITLE = 1 << 4;
     public static final int FLAG_BADGE_WITH_COMPONENT_NAME = 1 << 5;
     public static final int FLAG_ALLOW_PINNING = 1 << 6;
+    public static final int FLAG_SEARCH_IN_APP = 1 << 7;
 
     private String mFallbackPackageName;
     private int mFlags = 0;
diff --git a/src/com/android/launcher3/popup/ArrowPopup.java b/src/com/android/launcher3/popup/ArrowPopup.java
index 2230914..5a1e4bf 100644
--- a/src/com/android/launcher3/popup/ArrowPopup.java
+++ b/src/com/android/launcher3/popup/ArrowPopup.java
@@ -234,7 +234,7 @@
      * @param backgroundColor When Color.TRANSPARENT, we get color from {@link #mColorIds}.
      *                        Otherwise, we will use this color for all child views.
      */
-    private void assignMarginsAndBackgrounds(ViewGroup viewGroup, int backgroundColor) {
+    protected void assignMarginsAndBackgrounds(ViewGroup viewGroup, int backgroundColor) {
         int[] colors = null;
         if (backgroundColor == Color.TRANSPARENT) {
             // Lazily get the colors so they match the current wallpaper colors.
@@ -445,7 +445,7 @@
         animateOpen();
     }
 
-    private void setupForDisplay() {
+    protected void setupForDisplay() {
         setVisibility(View.INVISIBLE);
         mIsOpen = true;
         getPopupContainer().addView(this);
@@ -482,7 +482,7 @@
         mArrow.setVisibility(show && shouldAddArrow() ? VISIBLE : INVISIBLE);
     }
 
-    private void addArrow() {
+    protected void addArrow() {
         getPopupContainer().addView(mArrow);
         mArrow.setX(getX() + getArrowLeft());
 
@@ -686,12 +686,13 @@
         return getChildCount() > 0 ? getChildAt(0) : this;
     }
 
-    private void animateOpen() {
+    protected void animateOpen() {
         setVisibility(View.VISIBLE);
 
         mOpenCloseAnimator = getOpenCloseAnimator(true, OPEN_DURATION, OPEN_FADE_START_DELAY,
                 OPEN_FADE_DURATION, OPEN_CHILD_FADE_START_DELAY, OPEN_CHILD_FADE_DURATION,
                 DECELERATED_EASE);
+        onCreateOpenAnimation(mOpenCloseAnimator);
         mOpenCloseAnimator.addListener(new AnimatorListenerAdapter() {
             @Override
             public void onAnimationEnd(Animator animation) {
@@ -785,6 +786,11 @@
     }
 
     /**
+     * Called when creating the open transition allowing subclass can add additional animations.
+     */
+    protected void onCreateOpenAnimation(AnimatorSet anim) { }
+
+    /**
      * Called when creating the close transition allowing subclass can add additional animations.
      */
     protected void onCreateCloseAnimation(AnimatorSet anim) { }
diff --git a/src/com/android/launcher3/touch/ItemClickHandler.java b/src/com/android/launcher3/touch/ItemClickHandler.java
index 5e907a4..8d57d69 100644
--- a/src/com/android/launcher3/touch/ItemClickHandler.java
+++ b/src/com/android/launcher3/touch/ItemClickHandler.java
@@ -17,6 +17,7 @@
 
 import static com.android.launcher3.Launcher.REQUEST_BIND_PENDING_APPWIDGET;
 import static com.android.launcher3.Launcher.REQUEST_RECONFIGURE_APPWIDGET;
+import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_ALLAPPS_SEARCHINAPP_LAUNCH;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_APP_LAUNCH_TAP;
 import static com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_FOLDER_OPEN;
 import static com.android.launcher3.model.data.ItemInfoWithIcon.FLAG_DISABLED_BY_PUBLISHER;
@@ -286,7 +287,13 @@
                         Toast.LENGTH_SHORT).show();
             }
         }
-        launcher.getStatsLogManager().logger().withItemInfo(itemInfo).log(LAUNCHER_APP_LAUNCH_TAP);
+        if (itemInfo.hasFlags(SearchActionItemInfo.FLAG_SEARCH_IN_APP)) {
+            launcher.getStatsLogManager().logger().withItemInfo(itemInfo).log(
+                    LAUNCHER_ALLAPPS_SEARCHINAPP_LAUNCH);
+        } else {
+            launcher.getStatsLogManager().logger().withItemInfo(itemInfo).log(
+                    LAUNCHER_APP_LAUNCH_TAP);
+        }
     }
 
     private static void startAppShortcutOrInfoActivity(View v, ItemInfo item, Launcher launcher) {
diff --git a/src/com/android/launcher3/touch/LandscapePagedViewHandler.java b/src/com/android/launcher3/touch/LandscapePagedViewHandler.java
index a190f52..93e3ea7 100644
--- a/src/com/android/launcher3/touch/LandscapePagedViewHandler.java
+++ b/src/com/android/launcher3/touch/LandscapePagedViewHandler.java
@@ -390,7 +390,7 @@
     @Override
     public void measureGroupedTaskViewThumbnailBounds(View primarySnapshot, View secondarySnapshot,
             int parentWidth, int parentHeight,
-            SplitConfigurationOptions.StagedSplitBounds splitBoundsConfig, DeviceProfile dp) {
+            StagedSplitBounds splitBoundsConfig, DeviceProfile dp) {
         int spaceAboveSnapshot = dp.overviewTaskThumbnailTopMarginPx;
         int totalThumbnailHeight = parentHeight - spaceAboveSnapshot;
         int dividerBar = splitBoundsConfig.visualDividerBounds.width();
@@ -435,15 +435,18 @@
                 (FrameLayout.LayoutParams) primaryIconView.getLayoutParams();
         FrameLayout.LayoutParams secondaryIconParams =
                 new FrameLayout.LayoutParams(primaryIconParams);
+        int dividerBar = (splitConfig.appsStackedVertically ?
+                splitConfig.visualDividerBounds.height() :
+                splitConfig.visualDividerBounds.width());
 
         int primaryHeight = primarySnapshotBounds.height();
-        int secondaryHeight = secondarySnapshotBounds.height();
         primaryIconParams.gravity = (isRtl ? START : END) | TOP;
-        primaryIconView.setTranslationY((primaryHeight + taskIconHeight) / 2f );
+        primaryIconView.setTranslationY(primaryHeight - primaryIconView.getHeight() / 2f);
+        primaryIconView.setTranslationX(0);
 
         secondaryIconParams.gravity = (isRtl ? START : END) | TOP;
-        secondaryIconView.setTranslationY(primaryHeight
-                + ((secondaryHeight + taskIconHeight) / 2f));
+        secondaryIconView.setTranslationY(primaryHeight + taskIconHeight + dividerBar);
+        secondaryIconView.setTranslationX(0);
         primaryIconView.setLayoutParams(primaryIconParams);
         secondaryIconView.setLayoutParams(secondaryIconParams);
     }
diff --git a/src/com/android/launcher3/touch/PagedOrientationHandler.java b/src/com/android/launcher3/touch/PagedOrientationHandler.java
index 8112afd..2ff2feb 100644
--- a/src/com/android/launcher3/touch/PagedOrientationHandler.java
+++ b/src/com/android/launcher3/touch/PagedOrientationHandler.java
@@ -141,13 +141,12 @@
      * @param desiredStagePosition Which stage position (topLeft/rightBottom) we want to resize
      *                           outRect for
      */
-    void setSplitTaskSwipeRect(DeviceProfile dp, Rect outRect,
-            StagedSplitBounds splitInfo,
+    void setSplitTaskSwipeRect(DeviceProfile dp, Rect outRect, StagedSplitBounds splitInfo,
             @SplitConfigurationOptions.StagePosition int desiredStagePosition);
 
     void measureGroupedTaskViewThumbnailBounds(View primarySnapshot, View secondarySnapshot,
             int parentWidth, int parentHeight,
-            SplitConfigurationOptions.StagedSplitBounds splitBoundsConfig, DeviceProfile dp);
+            StagedSplitBounds splitBoundsConfig, DeviceProfile dp);
 
     // Overview TaskMenuView methods
     void setIconAndSnapshotParams(View iconView, int taskIconMargin, int taskIconHeight,
diff --git a/src/com/android/launcher3/touch/PortraitPagedViewHandler.java b/src/com/android/launcher3/touch/PortraitPagedViewHandler.java
index 576c6f5..8caf886 100644
--- a/src/com/android/launcher3/touch/PortraitPagedViewHandler.java
+++ b/src/com/android/launcher3/touch/PortraitPagedViewHandler.java
@@ -16,7 +16,6 @@
 
 package com.android.launcher3.touch;
 
-import static android.view.Gravity.BOTTOM;
 import static android.view.Gravity.CENTER_HORIZONTAL;
 import static android.view.Gravity.START;
 import static android.view.Gravity.TOP;
@@ -477,7 +476,7 @@
     @Override
     public void measureGroupedTaskViewThumbnailBounds(View primarySnapshot, View secondarySnapshot,
             int parentWidth, int parentHeight,
-            SplitConfigurationOptions.StagedSplitBounds splitBoundsConfig, DeviceProfile dp) {
+            StagedSplitBounds splitBoundsConfig, DeviceProfile dp) {
         int spaceAboveSnapshot = dp.overviewTaskThumbnailTopMarginPx;
         int totalThumbnailHeight = parentHeight - spaceAboveSnapshot;
         int dividerBar = (splitBoundsConfig.appsStackedVertically ?
@@ -535,23 +534,27 @@
                 (FrameLayout.LayoutParams) primaryIconView.getLayoutParams();
         FrameLayout.LayoutParams secondaryIconParams =
                 new FrameLayout.LayoutParams(primaryIconParams);
+        int dividerBar = (splitConfig.appsStackedVertically ?
+                splitConfig.visualDividerBounds.height() :
+                splitConfig.visualDividerBounds.width());
 
+        int primaryWidth = primarySnapshotBounds.width();
         if (deviceProfile.isLandscape) {
-            int primaryWidth = primarySnapshotBounds.width();
-            int secondaryWidth = secondarySnapshotBounds.width();
             primaryIconParams.gravity = TOP | START;
-            primaryIconView.setTranslationX((primaryWidth - taskIconHeight) / 2f );
+            primaryIconView.setTranslationX(primaryWidth - primaryIconView.getWidth());
+            primaryIconView.setTranslationY(0);
 
             secondaryIconParams.gravity = TOP | START;
-            secondaryIconView.setTranslationX(primaryWidth
-                    + ((secondaryWidth - taskIconHeight) / 2f));
-        } else {
-            primaryIconView.setTranslationX(0);
-            secondaryIconView.setTranslationX(0);
-            primaryIconView.setTranslationY(0);
+            secondaryIconView.setTranslationX(primaryWidth + dividerBar);
             secondaryIconView.setTranslationY(0);
-            secondaryIconParams.gravity = BOTTOM | CENTER_HORIZONTAL;
-            secondaryIconParams.bottomMargin = -(secondaryIconParams.topMargin + taskIconHeight);
+        } else {
+            primaryIconParams.gravity = TOP | CENTER_HORIZONTAL;
+            primaryIconView.setTranslationX(-(primaryIconView.getWidth()) / 2f);
+            primaryIconView.setTranslationY(0);
+
+            secondaryIconParams.gravity = TOP | CENTER_HORIZONTAL;
+            secondaryIconView.setTranslationX(secondaryIconView.getWidth() / 2f);
+            secondaryIconView.setTranslationY(0);
         }
         primaryIconView.setLayoutParams(primaryIconParams);
         secondaryIconView.setLayoutParams(secondaryIconParams);
diff --git a/src/com/android/launcher3/touch/SeascapePagedViewHandler.java b/src/com/android/launcher3/touch/SeascapePagedViewHandler.java
index d5851c8..a0dde22 100644
--- a/src/com/android/launcher3/touch/SeascapePagedViewHandler.java
+++ b/src/com/android/launcher3/touch/SeascapePagedViewHandler.java
@@ -19,6 +19,7 @@
 import static android.view.Gravity.CENTER_VERTICAL;
 import static android.view.Gravity.END;
 import static android.view.Gravity.START;
+import static android.view.Gravity.TOP;
 
 import static com.android.launcher3.touch.SingleAxisSwipeDetector.HORIZONTAL;
 import static com.android.launcher3.util.SplitConfigurationOptions.STAGE_POSITION_BOTTOM_OR_RIGHT;
@@ -36,7 +37,9 @@
 import com.android.launcher3.DeviceProfile;
 import com.android.launcher3.R;
 import com.android.launcher3.Utilities;
+import com.android.launcher3.util.SplitConfigurationOptions;
 import com.android.launcher3.util.SplitConfigurationOptions.SplitPositionOption;
+import com.android.launcher3.util.SplitConfigurationOptions.StagedSplitBounds;
 import com.android.launcher3.views.BaseDragLayer;
 
 import java.util.Collections;
@@ -127,6 +130,23 @@
         iconParams.topMargin = snapshotParams.topMargin / 2;
     }
 
+    @Override
+    public void setSplitIconParams(View primaryIconView, View secondaryIconView,
+            int taskIconHeight, Rect primarySnapshotBounds, Rect secondarySnapshotBounds,
+            boolean isRtl, DeviceProfile deviceProfile, StagedSplitBounds splitConfig) {
+        super.setSplitIconParams(primaryIconView, secondaryIconView, taskIconHeight,
+                primarySnapshotBounds, secondarySnapshotBounds, isRtl, deviceProfile, splitConfig);
+        FrameLayout.LayoutParams primaryIconParams =
+                (FrameLayout.LayoutParams) primaryIconView.getLayoutParams();
+        FrameLayout.LayoutParams secondaryIconParams =
+                (FrameLayout.LayoutParams) secondaryIconView.getLayoutParams();
+
+        primaryIconParams.gravity = (isRtl ? END : START) | TOP;
+        secondaryIconParams.gravity = (isRtl ? END : START) | TOP;
+        primaryIconView.setLayoutParams(primaryIconParams);
+        secondaryIconView.setLayoutParams(secondaryIconParams);
+    }
+
     /* ---------- The following are only used by TaskViewTouchHandler. ---------- */
 
     @Override
diff --git a/src/com/android/launcher3/util/SplitConfigurationOptions.java b/src/com/android/launcher3/util/SplitConfigurationOptions.java
index 0b083e3..6aef38f 100644
--- a/src/com/android/launcher3/util/SplitConfigurationOptions.java
+++ b/src/com/android/launcher3/util/SplitConfigurationOptions.java
@@ -85,6 +85,12 @@
         }
     }
 
+    /**
+     * NOTE: Engineers complained about too little ambiguity in the last survey, so there is a class
+     * with the same name/functionality in wm.shell.util (which launcher3 cannot be built against)
+     *
+     * If you make changes here, consider making the same changes there
+     */
     public static class StagedSplitBounds {
         public final Rect leftTopBounds;
         public final Rect rightBottomBounds;
@@ -100,10 +106,15 @@
          * the bounds were originally in
          */
         public final boolean appsStackedVertically;
+        public final int leftTopTaskId;
+        public final int rightBottomTaskId;
 
-        public StagedSplitBounds(Rect leftTopBounds, Rect rightBottomBounds) {
+        public StagedSplitBounds(Rect leftTopBounds, Rect rightBottomBounds, int leftTopTaskId,
+                int rightBottomTaskId) {
             this.leftTopBounds = leftTopBounds;
             this.rightBottomBounds = rightBottomBounds;
+            this.leftTopTaskId = leftTopTaskId;
+            this.rightBottomTaskId = rightBottomTaskId;
 
             if (rightBottomBounds.top > leftTopBounds.top) {
                 // vertical apps, horizontal divider
diff --git a/src/com/android/launcher3/views/OptionsPopupView.java b/src/com/android/launcher3/views/OptionsPopupView.java
index 9774c46..fc8b4b7 100644
--- a/src/com/android/launcher3/views/OptionsPopupView.java
+++ b/src/com/android/launcher3/views/OptionsPopupView.java
@@ -37,7 +37,6 @@
 import android.widget.Toast;
 
 import androidx.annotation.Nullable;
-import androidx.annotation.VisibleForTesting;
 import androidx.core.content.ContextCompat;
 
 import com.android.launcher3.AbstractFloatingView;
@@ -176,11 +175,6 @@
         return children;
     }
 
-    @VisibleForTesting
-    public static ArrowPopup getOptionsPopup(Launcher launcher) {
-        return launcher.findViewById(R.id.popup_container);
-    }
-
     /**
      * Returns the list of supported actions
      */
diff --git a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
index f29ac23..a9bcd67 100644
--- a/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
+++ b/tests/src/com/android/launcher3/ui/TaplTestsLauncher3.java
@@ -38,7 +38,6 @@
 import com.android.launcher3.tapl.FolderIcon;
 import com.android.launcher3.tapl.Widgets;
 import com.android.launcher3.tapl.Workspace;
-import com.android.launcher3.views.OptionsPopupView;
 import com.android.launcher3.widget.picker.WidgetsFullSheet;
 import com.android.launcher3.widget.picker.WidgetsRecyclerView;
 
@@ -102,8 +101,8 @@
         mDevice.pressMenu();
         mDevice.waitForIdle();
         executeOnLauncher(
-                launcher -> assertTrue("Launcher internal state didn't switch to Showing Menu",
-                        OptionsPopupView.getOptionsPopup(launcher) != null));
+                launcher -> assertNotNull("Launcher internal state didn't switch to Showing Menu",
+                        launcher.getOptionsPopup()));
         // Check that pressHome works when the menu is shown.
         mLauncher.pressHome();
     }
@@ -286,7 +285,7 @@
     }
 
     private boolean isOptionsPopupVisible(Launcher launcher) {
-        final ArrowPopup popup = OptionsPopupView.getOptionsPopup(launcher);
+        final ArrowPopup<?> popup = launcher.getOptionsPopup();
         return popup != null && popup.isShown();
     }