Prevent split on lockscreen checks directly using stage types

* Prevent directly using STAGE_TYPE_* ids. This will also help
with stages that are not starting with 0 and 1 for flex split.

Bug: 349828130
Test: Apps break on lock screen when expected
Flag: EXEMPT bug fix
Change-Id: I96b59ce1e7548625eccdbd01d84be1ab0117ea40

1 files changed, 24 insertions, 8 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 511e426cc681..9efbf6f770b1 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
@@ -109,6 +109,7 @@ import android.util.ArraySet;
 import android.util.IntArray;
 import android.util.Log;
 import android.util.Slog;
+import android.util.SparseIntArray;
 import android.view.Choreographer;
 import android.view.IRemoteAnimationFinishedCallback;
 import android.view.IRemoteAnimationRunner;
@@ -3011,11 +3012,18 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             final int transitType = info.getType();
             TransitionInfo.Change pipChange = null;
             int closingSplitTaskId = -1;
-            // This array tracks where we are sending stages (TO_BACK/TO_FRONT) in this transition.
-            // TODO (b/349828130): Update for n apps (needs to handle different indices than 0/1).
-            //  Also make sure having multiple changes per stage (2+ tasks in one stage) is being
-            //  handled properly.
-            int[] stageChanges = new int[2];
+            // This array tracks if we are sending stages TO_BACK/TO_FRONT in this transition.
+            // TODO (b/349828130): Also make sure having multiple changes per stage (2+ tasks in
+            //  one stage) is being handled properly.
+            SparseIntArray stageChanges = new SparseIntArray();
+            if (enableFlexibleSplit()) {
+                mStageOrderOperator.getActiveStages()
+                        .forEach(stage -> stageChanges.put(stage.getId(), -1));
+            } else {
+                stageChanges.put(STAGE_TYPE_MAIN, -1);
+                stageChanges.put(STAGE_TYPE_SIDE, -1);
+            }
+
 
             for (int iC = 0; iC < info.getChanges().size(); ++iC) {
                 final TransitionInfo.Change change = info.getChanges().get(iC);
@@ -3095,7 +3103,7 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                             || change.getMode() == TRANSIT_TO_FRONT)
                             && (stageOfTaskId == STAGE_TYPE_MAIN
                             || stageOfTaskId == STAGE_TYPE_SIDE)) {
-                        stageChanges[stageOfTaskId] = change.getMode();
+                        stageChanges.put(getStageOfTask(taskId), change.getMode());
                     }
                 }
             }
@@ -3125,8 +3133,16 @@ public class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             // If keyguard is active, check to see if we have all our stages showing. If one stage
             // was moved but not the other (which can happen with SHOW_ABOVE_LOCKED apps), we should
             // break split.
-            if (mKeyguardActive && stageChanges[STAGE_TYPE_MAIN] != stageChanges[STAGE_TYPE_SIDE]) {
-                dismissSplitKeepingLastActiveStage(EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
+            if (mKeyguardActive && stageChanges.size() > 0) {
+                int firstChangeMode = stageChanges.valueAt(0);
+                for (int i = 0; i < stageChanges.size(); i++) {
+                    int changeMode = stageChanges.valueAt(i);
+                    // Compare each changeMode to the first one. If any are different, break split.
+                    if (changeMode != firstChangeMode) {
+                        dismissSplitKeepingLastActiveStage(EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
+                        break;
+                    }
+                }
             }
 
             final ArraySet<StageTaskListener> dismissStages = record.getShouldDismissedStage();