Support moving task to an activated split stage in a specific position

Split screen might has been activated when other components like
Launcher or PIP task organizer invoking enterSplitScreen API. This makes
sure the indicated task will be moved to the correct split side after
split screen activated.

Fix: 206745122
Test: atest WMShellUnitTests
Test: manual check expanding PIP task to split behavior
Change-Id: I6210846c6f4443727e93fcf313539a34c0f3298a

6 files changed, 104 insertions, 48 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
index 625bcee673b5..d07fff3cce79 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/common/split/SplitLayout.java
@@ -34,6 +34,7 @@ import android.animation.Animator;
 import android.animation.AnimatorListenerAdapter;
 import android.animation.ValueAnimator;
 import android.annotation.IntDef;
+import android.annotation.NonNull;
 import android.app.ActivityManager;
 import android.content.Context;
 import android.content.res.Configuration;
@@ -61,6 +62,8 @@ import com.android.wm.shell.animation.Interpolators;
 import com.android.wm.shell.common.DisplayImeController;
 import com.android.wm.shell.common.DisplayInsetsController;
 
+import java.io.PrintWriter;
+
 /**
  * Records and handles layout of splits. Helps to calculate proper bounds when configuration or
  * divide position changes.
@@ -415,6 +418,19 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
         return bounds.width() > bounds.height();
     }
 
+    /** Reverse the split position. */
+    @SplitPosition
+    public static int reversePosition(@SplitPosition int position) {
+        switch (position) {
+            case SPLIT_POSITION_TOP_OR_LEFT:
+                return SPLIT_POSITION_BOTTOM_OR_RIGHT;
+            case SPLIT_POSITION_BOTTOM_OR_RIGHT:
+                return SPLIT_POSITION_TOP_OR_LEFT;
+            default:
+                return SPLIT_POSITION_UNDEFINED;
+        }
+    }
+
     /**
      * Return if this layout is landscape.
      */
@@ -502,6 +518,13 @@ public final class SplitLayout implements DisplayInsetsController.OnInsetsChange
         }
     }
 
