Fix how task info is retrieved for split screen

Previously this was grabbing the top two tasks from the
task list, but this isn't guaranteed to be what's on the
screen. Instead, when already in split, use the taskInfo
from the split controller.

Test: manual - 1) have stuff in split
               2) keep dragging different apps into split
               => verify the app icon shown in the non
                  highlighted side matches the app on that
                  side
Bug: 209504662
Change-Id: Ib959f44bc9eaebfb5e3f640d69cb6f3ddccedfcf

3 files changed, 49 insertions, 32 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
index 67f9062b0a66..4d981f6aff50 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/draganddrop/DragLayout.java
@@ -18,6 +18,7 @@ package com.android.wm.shell.draganddrop;
 
 import static android.app.StatusBarManager.DISABLE_NONE;
 
+import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_TOP_OR_LEFT;
 
 import android.animation.Animator;
@@ -165,41 +166,40 @@ public class DragLayout extends LinearLayout {
         mHasDropped = false;
         mCurrentTarget = null;
 
-        List<ActivityManager.RunningTaskInfo> tasks = null;
-        // Figure out the splashscreen info for the existing task(s).
-        try {
-            tasks = ActivityTaskManager.getService().getTasks(2,
-                            false /* filterOnlyVisibleRecents */,
-                            false /* keepIntentExtra */);
-        } catch (RemoteException e) {
-            // don't show an icon / will just use the defaults
-        }
-        if (tasks != null && !tasks.isEmpty()) {
-            ActivityManager.RunningTaskInfo taskInfo1 = tasks.get(0);
-            Drawable icon1 = mIconProvider.getIcon(taskInfo1.topActivityInfo);
-            int bgColor1 = getResizingBackgroundColor(taskInfo1);
-
-            boolean alreadyInSplit = mSplitScreenController != null
-                    && mSplitScreenController.isSplitScreenVisible();
-            if (alreadyInSplit && tasks.size() > 1) {
-                ActivityManager.RunningTaskInfo taskInfo2 = tasks.get(1);
-                Drawable icon2 = mIconProvider.getIcon(taskInfo2.topActivityInfo);
-                int bgColor2 = getResizingBackgroundColor(taskInfo2);
-
-                // figure out which task is on which side
-                int splitPosition1 = mSplitScreenController.getSplitPosition(taskInfo1.taskId);
-                boolean isTask1TopOrLeft = splitPosition1 == SPLIT_POSITION_TOP_OR_LEFT;
-                if (isTask1TopOrLeft) {
-                    mDropZoneView1.setAppInfo(bgColor1, icon1);
-                    mDropZoneView2.setAppInfo(bgColor2, icon2);
-                } else {
-                    mDropZoneView2.setAppInfo(bgColor1, icon1);
-                    mDropZoneView1.setAppInfo(bgColor2, icon2);
-                }
-            } else {
+        boolean alreadyInSplit = mSplitScreenController != null
+                && mSplitScreenController.isSplitScreenVisible();
+        if (!alreadyInSplit) {
+            List<ActivityManager.RunningTaskInfo> tasks = null;
+            // Figure out the splashscreen info for the existing task.
+            try {
+                tasks = ActivityTaskManager.getService().getTasks(1,
+                        false /* filterOnlyVisibleRecents */,
+                        false /* keepIntentExtra */);
+            } catch (RemoteException e) {
+                // don't show an icon / will just use the defaults
+            }
+            if (tasks != null && !tasks.isEmpty()) {
+                ActivityManager.RunningTaskInfo taskInfo1 = tasks.get(0);
+                Drawable icon1 = mIconProvider.getIcon(taskInfo1.topActivityInfo);
+                int bgColor1 = getResizingBackgroundColor(taskInfo1);
                 mDropZoneView1.setAppInfo(bgColor1, icon1);
                 mDropZoneView2.setAppInfo(bgColor1, icon1);
             }
+        } else {
+            // We're already in split so get taskInfo from the controller to populate icon / color.
+            ActivityManager.RunningTaskInfo topOrLeftTask =
+                    mSplitScreenController.getTaskInfo(SPLIT_POSITION_TOP_OR_LEFT);
+            ActivityManager.RunningTaskInfo bottomOrRightTask =
+                    mSplitScreenController.getTaskInfo(SPLIT_POSITION_BOTTOM_OR_RIGHT);
+            if (topOrLeftTask != null && bottomOrRightTask != null) {
+                Drawable topOrLeftIcon = mIconProvider.getIcon(topOrLeftTask.topActivityInfo);
+                int topOrLeftColor = getResizingBackgroundColor(topOrLeftTask);
+                Drawable bottomOrRightIcon = mIconProvider.getIcon(
+                        bottomOrRightTask.topActivityInfo);
+                int bottomOrRightColor = getResizingBackgroundColor(bottomOrRightTask);
+                mDropZoneView1.setAppInfo(topOrLeftColor, topOrLeftIcon);
+                mDropZoneView2.setAppInfo(bottomOrRightColor, bottomOrRightIcon);
+            }
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
index 05552aad2303..03dbc2b4ef53 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java
@@ -184,6 +184,15 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
         return mStageCoordinator.isSplitScreenVisible();
     }
 
+    @Nullable
+    public ActivityManager.RunningTaskInfo getTaskInfo(@SplitPosition int splitPosition) {
+        if (isSplitScreenVisible()) {
+            int taskId = mStageCoordinator.getTaskId(splitPosition);
+            return mTaskOrganizer.getRunningTaskInfo(taskId);
+        }
+        return null;
+    }
+
     public boolean isTaskInSplitScreen(int taskId) {
         return isSplitScreenVisible()
                 && mStageCoordinator.getStageOfTask(taskId) != STAGE_TYPE_UNDEFINED;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
index e30e6c537c93..8f4b2c878f7d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java
@@ -521,6 +521,14 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         return SplitLayout.reversePosition(mSideStagePosition);
     }
 
+    int getTaskId(@SplitPosition int splitPosition) {
+        if (mSideStagePosition == splitPosition) {
+            return mSideStage.getTopVisibleChildTaskId();
+        } else {
+            return mMainStage.getTopVisibleChildTaskId();
+        }
+    }
+
     void setSideStagePosition(@SplitPosition int sideStagePosition,
             @Nullable WindowContainerTransaction wct) {
         setSideStagePosition(sideStagePosition, true /* updateBounds */, wct);