Add path to allow task to replace Pipped task without breaking split

* StageCoordinator wasns't handling the case where one of the
splits goes into Pip but doesn't break the split
* We now preemptively guess that an enter transition will
occur when we get the transition that Pip's one of the split tasks
and opens another.
* This is done because Pipping a task out of split results in
2 manual transitions being created in MixedTransitionHelper and
running through StageCoordinator's startAnimation() call again

Bug: 326869501
Test: Repro steps in bug don't break split.
Regular split to pip still works

Change-Id: I2717b1e25d5ac84f77a7a7dc56cee51dcfa08683

8 files changed, 97 insertions, 20 deletions

diff --git a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/TransitionUtil.java b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/TransitionUtil.java
index dcd4062cb819..785e30d879d2 100644
--- a/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/TransitionUtil.java
+++ b/libs/WindowManager/Shell/shared/src/com/android/wm/shell/shared/TransitionUtil.java
@@ -69,8 +69,12 @@ public class TransitionUtil {
 
     /** Returns {@code true} if the transition is opening or closing mode. */
     public static boolean isOpenOrCloseMode(@TransitionInfo.TransitionMode int mode) {
-        return mode == TRANSIT_OPEN || mode == TRANSIT_CLOSE
-                || mode == TRANSIT_TO_FRONT || mode == TRANSIT_TO_BACK;
+        return isOpeningMode(mode) || mode == TRANSIT_CLOSE || mode == TRANSIT_TO_BACK;
+    }
+
+    /** Returns {@code true} if the transition is opening mode. */
+    public static boolean isOpeningMode(@TransitionInfo.TransitionMode int mode) {
+        return mode == TRANSIT_OPEN || mode == TRANSIT_TO_FRONT;
     }
 
     /** Returns {@code true} if the transition has a display change. */
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index e885262658f4..e1657f99639d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -854,7 +854,8 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
         mPipUiEventLoggerLogger.log(uiEventEnum);
 
         ProtoLog.d(ShellProtoLogGroup.WM_SHELL_PICTURE_IN_PICTURE,
-                "onTaskAppeared: %s, state=%s", mTaskInfo.topActivity, mPipTransitionState);
+                "onTaskAppeared: %s, state=%s, taskId=%s", mTaskInfo.topActivity,
+                mPipTransitionState, mTaskInfo.taskId);
         if (mPipTransitionState.getInSwipePipToHomeTransition()) {
             if (!mWaitForFixedRotation) {
                 onEndOfSwipePipToHomeTransition();
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 b10176df5826..8e97068dce84 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
@@ -47,6 +47,7 @@ import static com.android.wm.shell.common.split.SplitScreenUtils.reverseSplitPos
 import static com.android.wm.shell.common.split.SplitScreenUtils.splitFailureMessage;
 import static com.android.wm.shell.protolog.ShellProtoLogGroup.WM_SHELL_SPLIT_SCREEN;
 import static com.android.wm.shell.shared.TransitionUtil.isClosingType;
+import static com.android.wm.shell.shared.TransitionUtil.isOpeningMode;
 import static com.android.wm.shell.shared.TransitionUtil.isOpeningType;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_MAIN;
 import static com.android.wm.shell.splitscreen.SplitScreen.STAGE_TYPE_SIDE;
@@ -67,6 +68,7 @@ import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP;
 import static com.android.wm.shell.splitscreen.SplitScreenController.EXIT_REASON_UNKNOWN;
 import static com.android.wm.shell.splitscreen.SplitScreenController.exitReasonToString;
+import static com.android.wm.shell.transition.MixedTransitionHelper.getPipReplacingChange;
 import static com.android.wm.shell.transition.Transitions.ENABLE_SHELL_TRANSITIONS;
 import static com.android.wm.shell.transition.Transitions.TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE;
 import static com.android.wm.shell.transition.Transitions.TRANSIT_SPLIT_SCREEN_PAIR_OPEN;
@@ -2836,7 +2838,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             mSplitLayout.setFreezeDividerWindow(false);
             final StageChangeRecord record = new StageChangeRecord();
             final int transitType = info.getType();
-            boolean hasEnteringPip = false;
+            TransitionInfo.Change pipChange = null;
             for (int iC = 0; iC < info.getChanges().size(); ++iC) {
                 final TransitionInfo.Change change = info.getChanges().get(iC);
                 if (change.getMode() == TRANSIT_CHANGE
@@ -2847,7 +2849,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 }
 
                 if (mMixedHandler.isEnteringPip(change, transitType)) {
-                    hasEnteringPip = true;
+                    pipChange = change;
                 }
 
                 final ActivityManager.RunningTaskInfo taskInfo = change.getTaskInfo();
@@ -2899,9 +2901,19 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 }
             }
 
-            if (hasEnteringPip) {
+            if (pipChange != null) {
+                TransitionInfo.Change pipReplacingChange = getPipReplacingChange(info, pipChange,
+                        mMainStage.mRootTaskInfo.taskId, mSideStage.mRootTaskInfo.taskId,
+                        getSplitItemStage(pipChange.getLastParent()));
+                if (pipReplacingChange != null) {
+                    // Set an enter transition for when startAnimation gets called again
+                    mSplitTransitions.setEnterTransition(transition, /*remoteTransition*/ null,
+                            TRANSIT_SPLIT_SCREEN_OPEN_TO_SIDE, /*resizeAnim*/ false);
+                }
+
                 mMixedHandler.animatePendingEnterPipFromSplit(transition, info,
-                        startTransaction, finishTransaction, finishCallback);
+                        startTransaction, finishTransaction, finishCallback,
+                        pipReplacingChange != null);
                 notifySplitAnimationFinished();
                 return true;
             }
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 1e305c5dbbcf..f77c80d2ff3b 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
@@ -177,9 +177,11 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     @Override
     @CallSuper
     public void onTaskAppeared(ActivityManager.RunningTaskInfo taskInfo, SurfaceControl leash) {
-        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskAppeared: task=%d taskParent=%d rootTask=%d",
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskAppeared: taskId=%d taskParent=%d rootTask=%d "
+                        + "taskActivity=%s",
                 taskInfo.taskId, taskInfo.parentTaskId,
-                mRootTaskInfo != null ? mRootTaskInfo.taskId : -1);
+                mRootTaskInfo != null ? mRootTaskInfo.taskId : -1,
+                taskInfo.baseActivity);
         if (mRootTaskInfo == null) {
             mRootLeash = leash;
             mRootTaskInfo = taskInfo;
@@ -213,6 +215,8 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
     @Override
     @CallSuper
     public void onTaskInfoChanged(ActivityManager.RunningTaskInfo taskInfo) {
+        ProtoLog.d(WM_SHELL_SPLIT_SCREEN, "onTaskInfoChanged: taskId=%d taskAct=%s",
+                taskInfo.taskId, taskInfo.baseActivity);
         mWindowDecorViewModel.ifPresent(viewModel -> viewModel.onTaskInfoChanged(taskInfo));
         if (mRootTaskInfo.taskId == taskInfo.taskId) {
             // Inflates split decor view only when the root task is visible.
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 968b27b62490..bcacecbd8981 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
@@ -77,6 +77,7 @@ public class DefaultMixedHandler implements MixedTransitionHandler,
     private ActivityEmbeddingController mActivityEmbeddingController;
 
     abstract static class MixedTransition {
+        /** Entering Pip from split, breaks split. */
         static final int TYPE_ENTER_PIP_FROM_SPLIT = 1;
 
         /** Both the display and split-state (enter/exit) is changing */
@@ -103,6 +104,9 @@ public class DefaultMixedHandler implements MixedTransitionHandler,
         /** Enter pip from one of the Activity Embedding windows. */
         static final int TYPE_ENTER_PIP_FROM_ACTIVITY_EMBEDDING = 9;
 
+        /** Entering Pip from split, but replace the Pip stage instead of breaking split. */
+        static final int TYPE_ENTER_PIP_REPLACE_FROM_SPLIT = 10;
+
         /** The default animation for this mixed transition. */
         static final int ANIM_TYPE_DEFAULT = 0;
 
@@ -484,9 +488,11 @@ public class DefaultMixedHandler implements MixedTransitionHandler,
     // TODO(b/287704263): Remove when split/mixed are reversed.
     public boolean animatePendingEnterPipFromSplit(IBinder transition, TransitionInfo info,
             SurfaceControl.Transaction startT, SurfaceControl.Transaction finishT,
-            Transitions.TransitionFinishCallback finishCallback) {
-        final MixedTransition mixed = createDefaultMixedTransition(
-                MixedTransition.TYPE_ENTER_PIP_FROM_SPLIT, transition);
+            Transitions.TransitionFinishCallback finishCallback, boolean replacingPip) {
+        int type = replacingPip
+                ? MixedTransition.TYPE_ENTER_PIP_REPLACE_FROM_SPLIT
+                : MixedTransition.TYPE_ENTER_PIP_FROM_SPLIT;
+        final MixedTransition mixed = createDefaultMixedTransition(type, transition);
         mActiveTransitions.add(mixed);
         Transitions.TransitionFinishCallback callback = wct -> {
             mActiveTransitions.remove(mixed);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
index b028bd65b438..0ada74937df4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/DefaultMixedTransition.java
@@ -76,7 +76,12 @@ class DefaultMixedTransition extends DefaultMixedHandler.MixedTransition {
                             info, startTransaction, finishTransaction, finishCallback);
             case TYPE_ENTER_PIP_FROM_SPLIT ->
                     animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
-                            finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
+                            finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler,
+                            /*replacingPip*/ false);
+            case TYPE_ENTER_PIP_REPLACE_FROM_SPLIT ->
+                    animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
+                            finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler,
+                            /*replacingPip*/ true);
             case TYPE_KEYGUARD ->
                     animateKeyguard(this, info, startTransaction, finishTransaction, finishCallback,
                             mKeyguardHandler, mPipHandler);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
index ffc0b76b131d..e8b01b5880fb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/MixedTransitionHelper.java
@@ -23,11 +23,15 @@ import static android.window.TransitionInfo.FLAG_IS_WALLPAPER;
 import static com.android.wm.shell.common.split.SplitScreenConstants.FLAG_IS_DIVIDER_BAR;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_UNDEFINED;
 import static com.android.wm.shell.pip.PipAnimationController.ANIM_TYPE_ALPHA;
+import static com.android.wm.shell.shared.TransitionUtil.isOpeningMode;
+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_CHILD_TASK_ENTER_PIP;
 import static com.android.wm.shell.transition.DefaultMixedHandler.subCopy;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.view.SurfaceControl;
 import android.window.TransitionInfo;
 
@@ -45,7 +49,8 @@ public class MixedTransitionHelper {
             @NonNull SurfaceControl.Transaction finishTransaction,
             @NonNull Transitions.TransitionFinishCallback finishCallback,
             @NonNull Transitions player, @NonNull MixedTransitionHandler mixedHandler,
-            @NonNull PipTransitionController pipHandler, @NonNull StageCoordinator splitHandler) {
+            @NonNull PipTransitionController pipHandler, @NonNull StageCoordinator splitHandler,
+            boolean replacingPip) {
         ProtoLog.v(ShellProtoLogGroup.WM_SHELL_TRANSITIONS, " Animating a mixed transition for "
                 + "entering PIP while Split-Screen is foreground.");
         TransitionInfo.Change pipChange = null;
@@ -99,7 +104,7 @@ public class MixedTransitionHelper {
             // we need a separate one to send over to launcher.
             SurfaceControl.Transaction otherStartT = new SurfaceControl.Transaction();
             @SplitScreen.StageType int topStageToKeep = STAGE_TYPE_UNDEFINED;
-            if (splitHandler.isSplitScreenVisible()) {
+            if (splitHandler.isSplitScreenVisible() && !replacingPip) {
                 // 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) {
@@ -115,11 +120,12 @@ public class MixedTransitionHelper {
                         break;
                     }
                 }
+
+                // Let split update internal state for dismiss.
+                splitHandler.prepareDismissAnimation(topStageToKeep,
+                        EXIT_REASON_CHILD_TASK_ENTER_PIP, everythingElse, otherStartT,
+                        finishTransaction);
             }
-            // Let split update internal state for dismiss.
-            splitHandler.prepareDismissAnimation(topStageToKeep,
-                    EXIT_REASON_CHILD_TASK_ENTER_PIP, everythingElse, otherStartT,
-                    finishTransaction);
 
             // We are trying to accommodate launcher's close animation which can't handle the
             // divider-bar, so if split-handler is closing the divider-bar, just hide it and
@@ -152,6 +158,44 @@ public class MixedTransitionHelper {
         return true;
     }
 
+    /**
+     * Check to see if we're only closing split to enter pip or if we're replacing pip with
+     * another task. If we are replacing, this will return the change for the task we are replacing
+     * pip with
+     *
+     * @param info Any number of changes
+     * @param pipChange TransitionInfo.Change indicating the task that is being pipped
+     * @param splitMainStageRootId MainStage's rootTaskInfo's id
+     * @param splitSideStageRootId SideStage's rootTaskInfo's id
+     * @param lastPipSplitStage The last stage that {@param pipChange} was in
+     * @return The change from {@param info} that is replacing the {@param pipChange}, {@code null}
+     *         otherwise
+     */
+    @Nullable
+    public static TransitionInfo.Change getPipReplacingChange(TransitionInfo info,
+            TransitionInfo.Change pipChange, int splitMainStageRootId, int splitSideStageRootId,
+            @SplitScreen.StageType int lastPipSplitStage) {
+        int lastPipParentTask = -1;
+        if (lastPipSplitStage == STAGE_TYPE_MAIN) {
+            lastPipParentTask = splitMainStageRootId;
+        } else if (lastPipSplitStage == STAGE_TYPE_SIDE) {
+            lastPipParentTask = splitSideStageRootId;
+        }
+
+        for (int i = info.getChanges().size() - 1; i >= 0; --i) {
+            TransitionInfo.Change change = info.getChanges().get(i);
+            if (change == pipChange || !isOpeningMode(change.getMode())) {
+                // Ignore the change/task that's going into Pip or not opening
+                continue;
+            }
+
+            if (change.getTaskInfo().parentTaskId == lastPipParentTask) {
+                return change;
+            }
+        }
+        return null;
+    }
+
     private static boolean isHomeOpening(@NonNull TransitionInfo.Change change) {
         return change.getTaskInfo() != null
                 && change.getTaskInfo().getActivityType() == ACTIVITY_TYPE_HOME;
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
index d6e64cfaf4d5..9fc6702562bb 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/RecentsMixedTransition.java
@@ -142,7 +142,8 @@ class RecentsMixedTransition extends DefaultMixedHandler.MixedTransition {
                     && mSplitHandler.getSplitItemPosition(change.getLastParent())
                     != SPLIT_POSITION_UNDEFINED) {
                 return animateEnterPipFromSplit(this, info, startTransaction, finishTransaction,
-                        finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler);
+                        finishCallback, mPlayer, mMixedHandler, mPipHandler, mSplitHandler,
+                        /*replacingPip*/ false);
             }
         }