diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
index 87fa6f318f..bfdf156dbd 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarUIController.java
@@ -204,12 +204,11 @@ public class TaskbarUIController {
                         if (foundTaskView != null) {
                             // There is already a running app of this type, use that as second app.
                             // Get index of task (0 or 1), in case it's a GroupedTaskView
-                            int indexOfTask = foundTaskView.getIndexOfTask(foundTask.key.id);
                             TaskIdAttributeContainer taskAttributes =
-                                    foundTaskView.getTaskIdAttributeContainers()[indexOfTask];
+                                    foundTaskView.getTaskAttributesById(foundTask.key.id);
                             recents.confirmSplitSelect(
                                     foundTaskView,
-                                    taskAttributes.getTask(),
+                                    foundTask,
                                     taskAttributes.getIconView().getDrawable(),
                                     taskAttributes.getThumbnailView(),
                                     taskAttributes.getThumbnailView().getThumbnail(),
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 8c8c2ba8aa..e8c2050657 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -602,16 +602,18 @@ public class TaskView extends FrameLayout implements Reusable {
     }
 
     /**
-     * Finds the index of a given taskId within this TaskView, or -1 if the TaskView does not
-     * contain it. For grouped tasks (of two), this is 0 or 1; for solo tasks, it is 0.
+     * Returns the TaskIdAttributeContainer corresponding to a given taskId, or null if the TaskView
+     * does not contain a Task with that ID.
      */
-    public int getIndexOfTask(int taskId) {
-        for (int i = 0; i < mTaskIdContainer.length; i++) {
-            if (mTaskIdContainer[i] == taskId) {
-                return i;
+    @Nullable
+    public TaskIdAttributeContainer getTaskAttributesById(int taskId) {
+        for (TaskIdAttributeContainer attributes : mTaskIdAttributeContainer) {
+            if (attributes.getTask().key.id == taskId) {
+                return attributes;
             }
         }
-        return -1;
+
+        return null;
     }
 
     public TaskThumbnailView getThumbnail() {