Make TaskAnimationManager per-display

Test: locally tested on Tangor
Flag: com.android.launcher3.enable_overview_on_connected_displays
Bug: 402362875
Change-Id: I6acb166c200e8bd9198e2aa73e506b3c3414f526
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 20dfb10..152630a 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -75,7 +75,6 @@
 constructor(
     private val touchInteractionService: TouchInteractionService,
     private val overviewComponentObserver: OverviewComponentObserver,
-    private val taskAnimationManager: TaskAnimationManager,
     private val dispatcherProvider: DispatcherProvider = ProductionDispatchers,
     private val recentsDisplayModel: RecentsDisplayModel,
     private val focusState: FocusState,
@@ -457,6 +456,14 @@
                 }
             }
 
+        val displayId = gestureState.displayId
+        val taskAnimationManager =
+            recentsDisplayModel.getTaskAnimationManager(displayId)
+                ?: run {
+                    Log.e(TAG, "No TaskAnimationManager found for display $displayId")
+                    ActiveGestureProtoLogProxy.logOnTaskAnimationManagerNotAvailable(displayId)
+                    return false
+                }
         if (taskAnimationManager.isRecentsAnimationRunning) {
             command.setAnimationCallbacks(
                 taskAnimationManager.continueRecentsAnimation(gestureState)
diff --git a/quickstep/src/com/android/quickstep/TaskAnimationManager.java b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
index 1c7f23c..4f00381 100644
--- a/quickstep/src/com/android/quickstep/TaskAnimationManager.java
+++ b/quickstep/src/com/android/quickstep/TaskAnimationManager.java
@@ -77,6 +77,7 @@
 
     private boolean mRecentsAnimationStartPending = false;
     private boolean mShouldIgnoreMotionEvents = false;
+    private final int mDisplayId;
 
     private final TaskStackChangeListener mLiveTileRestartListener = new TaskStackChangeListener() {
         @Override
@@ -101,10 +102,13 @@
         }
     };
 
-    TaskAnimationManager(Context ctx, RecentsAnimationDeviceState deviceState) {
+    public TaskAnimationManager(Context ctx, RecentsAnimationDeviceState deviceState,
+            int displayId) {
         mCtx = ctx;
         mDeviceState = deviceState;
+        mDisplayId = displayId;
     }
+
     SystemUiProxy getSystemUiProxy() {
         return SystemUiProxy.INSTANCE.get(mCtx);
     }
@@ -489,6 +493,7 @@
 
     public void dump(String prefix, PrintWriter pw) {
         pw.println(prefix + "TaskAnimationManager:");
+        pw.println(prefix + "\tmDisplayId=" + mDisplayId);
 
         if (enableHandleDelayedGestureCallbacks()) {
             pw.println(prefix + "\tmRecentsAnimationStartPending=" + mRecentsAnimationStartPending);
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 2bc5626..6912ba7 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -557,7 +557,6 @@
     private OverviewComponentObserver mOverviewComponentObserver;
     private InputConsumerController mInputConsumer;
     private RecentsAnimationDeviceState mDeviceState;
-    private TaskAnimationManager mTaskAnimationManager;
 
     private @NonNull InputConsumer mUncheckedConsumer = InputConsumer.DEFAULT_NO_OP;
     private @NonNull InputConsumer mConsumer = InputConsumer.DEFAULT_NO_OP;
@@ -667,7 +666,7 @@
         if (ENABLE_GESTURE_NAV_ON_CONNECTED_DISPLAYS.isTrue()) {
             mInputMonitorDisplayModel = new InputMonitorDisplayModel(this);
         } else {
-            mInputMonitorCompat = new InputMonitorCompat("swipe-up", Display.DEFAULT_DISPLAY);
+            mInputMonitorCompat = new InputMonitorCompat("swipe-up", DEFAULT_DISPLAY);
             mInputEventReceiver = mInputMonitorCompat.getInputReceiver(Looper.getMainLooper(),
                     mMainChoreographer, this::onInputEvent);
         }
@@ -687,10 +686,9 @@
     public void onUserUnlocked() {
         Log.d(TAG, "onUserUnlocked: userId=" + getUserId()
                 + " instance=" + System.identityHashCode(this));
-        mTaskAnimationManager = new TaskAnimationManager(this, mDeviceState);
         mOverviewComponentObserver = OverviewComponentObserver.INSTANCE.get(this);
         mOverviewCommandHelper = new OverviewCommandHelper(this,
-                mOverviewComponentObserver, mTaskAnimationManager, mRecentsDisplayModel,
+                mOverviewComponentObserver, mRecentsDisplayModel,
                 SystemUiProxy.INSTANCE.get(this).getFocusState(), mTaskbarManager);
         mUserUnlocked = true;
         mInputConsumer.registerInputConsumer();
@@ -765,14 +763,18 @@
         if (LockedUserState.get(this).isUserUnlocked()) {
             long systemUiStateFlags = mDeviceState.getSystemUiStateFlags(displayId);
             mTaskbarManager.onSystemUiFlagsChanged(systemUiStateFlags, displayId);
-            if (displayId == Display.DEFAULT_DISPLAY) {
+            if (displayId == DEFAULT_DISPLAY) {
                 // The following don't care about non-default displays, at least for now. If they
                 // ever will, they should be taken care of.
                 SystemUiProxy.INSTANCE.get(this).setLastSystemUiStateFlags(systemUiStateFlags);
                 mOverviewComponentObserver.setHomeDisabled(mDeviceState.isHomeDisabled());
                 // TODO b/399371607 - Propagate to taskAnimationManager once overview is multi
                 //  display.
-                mTaskAnimationManager.onSystemUiFlagsChanged(lastSysUIFlags, systemUiStateFlags);
+                TaskAnimationManager taskAnimationManager =
+                        mRecentsDisplayModel.getTaskAnimationManager(displayId);
+                if (taskAnimationManager != null) {
+                    taskAnimationManager.onSystemUiFlagsChanged(lastSysUIFlags, systemUiStateFlags);
+                }
             }
         }
     }
@@ -864,11 +866,18 @@
         boolean isHoverActionWithoutConsumer = enableCursorHoverStates()
                 && isHoverActionWithoutConsumer(event);
 
+        TaskAnimationManager taskAnimationManager = mRecentsDisplayModel.getTaskAnimationManager(
+                displayId);
+        if (taskAnimationManager == null) {
+            Log.e(TAG, "TaskAnimationManager not available for displayId " + displayId);
+            ActiveGestureProtoLogProxy.logOnTaskAnimationManagerNotAvailable(displayId);
+            return;
+        }
         if (enableHandleDelayedGestureCallbacks()) {
             if (action == ACTION_DOWN || isHoverActionWithoutConsumer) {
-                mTaskAnimationManager.notifyNewGestureStart();
+                taskAnimationManager.notifyNewGestureStart();
             }
-            if (mTaskAnimationManager.shouldIgnoreMotionEvents()) {
+            if (taskAnimationManager.shouldIgnoreMotionEvents()) {
                 if (action == ACTION_DOWN || isHoverActionWithoutConsumer) {
                     ActiveGestureProtoLogProxy.logOnInputIgnoringFollowingEvents(displayId);
                 }
@@ -938,7 +947,7 @@
                         mDeviceState,
                         prevGestureState,
                         mGestureState,
-                        mTaskAnimationManager,
+                        taskAnimationManager,
                         inputMonitorCompat,
                         getSwipeUpHandlerFactory(),
                         this::onConsumerInactive,
@@ -1055,7 +1064,9 @@
             GestureState.TrackpadGestureType trackpadGestureType) {
         final GestureState gestureState;
         TopTaskTracker.CachedTaskInfo taskInfo;
-        if (mTaskAnimationManager.isRecentsAnimationRunning()) {
+        TaskAnimationManager taskAnimationManager = mRecentsDisplayModel.getTaskAnimationManager(
+                displayId);
+        if (taskAnimationManager != null && taskAnimationManager.isRecentsAnimationRunning()) {
             gestureState = new GestureState(
                     mOverviewComponentObserver, displayId, ActiveGestureLog.INSTANCE.getLogId());
             TopTaskTracker.CachedTaskInfo previousTaskInfo = previousGestureState.getRunningTask();
@@ -1105,7 +1116,7 @@
         mConsumer = mUncheckedConsumer = InputConsumerUtils.getDefaultInputConsumer(
                 displayId,
                 mUserUnlocked,
-                mTaskAnimationManager,
+                mRecentsDisplayModel.getTaskAnimationManager(displayId),
                 mTaskbarManager,
                 CompoundString.NO_OP);
         mGestureState = DEFAULT_STATE;
@@ -1192,13 +1203,11 @@
             if (createdOverviewContainer != null) {
                 createdOverviewContainer.getDeviceProfile().dump(this, "", pw);
             }
+            resource.getTaskAnimationManager().dump("\t", pw);
         }
         pw.println("\tmConsumer=" + mConsumer.getName());
         ActiveGestureLog.INSTANCE.dump("", pw);
         RecentsModel.INSTANCE.get(this).dump("", pw);
-        if (mTaskAnimationManager != null) {
-            mTaskAnimationManager.dump("", pw);
-        }
         mTaskbarManager.dumpLogs("", pw);
         DesktopVisibilityController.INSTANCE.get(this).dumpLogs("", pw);
         pw.println("ContextualSearchStateManager:");
@@ -1210,22 +1219,28 @@
 
     private AbsSwipeUpHandler createLauncherSwipeHandler(
             GestureState gestureState, long touchTimeMs) {
-        return new LauncherSwipeHandlerV2(this, mTaskAnimationManager,
-                gestureState, touchTimeMs, mTaskAnimationManager.isRecentsAnimationRunning(),
+        TaskAnimationManager taskAnimationManager = mRecentsDisplayModel.getTaskAnimationManager(
+                gestureState.getDisplayId());
+        return new LauncherSwipeHandlerV2(this, taskAnimationManager,
+                gestureState, touchTimeMs, taskAnimationManager.isRecentsAnimationRunning(),
                 mInputConsumer, MSDLPlayerWrapper.INSTANCE.get(this));
     }
 
     private AbsSwipeUpHandler createFallbackSwipeHandler(
             GestureState gestureState, long touchTimeMs) {
-        return new FallbackSwipeHandler(this, mTaskAnimationManager,
-                gestureState, touchTimeMs, mTaskAnimationManager.isRecentsAnimationRunning(),
+        TaskAnimationManager taskAnimationManager = mRecentsDisplayModel.getTaskAnimationManager(
+                gestureState.getDisplayId());
+        return new FallbackSwipeHandler(this, taskAnimationManager,
+                gestureState, touchTimeMs, taskAnimationManager.isRecentsAnimationRunning(),
                 mInputConsumer, MSDLPlayerWrapper.INSTANCE.get(this));
     }
 
     private AbsSwipeUpHandler createRecentsWindowSwipeHandler(
             GestureState gestureState, long touchTimeMs) {
-        return new RecentsWindowSwipeHandler(this, mTaskAnimationManager,
-                gestureState, touchTimeMs, mTaskAnimationManager.isRecentsAnimationRunning(),
+        TaskAnimationManager taskAnimationManager = mRecentsDisplayModel.getTaskAnimationManager(
+                gestureState.getDisplayId());
+        return new RecentsWindowSwipeHandler(this, taskAnimationManager,
+                gestureState, touchTimeMs, taskAnimationManager.isRecentsAnimationRunning(),
                 mInputConsumer, MSDLPlayerWrapper.INSTANCE.get(this));
     }
 
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
index 12dc177..0f611eb 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
@@ -18,6 +18,7 @@
 
 import android.content.Context
 import android.view.Display
+import android.view.Display.DEFAULT_DISPLAY
 import androidx.core.util.valueIterator
 import com.android.launcher3.dagger.ApplicationContext
 import com.android.launcher3.dagger.LauncherAppSingleton
@@ -26,6 +27,8 @@
 import com.android.launcher3.util.WallpaperColorHints
 import com.android.quickstep.DisplayModel
 import com.android.quickstep.FallbackWindowInterface
+import com.android.quickstep.RecentsAnimationDeviceState
+import com.android.quickstep.TaskAnimationManager
 import com.android.quickstep.dagger.QuickstepBaseAppComponent
 import com.android.quickstep.fallback.window.RecentsDisplayModel.RecentsDisplayResource
 import com.android.quickstep.fallback.window.RecentsWindowFlags.Companion.enableOverviewInWindow
@@ -55,8 +58,11 @@
     init {
         if (enableOverviewInWindow) {
             registerDisplayListener()
-            tracker.addCloseable { destroy() }
+        } else {
+            // Always create resource for default display
+            storeDisplayResource(DEFAULT_DISPLAY)
         }
+        tracker.addCloseable { destroy() }
     }
 
     override fun createDisplayResource(display: Display): RecentsDisplayResource {
@@ -75,6 +81,10 @@
         return getDisplayResource(displayId)?.fallbackWindowInterface
     }
 
+    fun getTaskAnimationManager(displayId: Int): TaskAnimationManager? {
+        return getDisplayResource(displayId)?.taskAnimationManager
+    }
+
     val activeDisplayResources: Iterable<RecentsDisplayResource>
         get() =
             object : Iterable<RecentsDisplayResource> {
@@ -86,12 +96,20 @@
         val displayContext: Context,
         val wallpaperColorHints: Int,
     ) : DisplayResource() {
-        val recentsWindowManager = RecentsWindowManager(displayContext, wallpaperColorHints)
-        val fallbackWindowInterface: FallbackWindowInterface =
-            FallbackWindowInterface(recentsWindowManager)
+        val recentsWindowManager =
+            if (enableOverviewInWindow) RecentsWindowManager(displayContext, wallpaperColorHints)
+            else null
+        val fallbackWindowInterface =
+            if (enableOverviewInWindow) FallbackWindowInterface(recentsWindowManager) else null
+        val taskAnimationManager =
+            TaskAnimationManager(
+                displayContext,
+                RecentsAnimationDeviceState.INSTANCE.get(displayContext),
+                displayId,
+            )
 
         override fun cleanup() {
-            recentsWindowManager.destroy()
+            recentsWindowManager?.destroy()
         }
 
         override fun dump(prefix: String, writer: PrintWriter) {
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowFlags.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowFlags.kt
index 9953154..d880774 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowFlags.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsWindowFlags.kt
@@ -29,8 +29,15 @@
         val enableFallbackOverviewInWindow: DesktopModeFlag =
             DesktopModeFlag(Flags::enableFallbackOverviewInWindow, false)
 
+        @JvmField
+        val enableOverviewOnConnectedDisplays: DesktopModeFlag =
+            DesktopModeFlag(Flags::enableOverviewOnConnectedDisplays, false)
+
         @JvmStatic
         val enableOverviewInWindow
-            get() = enableLauncherOverviewInWindow.isTrue || enableFallbackOverviewInWindow.isTrue
+            get() =
+                enableLauncherOverviewInWindow.isTrue ||
+                    enableFallbackOverviewInWindow.isTrue ||
+                    enableOverviewOnConnectedDisplays.isTrue
     }
 }
diff --git a/quickstep/src_protolog/com/android/quickstep/util/ActiveGestureProtoLogProxy.java b/quickstep/src_protolog/com/android/quickstep/util/ActiveGestureProtoLogProxy.java
index 18a5338..773a039 100644
--- a/quickstep/src_protolog/com/android/quickstep/util/ActiveGestureProtoLogProxy.java
+++ b/quickstep/src_protolog/com/android/quickstep/util/ActiveGestureProtoLogProxy.java
@@ -551,4 +551,13 @@
                 displayRotation, displaySize.flattenToString(), swipeRegion.toShortString(),
                 ohmRegion.toShortString(), gesturalHeight, largerGesturalHeight, reason);
     }
+
+    public static void logOnTaskAnimationManagerNotAvailable(int displayId) {
+        ActiveGestureLog.INSTANCE.addLog(new ActiveGestureLog.CompoundString(
+                "TaskAnimationManager not available for displayId=%d",
+                displayId));
+        if (!enableActiveGestureProtoLog() || !isProtoLogInitialized()) return;
+        ProtoLog.d(ACTIVE_GESTURE_LOG, "TaskAnimationManager not available for displayId=%d",
+                displayId);
+    }
 }
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
index 6fbbd59..7824000 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/AbsSwipeUpHandlerTestCase.java
@@ -49,6 +49,7 @@
 import android.platform.test.annotations.DisableFlags;
 import android.platform.test.annotations.EnableFlags;
 import android.platform.test.flag.junit.SetFlagsRule;
+import android.view.Display;
 import android.view.RemoteAnimationTarget;
 import android.view.SurfaceControl;
 import android.view.ViewTreeObserver;
@@ -189,7 +190,7 @@
     @Before
     public void setUpRecentsContainer() {
         mTaskAnimationManager = new TaskAnimationManager(mContext,
-                RecentsAnimationDeviceState.INSTANCE.get(mContext));
+                RecentsAnimationDeviceState.INSTANCE.get(mContext), Display.DEFAULT_DISPLAY);
         RecentsViewContainer recentsContainer = getRecentsContainer();
         RECENTS_VIEW recentsView = getRecentsView();
 
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
index 56c01f9..381ac68 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
@@ -67,7 +67,6 @@
                 OverviewCommandHelper(
                     touchInteractionService = mock(),
                     overviewComponentObserver = mock(),
-                    taskAnimationManager = mock(),
                     dispatcherProvider = TestDispatcherProvider(dispatcher),
                     recentsDisplayModel = mock(),
                     focusState = mock(),
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/TaskAnimationManagerTest.java b/quickstep/tests/multivalentTests/src/com/android/quickstep/TaskAnimationManagerTest.java
index 6e9885a..fd88a5c 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/TaskAnimationManagerTest.java
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/TaskAnimationManagerTest.java
@@ -26,6 +26,7 @@
 import android.app.ActivityOptions;
 import android.content.Context;
 import android.content.Intent;
+import android.view.Display;
 
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.SmallTest;
@@ -54,7 +55,7 @@
     public void setUp() {
         MockitoAnnotations.initMocks(this);
         mTaskAnimationManager = new TaskAnimationManager(mContext,
-                RecentsAnimationDeviceState.INSTANCE.get(mContext)) {
+                RecentsAnimationDeviceState.INSTANCE.get(mContext), Display.DEFAULT_DISPLAY) {
             @Override
             SystemUiProxy getSystemUiProxy() {
                 return mSystemUiProxy;
diff --git a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
index 3db4398..93b979c 100644
--- a/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
+++ b/quickstep/tests/src/com/android/quickstep/InputConsumerUtilsTest.java
@@ -126,8 +126,7 @@
 
     @Before
     public void setupTaskAnimationManager() {
-        mTaskAnimationManager = new TaskAnimationManager(
-                mContext, mDeviceState);
+        mTaskAnimationManager = new TaskAnimationManager(mContext, mDeviceState, mDisplayId);
     }
 
     @Before