Handle split screen accordingly after recent transition finished

When a recent transition finished in split screen, make sure to:
1. dismiss the split screen when going to home or launching another
   fullscreen app
2. restore the split screen when returning to a split pair
3. evict privious splitting tasks when entering to a different split
   pair

Bug: 228965939
Fix: 233612774
Test: atest WMShellUnitTests
Test: enable shell transition, trigger split screen and enter overview,
      verified it'll dismiss the split when swipe up to home or switch
      to another fullscreen app
Test: enable shell transition, trigger split screen and enter overview,
      verified it'll restore the divider bar when it returning to the
      original split pair or switching to a different split pair
Change-Id: Ifbfe3baf0abb33f9970fc1a0d36f301cf3b6993a
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
index f8221d5..f7057d4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenTransitions.java
@@ -68,6 +68,7 @@
     private IBinder mAnimatingTransition = null;
     OneShotRemoteHandler mPendingRemoteHandler = null;
     private OneShotRemoteHandler mActiveRemoteHandler = null;
+    private boolean mEnterTransitionMerged;
 
     private final Transitions.TransitionFinishCallback mRemoteFinishCB = this::onFinish;
 
@@ -229,6 +230,18 @@
         }
     }
 
+    void onTransitionMerged(@NonNull IBinder transition) {
+        // Once a pending enter transition got merged, make sure to append the reset of finishing
+        // operations to the finish transition.
+        if (transition == mPendingEnter) {
+            mFinishTransaction = mTransactionPool.acquire();
+            mStageCoordinator.finishEnterSplitScreen(mFinishTransaction);
+            mPendingEnter = null;
+            mPendingRemoteHandler = null;
+            mEnterTransitionMerged = true;
+        }
+    }
+
     void onFinish(WindowContainerTransaction wct, WindowContainerTransactionCallback wctCB) {
         if (!mAnimations.isEmpty()) return;
         if (mAnimatingTransition == mPendingEnter) {
@@ -238,18 +251,16 @@
             mPendingDismiss = null;
         }
         if (mAnimatingTransition == mPendingRecent) {
-            // If the clean-up wct is null when finishing recent transition, it indicates it's
-            // returning to home and thus no need to reorder tasks.
-            final boolean returnToHome = wct == null;
-            if (returnToHome) {
-                wct = new WindowContainerTransaction();
+            if (!mEnterTransitionMerged) {
+                if (wct == null) wct = new WindowContainerTransaction();
+                mStageCoordinator.onRecentTransitionFinished(wct, mFinishTransaction);
             }
-            mStageCoordinator.onRecentTransitionFinished(returnToHome, wct, mFinishTransaction);
             mPendingRecent = null;
         }
         mPendingRemoteHandler = null;
         mActiveRemoteHandler = null;
         mAnimatingTransition = null;
+        mEnterTransitionMerged = false;
 
         mOnFinish.run();
         if (mFinishTransaction != null) {
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 5671fec..ddc01b6 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
@@ -30,6 +30,7 @@
 import static android.view.WindowManager.transitTypeToString;
 import static android.window.TransitionInfo.FLAG_FIRST_CUSTOM;
 import static android.window.TransitionInfo.FLAG_IS_DISPLAY;
+import static android.window.WindowContainerTransaction.HierarchyOp.HIERARCHY_OP_TYPE_REORDER;
 
 import static com.android.wm.shell.common.split.SplitLayout.PARALLAX_ALIGN_CENTER;
 import static com.android.wm.shell.common.split.SplitScreenConstants.SPLIT_POSITION_BOTTOM_OR_RIGHT;
@@ -369,10 +370,15 @@
         sideOptions = sideOptions != null ? sideOptions : new Bundle();
         setSideStagePosition(sidePosition, wct);
 
+        if (mMainStage.isActive()) {
+            mMainStage.evictAllChildren(wct);
+            mSideStage.evictAllChildren(wct);
+        } else {
+            // Build a request WCT that will launch both apps such that task 0 is on the main stage
+            // while task 1 is on the side stage.
+            mMainStage.activate(wct, false /* reparent */);
+        }
         mSplitLayout.setDivideRatio(splitRatio);
-        // Build a request WCT that will launch both apps such that task 0 is on the main stage
-        // while task 1 is on the side stage.
-        mMainStage.activate(wct, false /* reparent */);
         updateWindowBounds(mSplitLayout, wct);
         wct.reorder(mRootTaskInfo.token, true);
 
@@ -1503,16 +1509,7 @@
 
     @Override
     public void onTransitionMerged(@NonNull IBinder transition) {
-        // Once the pending enter transition got merged, make sure to bring divider bar visible and
-        // clear the pending transition from cache to prevent mess-up the following state.
-        if (transition == mSplitTransitions.mPendingEnter) {
-            final SurfaceControl.Transaction t = mTransactionPool.acquire();
-            finishEnterSplitScreen(t);
-            mSplitTransitions.mPendingEnter = null;
-            mSplitTransitions.mPendingRemoteHandler = null;
-            t.apply();
-            mTransactionPool.release(t);
-        }
+        mSplitTransitions.onTransitionMerged(transition);
     }
 
     @Override
@@ -1713,8 +1710,6 @@
             logExitToStage(dismissReason, toStage == STAGE_TYPE_MAIN);
         }
 
-        addDividerBarToTransition(info, t, false /* show */);
-
         // Hide divider and dim layer on transition finished.
         setDividerVisibility(false, finishT);
         finishT.hide(mMainStage.mDimLayer);
@@ -1736,6 +1731,7 @@
             return false;
         }
 
+        addDividerBarToTransition(info, t, false /* show */);
         return true;
     }
 
