diff --git a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
index 71d4aa0..df99d27 100644
--- a/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
+++ b/quickstep/src/com/android/quickstep/views/GroupedTaskView.java
@@ -79,7 +79,7 @@
         mSecondaryTask = secondary;
         mTaskIdContainer[1] = secondary.key.id;
         mTaskIdAttributeContainer[1] = new TaskIdAttributeContainer(secondary, mSnapshotView2,
-                mIconView2, STAGE_POSITION_BOTTOM_OR_RIGHT);
+                STAGE_POSITION_BOTTOM_OR_RIGHT);
         mTaskIdAttributeContainer[0].setStagePosition(STAGE_POSITION_TOP_OR_LEFT);
         mSnapshotView2.bind(secondary);
         mSplitBoundsConfig = splitBoundsConfig;
@@ -120,6 +120,14 @@
         }
     }
 
+    protected boolean showTaskMenuWithContainer(IconView iconView) {
+        if (mActivity.getDeviceProfile().overviewShowAsGrid) {
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
+        } else {
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
+        }
+    }
+
     public void updateSplitBoundsConfig(StagedSplitBounds stagedSplitBounds) {
         mSplitBoundsConfig = stagedSplitBounds;
         invalidate();
diff --git a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
index f891e03..5059f8b 100644
--- a/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
+++ b/quickstep/src/com/android/quickstep/views/TaskMenuViewWithArrow.kt
@@ -61,13 +61,14 @@
 
     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(taskContainer.iconView, outPos)
+        popupContainer.getDescendantRectRelativeToSelf(taskView.iconView, outPos)
     }
 
     override fun onControllerInterceptTouchEvent(ev: MotionEvent?): Boolean {
@@ -90,6 +91,7 @@
             return false
         }
 
+        taskView = taskContainer.taskView
         this.taskContainer = taskContainer
         if (!populateMenu()) return false
         show()
@@ -109,8 +111,7 @@
     private fun addMenuOptions() {
         // Add the options
         TaskOverlayFactory
-            .getEnabledShortcuts(taskContainer.taskView, mActivityContext.deviceProfile,
-                    taskContainer)
+            .getEnabledShortcuts(taskView, mActivityContext.deviceProfile, taskContainer)
             .forEach { this.addMenuOption(it) }
 
         // Add the spaces between items
diff --git a/quickstep/src/com/android/quickstep/views/TaskView.java b/quickstep/src/com/android/quickstep/views/TaskView.java
index 6c27b96..e9a3779 100644
--- a/quickstep/src/com/android/quickstep/views/TaskView.java
+++ b/quickstep/src/com/android/quickstep/views/TaskView.java
@@ -539,7 +539,7 @@
         mTask = task;
         mTaskIdContainer[0] = mTask.key.id;
         mTaskIdAttributeContainer[0] = new TaskIdAttributeContainer(task, mSnapshotView,
-                mIconView, STAGE_POSITION_UNDEFINED);
+                STAGE_POSITION_UNDEFINED);
         mSnapshotView.bind(task);
         setOrientationState(orientedState);
     }
@@ -828,12 +828,10 @@
     }
 
     protected boolean showTaskMenuWithContainer(IconView iconView) {
-        TaskIdAttributeContainer menuContainer =
-                mTaskIdAttributeContainer[iconView == mIconView ? 0 : 1];
         if (mActivity.getDeviceProfile().overviewShowAsGrid) {
-            return TaskMenuViewWithArrow.Companion.showForTask(menuContainer);
+            return TaskMenuViewWithArrow.Companion.showForTask(mTaskIdAttributeContainer[0]);
         } else {
-            return TaskMenuView.showForTask(menuContainer);
+            return TaskMenuView.showForTask(mTaskIdAttributeContainer[0]);
         }
     }
 
@@ -1555,15 +1553,13 @@
     public class TaskIdAttributeContainer {
         private final TaskThumbnailView mThumbnailView;
         private final Task mTask;
-        private final IconView mIconView;
         /** Defaults to STAGE_POSITION_UNDEFINED if in not a split screen task view */
         private @SplitConfigurationOptions.StagePosition int mStagePosition;
 
         public TaskIdAttributeContainer(Task task, TaskThumbnailView thumbnailView,
-                IconView iconView, int stagePosition) {
+                int stagePosition) {
             this.mTask = task;
             this.mThumbnailView = thumbnailView;
-            this.mIconView = iconView;
             this.mStagePosition = stagePosition;
         }
 
@@ -1583,10 +1579,6 @@
             return TaskView.this;
         }
 
-        public IconView getIconView() {
-            return mIconView;
-        }
-
         public int getStagePosition() {
             return mStagePosition;
         }