Fix split dismissed after click pip to split button

When pip app could show when locked(e.g. Maps) and user click
expand to split button, the keyguard occluded callback usually
come before onStageVisibilityChanged then cause split exit by
keyguard occluded

Fix it by add more condition check, we should do it only when
device waking up, keyguard occluded and only the stage included
show when locked app is visible then exit split to show that app
full screen.

Bug: 202740657
Fix: 204527184
Test: manual
Test: pass existing tests
Change-Id: Id63e8dee42337b5dd4f767a866a965cb97761cf0

libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreen.java

@@ -87,6 +87,12 @@ public interface SplitScreen {
      */
     void onKeyguardVisibilityChanged(boolean showing);
 
+    /** Called when device waking up finished. */
+    void onFinishedWakingUp();
+
+    /** Called when device going to sleep finished. */
+    void onFinishedGoingToSleep();
+
     /** Get a string representation of a stage type */
     static String stageTypeToString(@StageType int stage) {
         switch (stage) {

libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/SplitScreenController.java

@@ -242,6 +242,14 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
         mStageCoordinator.onKeyguardVisibilityChanged(showing);
     }
 
+    public void onFinishedWakingUp() {
+        mStageCoordinator.onFinishedWakingUp();
+    }
+
+    public void onFinishedGoingToSleep() {
+        mStageCoordinator.onFinishedGoingToSleep();
+    }
+
     public void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
         mStageCoordinator.exitSplitScreenOnHide(exitSplitScreenOnHide);
     }
@@ -501,6 +509,20 @@ public class SplitScreenController implements DragAndDropPolicy.Starter,
                 SplitScreenController.this.onKeyguardVisibilityChanged(showing);
             });
         }
+
+        @Override
+        public void onFinishedWakingUp() {
+            mMainExecutor.execute(() -> {
+                SplitScreenController.this.onFinishedWakingUp();
+            });
+        }
+
+        @Override
+        public void onFinishedGoingToSleep() {
+            mMainExecutor.execute(() -> {
+                SplitScreenController.this.onFinishedGoingToSleep();
+            });
+        }
     }
 
     /**

libs/WindowManager/Shell/src/com/android/wm/shell/splitscreen/StageCoordinator.java

@@ -155,6 +155,7 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
     private boolean mShouldUpdateRecents;
     private boolean mExitSplitScreenOnHide;
     private boolean mKeyguardOccluded;
+    private boolean mDeviceSleep;
 
     @SplitScreen.StageType
     private int mDismissTop = NO_DISMISS;
@@ -537,6 +538,17 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
+    void onFinishedWakingUp() {
+        if (mMainStage.isActive()) {
+            exitSplitScreenIfKeyguardOccluded();
+        }
+        mDeviceSleep = false;
+    }
+
+    void onFinishedGoingToSleep() {
+        mDeviceSleep = true;
+    }
+
     void exitSplitScreenOnHide(boolean exitSplitScreenOnHide) {
         mExitSplitScreenOnHide = exitSplitScreenOnHide;
     }
@@ -565,6 +577,19 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         applyExitSplitScreen(childrenToTop, wct, exitReason);
     }
 
+    private void exitSplitScreenIfKeyguardOccluded() {
+        final boolean mainStageVisible = mMainStageListener.mVisible;
+        final boolean oneStageVisible = mainStageVisible ^ mSideStageListener.mVisible;
+        if (mDeviceSleep && mKeyguardOccluded && oneStageVisible) {
+            // Only the stages include show-when-locked activity is visible while keyguard occluded.
+            // Dismiss split because there's show-when-locked activity showing on top of keyguard.
+            // Also make sure the task contains show-when-locked activity remains on top after split
+            // dismissed.
+            final StageTaskListener toTop = mainStageVisible ? mMainStage : mSideStage;
+            exitSplitScreen(toTop, EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
+        }
+    }
+
     private void applyExitSplitScreen(StageTaskListener childrenToTop,
             WindowContainerTransaction wct, @ExitReason int exitReason) {
         mRecentTasks.ifPresent(recentTasks -> {
@@ -780,14 +805,8 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
             // like the cases keyguard showing or screen off.
                 exitSplitScreen(null /* childrenToTop */, EXIT_REASON_RETURN_HOME);
             }
-        } else if (mKeyguardOccluded) {
-            // At least one of the stages is visible while keyguard occluded. Dismiss split because
-            // there's show-when-locked activity showing on top of keyguard. Also make sure the
-            // task contains show-when-locked activity remains on top after split dismissed.
-            final StageTaskListener toTop =
-                    mainStageVisible ? mMainStage : (sideStageVisible ? mSideStage : null);
-            exitSplitScreen(toTop, EXIT_REASON_SCREEN_LOCKED_SHOW_ON_TOP);
         }
+        exitSplitScreenIfKeyguardOccluded();
 
         mSyncQueue.runInSync(t -> {
             // Same above, we only set root tasks and divider leash visibility when both stage

packages/SystemUI/src/com/android/systemui/wmshell/WMShell.java

@@ -260,6 +260,18 @@ public final class WMShell extends SystemUI
             }
         };
         mKeyguardUpdateMonitor.registerCallback(mSplitScreenKeyguardCallback);
+
+        mWakefulnessLifecycle.addObserver(new WakefulnessLifecycle.Observer() {
+            @Override
+            public void onFinishedWakingUp() {
+                splitScreen.onFinishedWakingUp();
+            }
+
+            @Override
+            public void onFinishedGoingToSleep() {
+                splitScreen.onFinishedGoingToSleep();
+            }
+        });
     }
 
     @VisibleForTesting