[CD] Fix predictive back system anims partially run on wrong display

Bug: 382774299
Test: Manual, i.e. verified that predictive back system animations do
      not add background/scrim windows on the wrong display.
Flag: com.android.window.flags.enable_multidisplay_trackpad_back_gesture
Change-Id: I833a7d93c5fe30287cb8ec2bf0b4eea803a659a1

3 files changed, 28 insertions, 8 deletions

diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationBackground.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationBackground.java
index 26f7b360e0fc..98cae5ae9296 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationBackground.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/BackAnimationBackground.java
@@ -19,6 +19,8 @@ package com.android.wm.shell.back;
 import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowInsetsController.APPEARANCE_LIGHT_STATUS_BARS;
 
+import static com.android.window.flags.Flags.enableMultidisplayTrackpadBackGesture;
+
 import android.annotation.NonNull;
 import android.annotation.Nullable;
 import android.graphics.Color;
@@ -59,9 +61,9 @@ public class BackAnimationBackground {
      * @param statusbarHeight The height of the statusbar (in px).
      */
     public void ensureBackground(Rect startRect, int color,
-            @NonNull SurfaceControl.Transaction transaction, int statusbarHeight) {
+            @NonNull SurfaceControl.Transaction transaction, int statusbarHeight, int displayId) {
         ensureBackground(startRect, color, transaction, statusbarHeight,
-                null /* cropBounds */, 0 /* cornerRadius */);
+                null /* cropBounds */, 0 /* cornerRadius */, displayId);
     }
 
     /**
@@ -76,7 +78,7 @@ public class BackAnimationBackground {
      */
     public void ensureBackground(Rect startRect, int color,
             @NonNull SurfaceControl.Transaction transaction, int statusbarHeight,
-            @Nullable Rect cropBounds, float cornerRadius) {
+            @Nullable Rect cropBounds, float cornerRadius, int displayId) {
         if (mBackgroundSurface != null) {
             return;
         }
@@ -91,7 +93,11 @@ public class BackAnimationBackground {
                 .setCallsite("BackAnimationBackground")
                 .setColorLayer();
 
-        mRootTaskDisplayAreaOrganizer.attachToDisplayArea(DEFAULT_DISPLAY, colorLayerBuilder);
+        if (enableMultidisplayTrackpadBackGesture()) {
+            mRootTaskDisplayAreaOrganizer.attachToDisplayArea(displayId, colorLayerBuilder);
+        } else {
+            mRootTaskDisplayAreaOrganizer.attachToDisplayArea(DEFAULT_DISPLAY, colorLayerBuilder);
+        }
         mBackgroundSurface = colorLayerBuilder.build();
         transaction.setColor(mBackgroundSurface, colorComponents)
                 .setLayer(mBackgroundSurface, BACKGROUND_LAYER)
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
index 7e5a82e640cc..6c41c975cf4d 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossActivityBackAnimation.kt
@@ -52,6 +52,7 @@ import com.android.internal.jank.Cuj
 import com.android.internal.policy.ScreenDecorationsUtils
 import com.android.internal.policy.SystemBarUtils
 import com.android.internal.protolog.ProtoLog
+import com.android.window.flags.Flags.enableMultidisplayTrackpadBackGesture
 import com.android.window.flags.Flags.predictiveBackTimestampApi
 import com.android.wm.shell.R
 import com.android.wm.shell.RootTaskDisplayAreaOrganizer
@@ -210,7 +211,8 @@ abstract class CrossActivityBackAnimation(
                 statusbarHeight,
                 if (closingTarget!!.windowConfiguration.tasksAreFloating())
                     closingTarget!!.localBounds else null,
-                cornerRadius
+                cornerRadius,
+                closingTarget!!.taskInfo.getDisplayId()
         )
         ensureScrimLayer()
         if (isLetterboxed && enteringHasSameLetterbox) {
@@ -409,7 +411,12 @@ abstract class CrossActivityBackAnimation(
                 .setOpaque(false)
                 .setHidden(false)
 
-        rootTaskDisplayAreaOrganizer.attachToDisplayArea(Display.DEFAULT_DISPLAY, scrimBuilder)
+        if (enableMultidisplayTrackpadBackGesture()) {
+            rootTaskDisplayAreaOrganizer.attachToDisplayArea(
+                closingTarget!!.taskInfo.getDisplayId(), scrimBuilder)
+        } else {
+            rootTaskDisplayAreaOrganizer.attachToDisplayArea(Display.DEFAULT_DISPLAY, scrimBuilder)
+        }
         scrimLayer = scrimBuilder.build()
         val colorComponents = floatArrayOf(0f, 0f, 0f)
         maxScrimAlpha = if (isDarkTheme) MAX_SCRIM_ALPHA_DARK else MAX_SCRIM_ALPHA_LIGHT
@@ -473,7 +480,13 @@ abstract class CrossActivityBackAnimation(
                 .setOpaque(true)
                 .setHidden(false)
 
-        rootTaskDisplayAreaOrganizer.attachToDisplayArea(Display.DEFAULT_DISPLAY, letterboxBuilder)
+        if (enableMultidisplayTrackpadBackGesture()) {
+            rootTaskDisplayAreaOrganizer.attachToDisplayArea(
+                closingTarget!!.taskInfo.getDisplayId(), letterboxBuilder)
+        } else {
+            rootTaskDisplayAreaOrganizer.attachToDisplayArea(
+                Display.DEFAULT_DISPLAY, letterboxBuilder)
+        }
         val layer = letterboxBuilder.build()
         val colorComponents =
             floatArrayOf(
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
index f48b3ffcd598..f5b0e359e019 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/back/CrossTaskBackAnimation.java
@@ -175,7 +175,8 @@ public class CrossTaskBackAnimation extends ShellBackAnimation {
 
         // Draw background.
         mBackground.ensureBackground(mClosingTarget.windowConfiguration.getBounds(),
-                BACKGROUNDCOLOR, mTransaction, mStatusbarHeight);
+                BACKGROUNDCOLOR, mTransaction, mStatusbarHeight,
+                mClosingTarget.taskInfo.getDisplayId());
         mInterWindowMargin = mContext.getResources()
                 .getDimension(R.dimen.cross_task_back_inter_window_margin);
         mVerticalMargin = mContext.getResources()