+    /** Dumps the current split bounds recorded in this layout. */
+    public void dump(@NonNull PrintWriter pw, String prefix) {
+        pw.println(prefix + "bounds1=" + mBounds1.toShortString());
+        pw.println(prefix + "dividerBounds=" + mDividerBounds.toShortString());
+        pw.println(prefix + "bounds2=" + mBounds2.toShortString());
+    }
+
     /** Handles layout change event. */
     public interface SplitLayoutHandler {
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
index f8c03044c5c9..1ba1d22e7831 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SideStage.java
@@ -50,10 +50,6 @@ class SideStage extends StageTaskListener {
         wct.setBounds(rootToken, rootBounds).reorder(rootToken, true /* onTop */);
     }
 
-    void addTask(ActivityManager.RunningTaskInfo task, WindowContainerTransaction wct) {
-        wct.reparent(task.token, mRootTaskInfo.token, true /* onTop*/);
-    }
-
     boolean removeAllTasks(WindowContainerTransaction wct, boolean toTop) {
         // No matter if the root task is empty or not, moving the root to bottom because it no
         // longer preserves visible child task.
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 7457be2d0871..207338fa122b 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
@@ -24,6 +24,8 @@ import static android.view.RemoteAnimationTarget.MODE_OPENING;
 import static com.android.wm.shell.common.ExecutorUtils.executeRemoteCallWithTaskPermission;
 import static com.android.wm.shell.common.split.SplitLayout.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.common.split.SplitLayout.SPLIT_POSITION_TOP_OR_LEFT;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_SIDE;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
 
 import android.app.ActivityManager;
 import android.app.ActivityTaskManager;
@@ -195,30 +197,17 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
     }
 
     public boolean moveToSideStage(int taskId, @SplitPosition int sideStagePosition) {
-        final ActivityManager.RunningTaskInfo task = mTaskOrganizer.getRunningTaskInfo(taskId);
-        if (task == null) {
-            throw new IllegalArgumentException("Unknown taskId" + taskId);
-        }
-        return moveToSideStage(task, sideStagePosition);
+        return moveToStage(taskId, STAGE_TYPE_SIDE, sideStagePosition,
+                new WindowContainerTransaction());
     }
 
-    public boolean moveToSideStage(int taskId, @SplitPosition int sideStagePosition,
-            WindowContainerTransaction wct) {
+    private boolean moveToStage(int taskId, @SplitScreen.StageType int stageType,
+            @SplitPosition int stagePosition, WindowContainerTransaction wct) {
         final ActivityManager.RunningTaskInfo task = mTaskOrganizer.getRunningTaskInfo(taskId);
         if (task == null) {
             throw new IllegalArgumentException("Unknown taskId" + taskId);
         }
-        return moveToSideStage(task, sideStagePosition, wct);
-    }
-
-    public boolean moveToSideStage(ActivityManager.RunningTaskInfo task,
-            @SplitPosition int sideStagePosition) {
-        return mStageCoordinator.moveToSideStage(task, sideStagePosition);
-    }
-
-    public boolean moveToSideStage(ActivityManager.RunningTaskInfo task,
-            @SplitPosition int sideStagePosition, WindowContainerTransaction wct) {
-        return mStageCoordinator.moveToSideStage(task, sideStagePosition, wct);
+        return mStageCoordinator.moveToStage(task, stageType, stagePosition, wct);
     }
 
     public boolean removeFromSideStage(int taskId) {
@@ -234,13 +223,14 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
     }
 
     public void enterSplitScreen(int taskId, boolean leftOrTop) {
-        moveToSideStage(taskId,
-                leftOrTop ? SPLIT_POSITION_TOP_OR_LEFT : SPLIT_POSITION_BOTTOM_OR_RIGHT);
+        enterSplitScreen(taskId, leftOrTop, new WindowContainerTransaction());
     }
 
     public void enterSplitScreen(int taskId, boolean leftOrTop, WindowContainerTransaction wct) {
-        moveToSideStage(taskId,
-                leftOrTop ? SPLIT_POSITION_TOP_OR_LEFT : SPLIT_POSITION_BOTTOM_OR_RIGHT, wct);
+        final int stageType = isSplitScreenVisible() ? STAGE_TYPE_UNDEFINED : STAGE_TYPE_SIDE;
+        final int stagePosition =
+                leftOrTop ? SPLIT_POSITION_TOP_OR_LEFT : SPLIT_POSITION_BOTTOM_OR_RIGHT;
+        moveToStage(taskId, stageType, stagePosition, wct);
     }
 
     public void exitSplitScreen(int toTopTaskId, @ExitReason int exitReason) {
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 d4941916850d..d0fc91ca8cf3 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
@@ -273,18 +273,31 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         return mSideStageListener.mVisible && mMainStageListener.mVisible;
     }
 
-    boolean moveToSideStage(ActivityManager.RunningTaskInfo task,
-            @SplitPosition int sideStagePosition) {
-        final WindowContainerTransaction wct = new WindowContainerTransaction();
-        return moveToSideStage(task, sideStagePosition, wct);
-    }
+    boolean moveToStage(ActivityManager.RunningTaskInfo task, @SplitScreen.StageType int stageType,
+            @SplitPosition int stagePosition, WindowContainerTransaction wct) {
+        StageTaskListener targetStage;
+        int sideStagePosition;
+        if (stageType == STAGE_TYPE_MAIN) {
+            targetStage = mMainStage;
+            sideStagePosition = SplitLayout.reversePosition(stagePosition);
+        } else if (stageType == STAGE_TYPE_SIDE) {
+            targetStage = mSideStage;
+            sideStagePosition = stagePosition;
+        } else {
+            if (mMainStage.isActive()) {
+                // If the split screen is activated, retrieves target stage based on position.
+                targetStage = stagePosition == mSideStagePosition ? mSideStage : mMainStage;
+                sideStagePosition = mSideStagePosition;
+            } else {
+                targetStage = mSideStage;
+                sideStagePosition = stagePosition;
+            }
+        }
 
-    boolean moveToSideStage(ActivityManager.RunningTaskInfo task,
-            @SplitPosition int sideStagePosition, WindowContainerTransaction wct) {
-        final WindowContainerTransaction evictWct = new WindowContainerTransaction();
         setSideStagePosition(sideStagePosition, wct);
-        mSideStage.evictAllChildren(evictWct);
-        mSideStage.addTask(task, wct);
+        final WindowContainerTransaction evictWct = new WindowContainerTransaction();
+        targetStage.evictAllChildren(evictWct);
+        targetStage.addTask(task, wct);
         if (!evictWct.isEmpty()) {
             wct.merge(evictWct, true /* transfer */);
         }
@@ -463,9 +476,7 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             case STAGE_TYPE_MAIN: {
                 if (position != SPLIT_POSITION_UNDEFINED) {
                     // Set the side stage opposite of what we want to the main stage.
-                    final int sideStagePosition = position == SPLIT_POSITION_TOP_OR_LEFT
-                            ? SPLIT_POSITION_BOTTOM_OR_RIGHT : SPLIT_POSITION_TOP_OR_LEFT;
-                    setSideStagePosition(sideStagePosition, wct);
+                    setSideStagePosition(SplitLayout.reversePosition(position), wct);
                 } else {
                     position = getMainStagePosition();
                 }
@@ -489,8 +500,7 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     @SplitLayout.SplitPosition
     int getMainStagePosition() {
-        return mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
-                ? SPLIT_POSITION_BOTTOM_OR_RIGHT : SPLIT_POSITION_TOP_OR_LEFT;
+        return SplitLayout.reversePosition(mSideStagePosition);
     }
 
     void setSideStagePosition(@SplitPosition int sideStagePosition,
@@ -870,8 +880,7 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     @Override
     public void onDoubleTappedDivider() {
-        setSideStagePosition(mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
-                ? SPLIT_POSITION_BOTTOM_OR_RIGHT : SPLIT_POSITION_TOP_OR_LEFT, null /* wct */);
+        setSideStagePosition(SplitLayout.reversePosition(mSideStagePosition), null /* wct */);
         mLogger.logSwap(getMainStagePosition(), mMainStage.getTopChildTaskUid(),
                 getSideStagePosition(), mSideStage.getTopChildTaskUid(),
                 mSplitLayout.isLandscape());
@@ -1296,11 +1305,16 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         pw.println(prefix + TAG + " mDisplayId=" + mDisplayId);
         pw.println(innerPrefix + "mDividerVisible=" + mDividerVisible);
         pw.println(innerPrefix + "MainStage");
+        pw.println(childPrefix + "stagePosition=" + getMainStagePosition());
         pw.println(childPrefix + "isActive=" + mMainStage.isActive());
         mMainStageListener.dump(pw, childPrefix);
         pw.println(innerPrefix + "SideStage");
+        pw.println(childPrefix + "stagePosition=" + getSideStagePosition());
         mSideStageListener.dump(pw, childPrefix);
-        pw.println(innerPrefix + "mSplitLayout=" + mSplitLayout);
+        if (mMainStage.isActive()) {
+            pw.println(innerPrefix + "SplitLayout");
+            mSplitLayout.dump(pw, childPrefix);
+        }
     }
 
     /**
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
index 62b8638a2582..5b0824567194 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageTaskListener.java
@@ -293,6 +293,10 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         }
     }
 
+    void addTask(ActivityManager.RunningTaskInfo task, WindowContainerTransaction wct) {
+        wct.reparent(task.token, mRootTaskInfo.token, true /* onTop*/);
+    }
+
     void setBounds(Rect bounds, WindowContainerTransaction wct) {
         wct.setBounds(mRootTaskInfo.token, bounds);
     }
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
index ad65c046be65..ef14d84c6a09 100644
--- a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/splitscreen/StageCoordinatorTests.java
@@ -19,11 +19,14 @@ package com.android.wm.shell.splitscreen;
 import static android.app.ActivityTaskManager.INVALID_TASK_ID;
 import static android.view.Display.DEFAULT_DISPLAY;
 
-import static com.android.internal.util.FrameworkStatsLog.SPLITSCREEN_UICHANGED__EXIT_REASON__RETURN_HOME;
 import static com.android.wm.shell.common.split.SplitLayout.SPLIT_POSITION_BOTTOM_OR_RIGHT;
 import static com.android.wm.shell.common.split.SplitLayout.SPLIT_POSITION_TOP_OR_LEFT;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_MAIN;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_SIDE;
+import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_UNDEFINED;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_RETURN_HOME;
 
+import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
@@ -49,7 +52,6 @@ import com.android.wm.shell.common.DisplayInsetsController;
 import com.android.wm.shell.common.SyncTransactionQueue;
 import com.android.wm.shell.common.TransactionPool;
 import com.android.wm.shell.common.split.SplitLayout;
-import com.android.wm.shell.recents.RecentTasksController;
 import com.android.wm.shell.transition.Transitions;
 
 import org.junit.Before;
@@ -110,12 +112,39 @@ public class StageCoordinatorTests extends ShellTestCase {
     }
 
     @Test
-    public void testMoveToSideStage() {
+    public void testMoveToStage() {
         final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
 
-        mStageCoordinator.moveToSideStage(task, SPLIT_POSITION_BOTTOM_OR_RIGHT);
+        mStageCoordinator.moveToStage(task, STAGE_TYPE_MAIN, SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                new WindowContainerTransaction());
+        verify(mMainStage).addTask(eq(task), any(WindowContainerTransaction.class));
+        assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getMainStagePosition());
 
+        mStageCoordinator.moveToStage(task, STAGE_TYPE_SIDE, SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                new WindowContainerTransaction());
         verify(mSideStage).addTask(eq(task), any(WindowContainerTransaction.class));
+        assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getSideStagePosition());
+    }
+
+    @Test
+    public void testMoveToUndefinedStage() {
+        final ActivityManager.RunningTaskInfo task = new TestRunningTaskInfoBuilder().build();
+
+        // Verify move to undefined stage while split screen not activated moves task to side stage.
+        when(mMainStage.isActive()).thenReturn(false);
+        mStageCoordinator.setSideStagePosition(SPLIT_POSITION_TOP_OR_LEFT, null);
+        mStageCoordinator.moveToStage(task, STAGE_TYPE_UNDEFINED, SPLIT_POSITION_BOTTOM_OR_RIGHT,
+                new WindowContainerTransaction());
+        verify(mSideStage).addTask(eq(task), any(WindowContainerTransaction.class));
+        assertEquals(SPLIT_POSITION_BOTTOM_OR_RIGHT, mStageCoordinator.getSideStagePosition());
+
+        // Verify move to undefined stage after split screen activated moves task based on position.
+        when(mMainStage.isActive()).thenReturn(true);
+        assertEquals(SPLIT_POSITION_TOP_OR_LEFT, mStageCoordinator.getMainStagePosition());
+        mStageCoordinator.moveToStage(task, STAGE_TYPE_UNDEFINED, SPLIT_POSITION_TOP_OR_LEFT,
+                new WindowContainerTransaction());
+        verify(mMainStage).addTask(eq(task), any(WindowContainerTransaction.class));
+        assertEquals(SPLIT_POSITION_TOP_OR_LEFT, mStageCoordinator.getMainStagePosition());
     }
 
     @Test