Keep split on top when going to Pip if not swipe to home

* If main or side stage have a child, we move that stage to
be on top when performing the split to pip transition.
* We also prevent the stage with a child from dismissing in
DefaultMixedHandler#animateEnterPipFromSplit
* In both cases we use isSplitScreenVisible() as our proxy to determine
if we're dismissing one of the split pairs vs swiping up to home,
since in the swipe up case split becomes invisible and then Pip requests
to be shown as a separate transition

Bug: 293254265
Test: Steps don't repro. Pip shows up and hides when it's supposed to.
There's a KI quickswitch issue but happens w/o this change
(b/303722342)

(cherry picked from commit 8003d4cb028d484d1d4f11451e5c2620bbf09686)
(cherry picked from https://googleplex-android-review.googlesource.com/q/commit:ed9e0a512bc2b483907c69cc9358411ef9e07be7)
Merged-In: I7892ac1612ab2dccdd81c0bece987aab3ad69194
Change-Id: I7892ac1612ab2dccdd81c0bece987aab3ad69194

2 files changed, 54 insertions, 3 deletions

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 c95c9f099f20..1c792395d22d 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
@@ -2244,6 +2244,25 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         return SPLIT_POSITION_UNDEFINED;
     }
 
+    /**
+     * Returns the {@link StageType} where {@param token} is being used
+     * {@link SplitScreen#STAGE_TYPE_UNDEFINED} otherwise
+     */
+    @StageType
+    public int getSplitItemStage(@Nullable WindowContainerToken token) {
+        if (token == null) {
+            return STAGE_TYPE_UNDEFINED;
+        }
+
+        if (mMainStage.containsToken(token)) {
+            return STAGE_TYPE_MAIN;
+        } else if (mSideStage.containsToken(token)) {
+            return STAGE_TYPE_SIDE;
+        }
+
+        return STAGE_TYPE_UNDEFINED;
+    }
+
     @Override
     public void setLayoutOffsetTarget(int offsetX, int offsetY, SplitLayout layout) {
         final StageTaskListener topLeftStage =
@@ -2491,7 +2510,16 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 mRecentTasks.ifPresent(
                         recentTasks -> recentTasks.removeSplitPair(triggerTask.taskId));
             }
-            prepareExitSplitScreen(STAGE_TYPE_UNDEFINED, outWCT);
+            @StageType int topStage = STAGE_TYPE_UNDEFINED;
+            if (isSplitScreenVisible()) {
+                // Get the stage where a child exists to keep that stage onTop
+                if (mMainStage.getChildCount() != 0 && mSideStage.getChildCount() == 0) {
+                    topStage = STAGE_TYPE_MAIN;
+                } else if (mSideStage.getChildCount() != 0 && mMainStage.getChildCount() == 0) {
+                    topStage = STAGE_TYPE_SIDE;
+                }
+            }
+            prepareExitSplitScreen(topStage, outWCT);
         }
     }
 
@@ -2908,7 +2936,11 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         return SPLIT_POSITION_UNDEFINED;
     }
 
-    /** Synchronize split-screen state with transition and make appropriate preparations. */
+    /**
+     * Synchronize split-screen state with transition and make appropriate preparations.
+     * @param toStage The stage that will not be dismissed. If set to
+     *        {@link SplitScreen#STAGE_TYPE_UNDEFINED} then both stages will be dismissed
+     */
     public void prepareDismissAnimation(@StageType int toStage, @ExitReason int dismissReason,
             @NonNull TransitionInfo info, @NonNull SurfaceControl.Transaction t,
             @NonNull SurfaceControl.Transaction finishT) {
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
index 4897c4ee7dd6..f0bb665f8082 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedHandler.java
@@ -50,6 +50,7 @@ import com.android.wm.shell.keyguard.KeyguardTransitionHandler;
 import com.android.wm.shell.pip.PipTransitionController;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
 import com.android.wm.shell.recents.RecentsTransitionHandler;
+import com.android.wm.shell.splitscreen.SplitScreen;
 import com.android.wm.shell.splitscreen.SplitScreenController;
 import com.android.wm.shell.splitscreen.StageCoordinator;
 import com.android.wm.shell.sysui.ShellInit;
@@ -511,8 +512,26 @@ public class DefaultMixedHandler implements Transitions.TransitionHandler,
             // make a new startTransaction because pip's startEnterAnimation "consumes" it so
             // we need a separate one to send over to launcher.
             SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
+            @SplitScreen.StageType int topStageToKeep = STAGE_TYPE_UNDEFINED;
+            if (mSplitHandler.isSplitScreenVisible()) {
+                // The non-going home case, we could be pip-ing one of the split stages and keep
+                // showing the other
+                for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+                    TransitionInfo.Change change = info.getChanges().get(i);
+                    if (change == pipChange) {
+                        // Ignore the change/task that's going into Pip
+                        continue;
+                    }
+                    @SplitScreen.StageType int splitItemStage =
+                            mSplitHandler.getSplitItemStage(change.getLastParent());
+                    if (splitItemStage != STAGE_TYPE_UNDEFINED) {
+                        topStageToKeep = splitItemStage;
+                        break;
+                    }
+                }
+            }
             // Let split update internal state for dismiss.
-            mSplitHandler.prepareDismissAnimation(STAGE_TYPE_UNDEFINED,
+            mSplitHandler.prepareDismissAnimation(topStageToKeep,
                     EXIT_REASON_CHILD_TASK_ENTER_PIP, everythingElse, otherStartT,
                     finishTransaction);