Exit split when fold only if user interacted

If user fold the phone but do not interact and unfold later, we
should keep split status which follow current UX design.

Fix: 288331854
Test: manual
Test: pass existing tests
Change-Id: I9f008c3b51e460fcbd25f20bd45e84ca4e7ae6f1

5 files changed, 64 insertions, 27 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 6961dabbb4b2..693bf6c02a84 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
@@ -219,6 +219,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     private boolean mIsDropEntering;
     private boolean mIsExiting;
     private boolean mIsRootTranslucent;
+    @VisibleForTesting
+    int mTopStageAfterFoldDismiss;
 
     private DefaultMixedHandler mMixedHandler;
     private final Toast mSplitUnsupportedToast;
@@ -1310,6 +1312,24 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             final StageTaskListener toTop = mainStageVisible ? mMainStage : mSideStage;
             exitSplitScreen(toTop, EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
         }
+
+        // Dismiss split if the flag record any side of stages.
+        if (mTopStageAfterFoldDismiss != STAGE_TYPE_UNDEFINED) {
+            if (ENABLE_SHELL_TRANSITIONS) {
+                // Need manually clear here due to this transition might be aborted due to keyguard
+                // on top and lead to no visible change.
+                clearSplitPairedInRecents(EXIT_REASON_DEVICE_FOLDED);
+                final WindowContainerTransaction wct = new WindowContainerTransaction();
+                prepareExitSplitScreen(mTopStageAfterFoldDismiss, wct);
+                mSplitTransitions.startDismissTransition(wct, this,
+                        mTopStageAfterFoldDismiss, EXIT_REASON_DEVICE_FOLDED);
+            } else {
+                exitSplitScreen(
+                        mTopStageAfterFoldDismiss == STAGE_TYPE_MAIN ? mMainStage : mSideStage,
+                        EXIT_REASON_DEVICE_FOLDED);
+            }
+            mTopStageAfterFoldDismiss = STAGE_TYPE_UNDEFINED;
+        }
     }
 
     void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
@@ -1346,15 +1366,8 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         if (!mMainStage.isActive() || mIsExiting) return;
 
         onSplitScreenExit();
+        clearSplitPairedInRecents(exitReason);
 
-        mRecentTasks.ifPresent(recentTasks -> {
-            // Notify recents if we are exiting in a way that breaks the pair, and disable further
-            // updates to splits in the recents until we enter split again
-            if (shouldBreakPairedTaskInRecents(exitReason) && mShouldUpdateRecents) {
-                recentTasks.removeSplitPair(mMainStage.getTopVisibleChildTaskId());
-                recentTasks.removeSplitPair(mSideStage.getTopVisibleChildTaskId());
-            }
-        });
         mShouldUpdateRecents = false;
         mIsDividerRemoteAnimating = false;
         mSplitRequest = null;
@@ -1481,6 +1494,17 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
+    private void clearSplitPairedInRecents(@ExitReason int exitReason) {
+        if (!shouldBreakPairedTaskInRecents(exitReason) || !mShouldUpdateRecents) return;
+
+        mRecentTasks.ifPresent(recentTasks -> {
+            // Notify recents if we are exiting in a way that breaks the pair, and disable further
+            // updates to splits in the recents until we enter split again
+            mMainStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
+            mSideStage.doForAllChildTasks(taskId -> recentTasks.removeSplitPair(taskId));
+        });
+    }
+
     /**
      * Unlike exitSplitScreen, this takes a stagetype vs an actual stage-reference and populates
      * an existing WindowContainerTransaction (rather than applying immediately). This is intended
@@ -2208,7 +2232,11 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
-    void updateSurfaces(SurfaceControl.Transaction transaction) {
+    /**
+     * Update surfaces of the split screen layout based on the current state
+     * @param transaction to write the updates to
+     */
+    public void updateSurfaces(SurfaceControl.Transaction transaction) {
         updateSurfaceBounds(mSplitLayout, transaction, /* applyResizingOffset */ false);
         mSplitLayout.update(transaction);
     }
