From da9213791e915ea0483702949809cdc18ffc8cca Mon Sep 17 00:00:00 2001
From: Vinit Nayak <peanutbutter@google.com>
Date: Mon, 6 Jun 2022 09:18:42 -0700
Subject: [PATCH] Ensure null elements aren't added to TaskView shortcut list

Bug: 234874147
Change-Id: I5ed84fdc64128264e5808b7d3b019c47e3740469
---
 .../quickstep/TaskShortcutFactory.java        | 30 +++++++++++++------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
index 825cf2ca8f..2af6e0be6f 100644
--- a/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
+++ b/quickstep/src/com/android/quickstep/TaskShortcutFactory.java
@@ -74,6 +74,15 @@ public interface TaskShortcutFactory {
         return false;
     }
 
+    /** @return a singleton list if the provided shortcut is non-null, null otherwise */
+    @Nullable
+    default List<SystemShortcut> createSingletonShortcutList(@Nullable SystemShortcut shortcut) {
+        if (shortcut != null) {
+            return Collections.singletonList(shortcut);
+        }
+        return null;
+    }
+
     TaskShortcutFactory APP_INFO = new TaskShortcutFactory() {
         @Override
         public List<SystemShortcut> getShortcuts(BaseDraggingActivity activity,
@@ -373,9 +382,10 @@ public interface TaskShortcutFactory {
         @Override
         public List<SystemShortcut> getShortcuts(BaseDraggingActivity activity,
                 TaskIdAttributeContainer taskContainer) {
-            return Collections.singletonList(
+            SystemShortcut<BaseDraggingActivity> wellbeingShortcut =
                     WellbeingModel.SHORTCUT_FACTORY.getShortcut(activity,
-                            taskContainer.getItemInfo(), taskContainer.getTaskView()));
+                            taskContainer.getItemInfo(), taskContainer.getTaskView());
+            return createSingletonShortcutList(wellbeingShortcut);
         }
     };
 
@@ -383,10 +393,10 @@ public interface TaskShortcutFactory {
         @Override
         public List<SystemShortcut> getShortcuts(BaseDraggingActivity activity,
                 TaskIdAttributeContainer taskContainer) {
-            return Collections.singletonList(
-                    taskContainer.getThumbnailView().getTaskOverlay()
-                            .getScreenshotShortcut(activity, taskContainer.getItemInfo(),
-                                    taskContainer.getTaskView()));
+            SystemShortcut screenshotShortcut = taskContainer.getThumbnailView().getTaskOverlay()
+                    .getScreenshotShortcut(activity, taskContainer.getItemInfo(),
+                            taskContainer.getTaskView());
+            return createSingletonShortcutList(screenshotShortcut);
         }
     };
 
@@ -394,10 +404,12 @@ public interface TaskShortcutFactory {
         @Override
         public List<SystemShortcut> getShortcuts(BaseDraggingActivity activity,
                 TaskIdAttributeContainer taskContainer) {
+            SystemShortcut modalStateSystemShortcut =
+                    taskContainer.getThumbnailView().getTaskOverlay()
+                            .getModalStateSystemShortcut(
+                                    taskContainer.getItemInfo(), taskContainer.getTaskView());
             if (ENABLE_OVERVIEW_SELECTIONS.get()) {
-                return Collections.singletonList(taskContainer.getThumbnailView().getTaskOverlay()
-                        .getModalStateSystemShortcut(
-                                        taskContainer.getItemInfo(), taskContainer.getTaskView()));
+                return createSingletonShortcutList(modalStateSystemShortcut);
             }
             return null;
         }