Collect tasks when windowing mode of parent display changes

When the windowing mode of a display is updated, the tasks in it
can change, but it's not collected now.

Bug: 225295924
Test: manual
Change-Id: Ief58575e99f51afe6ecfe79dd270d67863cfa0ea
diff --git a/services/core/java/com/android/server/wm/DisplayRotation.java b/services/core/java/com/android/server/wm/DisplayRotation.java
index b9d8319..3d91921 100644
--- a/services/core/java/com/android/server/wm/DisplayRotation.java
+++ b/services/core/java/com/android/server/wm/DisplayRotation.java
@@ -549,7 +549,7 @@
             // Go through all tasks and collect them before the rotation
             // TODO(shell-transitions): move collect() to onConfigurationChange once wallpaper
             //       handling is synchronized.
-            mDisplayContent.mTransitionController.collectForDisplayChange(mDisplayContent,
+            mDisplayContent.mTransitionController.collectForDisplayAreaChange(mDisplayContent,
                     null /* use collecting transition */);
         }
         mService.mAtmService.deferWindowLayout();
diff --git a/services/core/java/com/android/server/wm/PhysicalDisplaySwitchTransitionLauncher.java b/services/core/java/com/android/server/wm/PhysicalDisplaySwitchTransitionLauncher.java
index d209f08..64749cf 100644
--- a/services/core/java/com/android/server/wm/PhysicalDisplaySwitchTransitionLauncher.java
+++ b/services/core/java/com/android/server/wm/PhysicalDisplaySwitchTransitionLauncher.java
@@ -101,7 +101,7 @@
 
         if (t != null) {
             mDisplayContent.mAtmService.startLaunchPowerMode(POWER_MODE_REASON_CHANGE_DISPLAY);
-            mTransitionController.collectForDisplayChange(mDisplayContent, t);
+            mTransitionController.collectForDisplayAreaChange(mDisplayContent, t);
             mTransition = t;
         }
     }
diff --git a/services/core/java/com/android/server/wm/TransitionController.java b/services/core/java/com/android/server/wm/TransitionController.java
index 88572a9..a02be25 100644
--- a/services/core/java/com/android/server/wm/TransitionController.java
+++ b/services/core/java/com/android/server/wm/TransitionController.java
@@ -460,23 +460,26 @@
      * Collects the window containers which need to be synced with the changing display (e.g.
      * rotating) to the given transition or the current collecting transition.
      */
-    void collectForDisplayChange(@NonNull DisplayContent dc, @Nullable Transition incoming) {
+    void collectForDisplayAreaChange(@NonNull DisplayArea<?> wc, @Nullable Transition incoming) {
         if (incoming == null) incoming = mCollectingTransition;
         if (incoming == null) return;
         final Transition transition = incoming;
         // Collect all visible tasks.
-        dc.forAllLeafTasks(task -> {
+        wc.forAllLeafTasks(task -> {
             if (task.isVisible()) {
                 transition.collect(task);
             }
         }, true /* traverseTopToBottom */);
         // Collect all visible non-app windows which need to be drawn before the animation starts.
-        dc.forAllWindows(w -> {
-            if (w.mActivityRecord == null && w.isVisible() && !isCollecting(w.mToken)
-                    && dc.shouldSyncRotationChange(w)) {
-                transition.collect(w.mToken);
-            }
-        }, true /* traverseTopToBottom */);
+        final DisplayContent dc = wc.asDisplayContent();
+        if (dc != null) {
+            wc.forAllWindows(w -> {
+                if (w.mActivityRecord == null && w.isVisible() && !isCollecting(w.mToken)
+                        && dc.shouldSyncRotationChange(w)) {
+                    transition.collect(w.mToken);
+                }
+            }, true /* traverseTopToBottom */);
+        }
     }
 
     /** @see Transition#mStatusBarTransitionDelay */
diff --git a/services/core/java/com/android/server/wm/WindowOrganizerController.java b/services/core/java/com/android/server/wm/WindowOrganizerController.java
index 6bb5ece..ee64354 100644
--- a/services/core/java/com/android/server/wm/WindowOrganizerController.java
+++ b/services/core/java/com/android/server/wm/WindowOrganizerController.java
@@ -397,7 +397,7 @@
                     // Go through all tasks and collect them before the rotation
                     // TODO(shell-transitions): move collect() to onConfigurationChange once
                     //       wallpaper handling is synchronized.
-                    dc.mTransitionController.collectForDisplayChange(dc, transition);
+                    dc.mTransitionController.collectForDisplayAreaChange(dc, transition);
                     dc.sendNewConfiguration();
                     effects |= TRANSACT_EFFECTS_LIFECYCLE;
                 }
@@ -421,6 +421,15 @@
                     addToSyncSet(syncId, wc);
                 }
                 if (transition != null) transition.collect(wc);
+                final DisplayArea da = wc.asDisplayArea();
+                // Only check DisplayArea here as a similar thing is done for DisplayContent above.
+                if (da != null && wc.asDisplayContent() == null
+                        && entry.getValue().getWindowingMode() != da.getWindowingMode()) {
+                    // Go through all tasks and collect them before changing the windowing mode of a
+                    // display-level container.
+                    // TODO(shell-transitions): handle this more elegantly.
+                    da.mTransitionController.collectForDisplayAreaChange(da, transition);
+                }
 
                 if ((entry.getValue().getChangeMask()
                         & WindowContainerTransaction.Change.CHANGE_FORCE_NO_PIP) != 0) {