Handle all stage root tasks dismissed for flex split

* Access stages via stageOperator instead of main/side directly
* Currently we always deactivate main even if that stage isn't to top,
with flex let's try deactivating everything if we aren't keeping any
stage on top

Bug: 333270112
Flag: com.android.wm.shell.enable_flexible_split
Test: Swiping up to dismiss split pair in overview doesn't crash
Change-Id: Ic5281bc3ed2b9f70506e424208cccb9b206c1582

1 files changed, 24 insertions, 9 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 45ecfa95b494..096a8da56765 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
@@ -494,8 +494,15 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         if (enableFlexibleSplit()) {
             StageTaskListener stageToDeactivate = mStageOrderOperator.getAllStages().stream()
                     .filter(stage -> stage.getId() == stageToTop)
-                    .findFirst().orElseThrow();
-            stageToDeactivate.deactivate(wct, true /*toTop*/);
+                    .findFirst().orElse(null);
+            if (stageToDeactivate != null) {
+                stageToDeactivate.deactivate(wct, true /*toTop*/);
+            } else {
+                // If no one stage is meant to go to the top, deactivate all stages to move any
+                // child tasks out from under their respective stage root tasks.
+                mStageOrderOperator.getAllStages().forEach(stage ->
+                        stage.deactivate(wct, false /*reparentTasksToTop*/));
+            }
             mStageOrderOperator.onExitingSplit();
         } else {
             mMainStage.deactivate(wct, stageToTop == STAGE_TYPE_MAIN);
@@ -1835,13 +1842,21 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
         if (present) {
             updateRecentTasksSplitPair();
-        } else if (mMainStage.getChildCount() == 0 && mSideStage.getChildCount() == 0) {
-            mRecentTasks.ifPresent(recentTasks -> {
-                // remove the split pair mapping from recentTasks, and disable further updates
-                // to splits in the recents until we enter split again.
-                recentTasks.removeSplitPair(taskId);
-            });
-            dismissSplitScreen(-1, EXIT_REASON_ROOT_TASK_VANISHED);
+        } else {
+            // TODO (b/349828130): Test b/333270112 for flex split (launch adjacent for flex
+            //  currently not working)
+            boolean allRootsEmpty = enableFlexibleSplit()
+                    ? runForActiveStagesAllMatch(stageTaskListener ->
+                        stageTaskListener.getChildCount() == 0)
+                    : mMainStage.getChildCount() == 0 && mSideStage.getChildCount() == 0;
+            if (allRootsEmpty) {
+                mRecentTasks.ifPresent(recentTasks -> {
+                    // remove the split pair mapping from recentTasks, and disable further updates
+                    // to splits in the recents until we enter split again.
+                    recentTasks.removeSplitPair(taskId);
+                });
+                dismissSplitScreen(INVALID_TASK_ID, EXIT_REASON_ROOT_TASK_VANISHED);
+            }
         }
 
         for (int i = mListeners.size() - 1; i >= 0; --i) {