@@ -1745,26 +1741,26 @@
         return true;
     }
 
-    void onRecentTransitionFinished(boolean returnToHome, WindowContainerTransaction wct,
+    void onRecentTransitionFinished(WindowContainerTransaction wct,
             SurfaceControl.Transaction finishT) {
-        // Exclude the case that the split screen has been dismissed already.
-        if (!mMainStage.isActive()) {
-            // The latest split dismissing transition might be a no-op transition and thus won't
-            // callback startAnimation, update split visibility here to cover this kind of no-op
-            // transition case.
-            setSplitsVisible(false);
-            return;
+        // Check if the recent transition is finished by returning to the current split so we can
+        // restore the divider bar.
+        for (int i = 0; i < wct.getHierarchyOps().size(); ++i) {
+            final WindowContainerTransaction.HierarchyOp op = wct.getHierarchyOps().get(i);
+            final IBinder container = op.getContainer();
+            if (op.getType() == HIERARCHY_OP_TYPE_REORDER && op.getToTop()
+                    && (mMainStage.containsContainer(container)
+                    || mSideStage.containsContainer(container))) {
+                setDividerVisibility(true, finishT);
+                return;
+            }
         }
 
-        if (returnToHome) {
-            // When returning to home from recent apps, the splitting tasks are already hidden, so
-            // append the reset of dismissing operations into the clean-up wct.
-            prepareExitSplitScreen(STAGE_TYPE_UNDEFINED, wct);
-            setSplitsVisible(false);
-            logExit(EXIT_REASON_RETURN_HOME);
-        } else {
-            setDividerVisibility(true, finishT);
-        }
+        // Dismiss the split screen is it's not returning to split.
+        prepareExitSplitScreen(STAGE_TYPE_UNDEFINED, wct);
+        setSplitsVisible(false);
+        setDividerVisibility(false, finishT);
+        logExit(EXIT_REASON_UNKNOWN);
     }
 
     private void addDividerBarToTransition(@NonNull TransitionInfo info,
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 4b12eb8..1ec228c5 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
@@ -31,6 +31,7 @@
 import android.content.Context;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.os.IBinder;
 import android.util.SparseArray;
 import android.view.SurfaceControl;
 import android.view.SurfaceSession;
@@ -47,6 +48,7 @@
 import com.android.wm.shell.splitscreen.SplitScreen.StageType;
 
 import java.io.PrintWriter;
+import java.util.function.Predicate;
 
 /**
  * Base class that handle common task org. related for split-screen stages.
@@ -119,63 +121,53 @@
     }
 
     boolean containsToken(WindowContainerToken token) {
-        if (token.equals(mRootTaskInfo.token)) {
-            return true;
-        }
+        return contains(t -> t.token.equals(token));
+    }
 
-        for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
-            if (token.equals(mChildrenTaskInfo.valueAt(i).token)) {
-                return true;
-            }
-        }
-
-        return false;
+    boolean containsContainer(IBinder binder) {
+        return contains(t -> t.token.asBinder() == binder);
     }
 
     /**
      * Returns the top visible child task's id.
      */
     int getTopVisibleChildTaskId() {
-        for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
-            final ActivityManager.RunningTaskInfo info = mChildrenTaskInfo.valueAt(i);
-            if (info.isVisible) {
-                return info.taskId;
-            }
-        }
-        return INVALID_TASK_ID;
+        final ActivityManager.RunningTaskInfo taskInfo = getChildTaskInfo(t -> t.isVisible);
+        return taskInfo != null ? taskInfo.taskId : INVALID_TASK_ID;
     }
 
     /**
      * Returns the top activity uid for the top child task.
      */
     int getTopChildTaskUid() {
-        for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
-            final ActivityManager.RunningTaskInfo info = mChildrenTaskInfo.valueAt(i);
-            if (info.topActivityInfo == null) {
-                continue;
-            }
-            return info.topActivityInfo.applicationInfo.uid;
-        }
-        return 0;
+        final ActivityManager.RunningTaskInfo taskInfo =
+                getChildTaskInfo(t -> t.topActivityInfo != null);
+        return taskInfo != null ? taskInfo.topActivityInfo.applicationInfo.uid : 0;
     }
 
     /** @return {@code true} if this listener contains the currently focused task. */
     boolean isFocused() {
-        if (mRootTaskInfo == null) {
-            return false;
-        }
+        return contains(t -> t.isFocused);
+    }
 
-        if (mRootTaskInfo.isFocused) {
+    private boolean contains(Predicate<ActivityManager.RunningTaskInfo> predicate) {
+        if (mRootTaskInfo != null && predicate.test(mRootTaskInfo)) {
             return true;
         }
 
+        return getChildTaskInfo(predicate) != null;
+    }
+
+    @Nullable
+    private ActivityManager.RunningTaskInfo getChildTaskInfo(
+            Predicate<ActivityManager.RunningTaskInfo> predicate) {
         for (int i = mChildrenTaskInfo.size() - 1; i >= 0; --i) {
-            if (mChildrenTaskInfo.valueAt(i).isFocused) {
-                return true;
+            final ActivityManager.RunningTaskInfo taskInfo = mChildrenTaskInfo.valueAt(i);
+            if (predicate.test(taskInfo)) {
+                return taskInfo;
             }
         }
-
-        return false;
+        return null;
     }
 
     @Override