Fix two bugs with thumbnail screenshotting for split apps

This patch fixes two bugs with the thumbnail screenshotting flow for split apps. There is no currently filed bug behavior for these bugs, but they look wrong in the code -- this is the proposed fix/refactor.

Fixes: 285191462
Test: Still passes presubmit
Change-Id: I43eefa334d66676b8175919e266b82947f1837e5
diff --git a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
index 87fbcae..6ab3bbe 100644
--- a/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
+++ b/quickstep/src/com/android/quickstep/AbsSwipeUpHandler.java
@@ -294,7 +294,8 @@
 
     private boolean mContinuingLastGesture;
 
-    private ThumbnailData mTaskSnapshot;
+    // Cache of recently-updated task snapshots, mapping task id to ThumbnailData
+    private HashMap<Integer, ThumbnailData> mTaskSnapshotCache = new HashMap<>();
 
     // Used to control launcher components throughout the swipe gesture.
     private AnimatorControllerWithResistance mLauncherTransitionController;
@@ -1908,7 +1909,7 @@
         mActivityInitListener.unregister();
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(
                 mActivityRestartListener);
-        mTaskSnapshot = null;
+        mTaskSnapshotCache.clear();
     }
 
     private void invalidateHandler() {
@@ -1926,7 +1927,7 @@
         mActivityInitListener.unregister();
         TaskStackChangeListeners.getInstance().unregisterTaskStackListener(
                 mActivityRestartListener);
-        mTaskSnapshot = null;
+        mTaskSnapshotCache.clear();
     }
 
     private void invalidateHandlerWithLauncher() {
@@ -1983,35 +1984,40 @@
         } else {
             final int runningTaskId = mGestureState.getRunningTaskId();
             boolean finishTransitionPosted = false;
+            // If we already have cached screenshot(s) from running tasks, skip update
+            boolean shouldUpdate = false;
+            int[] runningTaskIds = mIsSwipeForSplit
+                    ? TopTaskTracker.INSTANCE.get(mContext).getRunningSplitTaskIds()
+                    : new int[]{runningTaskId};
+            for (int id : runningTaskIds) {
+                if (!mTaskSnapshotCache.containsKey(id)) {
+                    shouldUpdate = true;
+                    break;
+                }
+            }
+
             if (mRecentsAnimationController != null) {
                 // Update the screenshot of the task
-                if (mTaskSnapshot == null) {
+                if (shouldUpdate) {
                     UI_HELPER_EXECUTOR.execute(() -> {
                         if (mRecentsAnimationController == null) return;
-                        final ThumbnailData taskSnapshot =
-                                mRecentsAnimationController.screenshotTask(runningTaskId);
-                        // If split case, we should update all split tasks snapshot
-                        if (mIsSwipeForSplit) {
-                            int[] splitTaskIds = TopTaskTracker.INSTANCE.get(
-                                    mContext).getRunningSplitTaskIds();
-                            for (int i = 0; i < splitTaskIds.length; i++) {
-                                // Skip running one because done above.
-                                if (splitTaskIds[i] == runningTaskId) continue;
-
-                                mRecentsAnimationController.screenshotTask(splitTaskIds[i]);
-                            }
+                        for (int id : runningTaskIds) {
+                            mTaskSnapshotCache.put(
+                                    id, mRecentsAnimationController.screenshotTask(id));
                         }
+
                         MAIN_EXECUTOR.execute(() -> {
-                            mTaskSnapshot = taskSnapshot;
-                            if (!updateThumbnail(runningTaskId, false /* refreshView */)) {
+                            if (!updateThumbnail(false /* refreshView */)) {
                                 setScreenshotCapturedState();
                             }
                         });
                     });
                     return;
                 }
-                finishTransitionPosted = updateThumbnail(runningTaskId, false /* refreshView */);
+
+                finishTransitionPosted = updateThumbnail(false /* refreshView */);
             }
+
             if (!finishTransitionPosted) {
                 setScreenshotCapturedState();
             }
@@ -2019,26 +2025,26 @@
     }
 
     // Returns whether finish transition was posted.
-    private boolean updateThumbnail(int runningTaskId, boolean refreshView) {
-        boolean finishTransitionPosted = false;
-        final TaskView taskView;
+    private boolean updateThumbnail(boolean refreshView) {
         if (mGestureState.getEndTarget() == HOME
                 || mGestureState.getEndTarget() == NEW_TASK
                 || mGestureState.getEndTarget() == ALL_APPS
                 || mRecentsView == null) {
             // Capture the screenshot before finishing the transition to home or quickswitching to
             // ensure it's taken in the correct orientation, but no need to update the thumbnail.
-            taskView = null;
-        } else {
-            taskView = mRecentsView.updateThumbnail(runningTaskId, mTaskSnapshot, refreshView);
+            return false;
         }
-        if (taskView != null && refreshView && !mCanceled) {
+
+        boolean finishTransitionPosted = false;
+        TaskView updatedTaskView = mRecentsView.updateThumbnail(mTaskSnapshotCache, refreshView);
+        if (updatedTaskView != null && refreshView && !mCanceled) {
             // Defer finishing the animation until the next launcher frame with the
             // new thumbnail
-            finishTransitionPosted = ViewUtils.postFrameDrawn(taskView,
+            finishTransitionPosted = ViewUtils.postFrameDrawn(updatedTaskView,
                     () -> mStateCallback.setStateOnUiThread(STATE_SCREENSHOT_CAPTURED),
                     this::isCanceled);
         }
+
         return finishTransitionPosted;
     }
 
diff --git a/quickstep/src/com/android/quickstep/views/RecentsView.java b/quickstep/src/com/android/quickstep/views/RecentsView.java
index bc6451c..ea027bf 100644
--- a/quickstep/src/com/android/quickstep/views/RecentsView.java
+++ b/quickstep/src/com/android/quickstep/views/RecentsView.java
@@ -211,6 +211,7 @@
 import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
@@ -992,16 +993,34 @@
     }
 
     /**
-     * Update the thumbnail of the task.
+     * Update the thumbnail(s) of the relevant TaskView.
      * @param refreshNow Refresh immediately if it's true.
      */
     @Nullable
-    public TaskView updateThumbnail(int taskId, ThumbnailData thumbnailData, boolean refreshNow) {
-        TaskView taskView = getTaskViewByTaskId(taskId);
-        if (taskView != null) {
-            taskView.getThumbnail().setThumbnail(taskView.getTask(), thumbnailData, refreshNow);
+    public TaskView updateThumbnail(
+            HashMap<Integer, ThumbnailData> thumbnailData, boolean refreshNow) {
+        TaskView updatedTaskView = null;
+        for (Map.Entry<Integer, ThumbnailData> entry : thumbnailData.entrySet()) {
+            Integer id = entry.getKey();
+            ThumbnailData thumbnail = entry.getValue();
+            TaskView taskView = getTaskViewByTaskId(id);
+            if (taskView == null) {
+                continue;
+            }
+            // taskView could be a GroupedTaskView, so select the relevant task by ID
+            TaskIdAttributeContainer taskAttributes = taskView.getTaskAttributesById(id);
+            if (taskAttributes == null) {
+                continue;
+            }
+            Task task = taskAttributes.getTask();
+            TaskThumbnailView taskThumbnailView = taskAttributes.getThumbnailView();
+            taskThumbnailView.setThumbnail(task, thumbnail, refreshNow);
+            // thumbnailData can contain 1-2 ids, but they should correspond to the same
+            // TaskView, so overwriting is ok
+            updatedTaskView = taskView;
         }
-        return taskView;
+
+        return updatedTaskView;
     }
 
     @Override