Dismiss staged split after folded the device

Bug: 188010345
Bug: 184795211
Test: atest WMShellUnitTests
Test: observed staged split will be dismissed after device folded

Change-Id: I00ef5ac179fc506f883f8ca519a901273427e595

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

@@ -44,6 +44,7 @@ import android.annotation.Nullable;
 import android.app.ActivityManager;
 import android.content.Context;
 import android.graphics.Rect;
+import android.hardware.devicestate.DeviceStateManager;
 import android.os.Bundle;
 import android.os.IBinder;
 import android.util.Log;
@@ -157,6 +158,10 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
                 mSurfaceSession);
         mDisplayImeController = displayImeController;
         mRootTDAOrganizer.registerListener(displayId, this);
+        final DeviceStateManager deviceStateManager =
+                mContext.getSystemService(DeviceStateManager.class);
+        deviceStateManager.registerCallback(taskOrganizer.getExecutor(),
+                new DeviceStateManager.FoldStateListener(mContext, this::onFoldedStateChanged));
         mSplitTransitions = new SplitScreenTransitions(transactionPool, transitions,
                 mOnTransitionAnimationComplete);
         transitions.addHandler(this);
@@ -585,6 +590,12 @@ class StageCoordinator implements SplitLayout.SplitLayoutHandler,
         }
     }
 
+    private void onFoldedStateChanged(boolean folded) {
+        if (folded && mMainStage.isActive()) {
+            exitSplitScreen();
+        }
+    }
+
     private Rect getSideStageBounds() {
         return mSideStagePosition == SPLIT_POSITION_TOP_OR_LEFT
                 ? mSplitLayout.getBounds1() : mSplitLayout.getBounds2();