@@ -2227,26 +2255,18 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
 
     @VisibleForTesting
     void onFoldedStateChanged(boolean folded) {
-        int topStageAfterFoldDismiss = STAGE_TYPE_UNDEFINED;
+        mTopStageAfterFoldDismiss = STAGE_TYPE_UNDEFINED;
         if (!folded) return;
 
-        if (!mMainStage.isActive()) return;
+        if (!isSplitActive() || !isSplitScreenVisible()) return;
 
+        // To avoid split dismiss when user fold the device and unfold to use later, we only
+        // record the flag here and try to dismiss on wakeUp callback to ensure split dismiss
+        // when user interact on phone folded.
         if (mMainStage.isFocused()) {
-            topStageAfterFoldDismiss = STAGE_TYPE_MAIN;
+            mTopStageAfterFoldDismiss = STAGE_TYPE_MAIN;
         } else if (mSideStage.isFocused()) {
-            topStageAfterFoldDismiss = STAGE_TYPE_SIDE;
-        }
-
-        if (ENABLE_SHELL_TRANSITIONS) {
-            final WindowContainerTransaction wct = new WindowContainerTransaction();
-            prepareExitSplitScreen(topStageAfterFoldDismiss, wct);
-            mSplitTransitions.startDismissTransition(wct, this,
-                    topStageAfterFoldDismiss, EXIT_REASON_DEVICE_FOLDED);
-        } else {
-            exitSplitScreen(
-                    topStageAfterFoldDismiss == STAGE_TYPE_MAIN ? mMainStage : mSideStage,
-                    EXIT_REASON_DEVICE_FOLDED);
+            mTopStageAfterFoldDismiss = STAGE_TYPE_SIDE;
         }
     }
 
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 e2c55e4c63ea..af7bf360f036 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
@@ -55,6 +55,7 @@ import com.android.wm.shell.windowdecor.WindowDecorViewModel;
 
 import java.io.PrintWriter;
 import java.util.Optional;
+import java.util.function.Consumer;
 import java.util.function.Predicate;
 
 /**
@@ -347,6 +348,13 @@ class StageTaskListener implements ShellTaskOrganizer.TaskListener {
         wct.reorder(mChildrenTaskInfo.get(taskId).token, onTop /* onTop */);
     }
 
+    void doForAllChildTasks(Consumer<Integer> consumer) {
+        for (int i = mChildrenTaskInfo.size() - 1; i >= 0; i--) {
+            final ActivityManager.RunningTaskInfo taskInfo = mChildrenTaskInfo.valueAt(i);
+            consumer.accept(taskInfo.taskId);
+        }
+    }
+
     /** Collects all the current child tasks and prepares transaction to evict them to display. */
     void evictAllChildren(WindowContainerTransaction wct) {
         for (int i = mChildrenTaskInfo.size() - 1; i >= 0; i--) {
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 d9edde16a863..c5c22ded8b71 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
@@ -704,6 +704,9 @@ public class DefaultMixedHandler implements Transitions.TransitionHandler,
         if (mPipHandler != null) {
             mPipHandler.syncPipSurfaceState(info, startTransaction, finishTransaction);
         }
+        if (mSplitHandler != null && mSplitHandler.isSplitActive()) {
+            mSplitHandler.updateSurfaces(startTransaction);
+        }
         return mUnfoldHandler.startAnimation(
                 mixed.mTransition, info, startTransaction, finishTransaction, finishCB);
     }
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java
index 21994a997be5..bb5d54652460 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/unfold/animation/SplitTaskUnfoldAnimator.java
@@ -270,7 +270,6 @@ public class SplitTaskUnfoldAnimator implements UnfoldTaskAnimator,
     @Override
     public void prepareStartTransaction(Transaction transaction) {
         mUnfoldBackgroundController.ensureBackground(transaction);
-        mSplitScreenController.get().get().updateSplitScreenSurfaces(transaction);
     }
 
     @Override
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 e59d09cd1ee1..cc4db22e772c 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
@@ -347,13 +347,20 @@ public class StageCoordinatorTests extends ShellTestCase {
     }
 
     @Test
-    public void testExitSplitScreenAfterFolded() {
-        when(mMainStage.isActive()).thenReturn(true);
+    public void testExitSplitScreenAfterFoldedAndWakeUp() {
         when(mMainStage.isFocused()).thenReturn(true);
         when(mMainStage.getTopVisibleChildTaskId()).thenReturn(INVALID_TASK_ID);
+        mSideStage.mRootTaskInfo = new TestRunningTaskInfoBuilder().setVisible(true).build();
+        mMainStage.mRootTaskInfo = new TestRunningTaskInfoBuilder().setVisible(true).build();
+        when(mStageCoordinator.isSplitActive()).thenReturn(true);
+        when(mStageCoordinator.isSplitScreenVisible()).thenReturn(true);
 
         mStageCoordinator.onFoldedStateChanged(true);
 
+        assertEquals(mStageCoordinator.mTopStageAfterFoldDismiss, STAGE_TYPE_MAIN);
+
+        mStageCoordinator.onFinishedWakingUp();
+
         if (Transitions.ENABLE_SHELL_TRANSITIONS) {
             verify(mTaskOrganizer).startNewTransition(eq(TRANSIT_SPLIT_DISMISS), notNull());
         } else {