Merge "Updates the app icon based on the TaskData" into main
diff --git a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
index 23065b5..5afc5ed 100644
--- a/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
+++ b/quickstep/src/com/android/launcher3/taskbar/KeyboardQuickSwitchController.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.taskbar;
 
+import static com.android.launcher3.Flags.enableAltTabKqsOnConnectedDisplays;
+
 import android.content.ComponentName;
 import android.content.pm.ActivityInfo;
 import android.view.MotionEvent;
@@ -272,11 +274,26 @@
     }
 
     private void processLoadedTasksOnDesktop(List<GroupTask> tasks, Set<Integer> taskIdsToExclude) {
-        // Find the single desktop task that contains a grouping of desktop tasks
-        DesktopTask desktopTask = findDesktopTask(tasks);
+        // Find all desktop tasks.
+        List<DesktopTask> desktopTasks = tasks.stream()
+                .filter(t -> t instanceof DesktopTask)
+                .map(t -> (DesktopTask) t)
+                .toList();
 
-        if (desktopTask != null) {
-            mTasks = desktopTask.getTasks().stream()
+        // Apps on the connected displays seem to be in different Desktop tasks even with the
+        // multiple desktops flag disabled. So, until multiple desktops is implemented the following
+        // should help with team-fooding Alt+tab on connected displays. Post multiple desktop,
+        // further changes maybe required to support launching selected desktops.
+        if (enableAltTabKqsOnConnectedDisplays()) {
+            mTasks = desktopTasks.stream()
+                    .flatMap(t -> t.getTasks().stream())
+                    .map(SingleTask::new)
+                    .filter(task -> !shouldExcludeTask(task, taskIdsToExclude))
+                    .collect(Collectors.toList());
+
+            mNumHiddenTasks = Math.max(0, tasks.size() - desktopTasks.size());
+        } else if (!desktopTasks.isEmpty()) {
+            mTasks = desktopTasks.get(0).getTasks().stream()
                     .map(SingleTask::new)
                     .filter(task -> !shouldExcludeTask(task, taskIdsToExclude))
                     .collect(Collectors.toList());
@@ -289,14 +306,6 @@
         }
     }
 
-    @Nullable
-    private DesktopTask findDesktopTask(List<GroupTask> tasks) {
-        return (DesktopTask) tasks.stream()
-                .filter(t -> t instanceof DesktopTask)
-                .findFirst()
-                .orElse(null);
-    }
-
     void closeQuickSwitchView() {
         closeQuickSwitchView(true);
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
index 9d1fc15..a6d3cde 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarActivityContext.java
@@ -24,6 +24,8 @@
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
 import static android.window.SplashScreen.SPLASH_SCREEN_STYLE_UNDEFINED;
 
+import static androidx.annotation.VisibleForTesting.PACKAGE_PRIVATE;
+
 import static com.android.launcher3.AbstractFloatingView.TYPE_ALL;
 import static com.android.launcher3.AbstractFloatingView.TYPE_ON_BOARD_POPUP;
 import static com.android.launcher3.AbstractFloatingView.TYPE_REBIND_SAFE;
@@ -2008,7 +2010,8 @@
         return mControllers.uiController.isIconAlignedWithHotseat();
     }
 
-    @VisibleForTesting
+    // TODO(b/395061396): Remove `otherwise` when overview in widow is enabled.
+    @VisibleForTesting(otherwise = PACKAGE_PRIVATE)
     public TaskbarControllers getControllers() {
         return mControllers;
     }
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index ea2dec1..028e9e7 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -32,6 +32,7 @@
 import static com.android.launcher3.util.FlagDebugUtils.formatFlagChange;
 import static com.android.quickstep.util.SystemActionConstants.ACTION_SHOW_TASKBAR;
 import static com.android.quickstep.util.SystemActionConstants.SYSTEM_ACTION_ID_TASKBAR;
+import static com.android.window.flags.Flags.enableTaskbarConnectedDisplays;
 
 import android.annotation.SuppressLint;
 import android.app.PendingIntent;
@@ -71,6 +72,7 @@
 import com.android.quickstep.AllAppsActionManager;
 import com.android.quickstep.RecentsActivity;
 import com.android.quickstep.SystemUiProxy;
+import com.android.quickstep.fallback.window.RecentsDisplayModel;
 import com.android.quickstep.fallback.window.RecentsWindowManager;
 import com.android.quickstep.util.ContextualSearchInvoker;
 import com.android.quickstep.views.RecentsViewContainer;
@@ -182,6 +184,7 @@
             new SimpleBroadcastReceiver(UI_HELPER_EXECUTOR, this::showTaskbarFromBroadcast);
 
     private final AllAppsActionManager mAllAppsActionManager;
+    private final RecentsDisplayModel mRecentsDisplayModel;
 
     private final Runnable mActivityOnDestroyCallback = new Runnable() {
         @Override
@@ -242,10 +245,12 @@
     public TaskbarManager(
             Context context,
             AllAppsActionManager allAppsActionManager,
-            TaskbarNavButtonCallbacks navCallbacks) {
+            TaskbarNavButtonCallbacks navCallbacks,
+            RecentsDisplayModel recentsDisplayModel) {
         mParentContext = context;
         createWindowContext(context.getDisplayId());
         mAllAppsActionManager = allAppsActionManager;
+        mRecentsDisplayModel = recentsDisplayModel;
         if (enableTaskbarNoRecreate()) {
             createTaskbarRootLayout(getDefaultDisplayId());
         }
@@ -487,6 +492,18 @@
         return null;
     }
 
+    /** Creates a {@link TaskbarUIController} to use with non default displays. */
+    private TaskbarUIController createTaskbarUIControllerForNonDefaultDisplay(int displayId) {
+        if (RecentsDisplayModel.enableOverviewInWindow()) {
+            RecentsViewContainer rvc = mRecentsDisplayModel.getRecentsWindowManager(displayId);
+            if (rvc != null) {
+                return createTaskbarUIControllerForRecentsViewContainer(rvc);
+            }
+        }
+
+        return new TaskbarUIController();
+    }
+
     /**
      * Creates a {@link TaskbarUIController} to use while the given StatefulActivity is active.
      */
@@ -561,7 +578,11 @@
             mSharedState.allAppsVisible = mSharedState.allAppsVisible && isLargeScreenTaskbar;
             taskbar.init(mSharedState);
 
-            if (mRecentsViewContainer != null) {
+            // Non default displays should not use LauncherTaskbarUIController as they shouldn't
+            // have access to the Launcher activity.
+            if (enableTaskbarConnectedDisplays() && !isDefaultDisplay(displayId)) {
+                taskbar.setUIController(createTaskbarUIControllerForNonDefaultDisplay(displayId));
+            } else if (mRecentsViewContainer != null) {
                 taskbar.setUIController(
                         createTaskbarUIControllerForRecentsViewContainer(mRecentsViewContainer));
             }
@@ -829,6 +850,23 @@
     }
 
     /**
+     * Returns the {@link TaskbarUIController} associated with the given display ID.
+     * TODO(b/395061396): Remove this method when overview in widow is enabled.
+     *
+     * @param displayId The ID of the display to retrieve the taskbar for.
+     * @return The {@link TaskbarUIController} for the specified display, or
+     *         {@code null} if no taskbar is associated with that display.
+     */
+    @Nullable
+    public TaskbarUIController getUIControllerForDisplay(int displayId) {
+        if (!mTaskbars.contains(displayId)) {
+            return null;
+        }
+
+        return getTaskbarForDisplay(displayId).getControllers().uiController;
+    }
+
+    /**
      * Retrieves whether RootLayout was added to window for specific display, or false if no
      * such mapping has been made.
      *
diff --git a/quickstep/src/com/android/quickstep/FocusState.kt b/quickstep/src/com/android/quickstep/FocusState.kt
index ba3991f..7c6aa5b 100644
--- a/quickstep/src/com/android/quickstep/FocusState.kt
+++ b/quickstep/src/com/android/quickstep/FocusState.kt
@@ -27,7 +27,10 @@
 class FocusState {
 
     var focusedDisplayId = DEFAULT_DISPLAY
-        private set
+        private set(value) {
+            field = value
+            listeners.forEach { it.onFocusedDisplayChanged(value) }
+        }
 
     private var listeners = mutableSetOf<FocusChangeListener>()
 
@@ -40,9 +43,7 @@
             transitions?.setFocusTransitionListener(
                 object : Stub() {
                     override fun onFocusedDisplayChanged(displayId: Int) {
-                        Executors.MAIN_EXECUTOR.execute {
-                            listeners.forEach { it.onFocusedDisplayChanged(displayId) }
-                        }
+                        Executors.MAIN_EXECUTOR.execute { focusedDisplayId = displayId }
                     }
                 }
             )
diff --git a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
index 94d115b..42aa86e 100644
--- a/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
+++ b/quickstep/src/com/android/quickstep/OverviewCommandHelper.kt
@@ -28,6 +28,7 @@
 import androidx.annotation.UiThread
 import androidx.annotation.VisibleForTesting
 import com.android.internal.jank.Cuj
+import com.android.launcher3.Flags.enableAltTabKqsOnConnectedDisplays
 import com.android.launcher3.Flags.enableFallbackOverviewInWindow
 import com.android.launcher3.Flags.enableLargeDesktopWindowingTile
 import com.android.launcher3.Flags.enableLauncherOverviewInWindow
@@ -38,6 +39,8 @@
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_3_BUTTON
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_KEYBOARD_QUICK_SWITCH
 import com.android.launcher3.logging.StatsLogManager.LauncherEvent.LAUNCHER_OVERVIEW_SHOW_OVERVIEW_FROM_KEYBOARD_SHORTCUT
+import com.android.launcher3.taskbar.TaskbarManager
+import com.android.launcher3.taskbar.TaskbarUIController
 import com.android.launcher3.util.Executors
 import com.android.launcher3.util.RunnableList
 import com.android.launcher3.util.coroutines.DispatcherProvider
@@ -48,6 +51,7 @@
 import com.android.quickstep.OverviewCommandHelper.CommandType.KEYBOARD_INPUT
 import com.android.quickstep.OverviewCommandHelper.CommandType.SHOW
 import com.android.quickstep.OverviewCommandHelper.CommandType.TOGGLE
+import com.android.quickstep.fallback.window.RecentsDisplayModel
 import com.android.quickstep.util.ActiveGestureLog
 import com.android.quickstep.util.ActiveGestureProtoLogProxy
 import com.android.quickstep.views.RecentsView
@@ -72,6 +76,9 @@
     private val overviewComponentObserver: OverviewComponentObserver,
     private val taskAnimationManager: TaskAnimationManager,
     private val dispatcherProvider: DispatcherProvider = ProductionDispatchers,
+    private val recentsDisplayModel: RecentsDisplayModel,
+    private val focusState: FocusState,
+    private val taskbarManager: TaskbarManager,
 ) {
     private val coroutineScope = CoroutineScope(SupervisorJob() + dispatcherProvider.background)
 
@@ -291,15 +298,51 @@
                 deviceProfile != null &&
                 (deviceProfile.isTablet || deviceProfile.isTwoPanels)
 
+        val focusedDisplayId = focusState.focusedDisplayId
+        val focusedDisplayUIController: TaskbarUIController? =
+            if (RecentsDisplayModel.enableOverviewInWindow()) {
+                Log.d(
+                    TAG,
+                    "Querying RecentsDisplayModel for TaskbarUIController for display: $focusedDisplayId",
+                )
+                recentsDisplayModel.getRecentsWindowManager(focusedDisplayId)?.taskbarUIController
+            } else {
+                Log.d(
+                    TAG,
+                    "Querying TaskbarManager for TaskbarUIController for display: $focusedDisplayId",
+                )
+                // TODO(b/395061396): Remove this path when overview in widow is enabled.
+                taskbarManager.getUIControllerForDisplay(focusedDisplayId)
+            }
+        Log.d(
+            TAG,
+            "TaskbarUIController for display $focusedDisplayId was" +
+                "${if (focusedDisplayUIController == null) " not" else ""} found",
+        )
+
         when (command.type) {
             HIDE -> {
                 if (!allowQuickSwitch) return true
-                keyboardTaskFocusIndex = uiController!!.launchFocusedTask()
+                keyboardTaskFocusIndex =
+                    if (
+                        enableAltTabKqsOnConnectedDisplays() && focusedDisplayUIController != null
+                    ) {
+                        focusedDisplayUIController.launchFocusedTask()
+                    } else {
+                        uiController!!.launchFocusedTask()
+                    }
+
                 if (keyboardTaskFocusIndex == -1) return true
             }
             KEYBOARD_INPUT ->
                 if (allowQuickSwitch) {
-                    uiController!!.openQuickSwitchView()
+                    if (
+                        enableAltTabKqsOnConnectedDisplays() && focusedDisplayUIController != null
+                    ) {
+                        focusedDisplayUIController.openQuickSwitchView()
+                    } else {
+                        uiController!!.openQuickSwitchView()
+                    }
                     return true
                 } else {
                     keyboardTaskFocusIndex = 0
diff --git a/quickstep/src/com/android/quickstep/SystemUiProxy.kt b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
index d5382ad..a1ac39e 100644
--- a/quickstep/src/com/android/quickstep/SystemUiProxy.kt
+++ b/quickstep/src/com/android/quickstep/SystemUiProxy.kt
@@ -149,7 +149,7 @@
     private var backToLauncherRunner: IRemoteAnimationRunner? = null
     private var dragAndDrop: IDragAndDrop? = null
     val homeVisibilityState = HomeVisibilityState()
-    private val focusState = FocusState()
+    val focusState = FocusState()
 
     // Used to dedupe calls to SystemUI
     private var lastShelfHeight = 0
diff --git a/quickstep/src/com/android/quickstep/TouchInteractionService.java b/quickstep/src/com/android/quickstep/TouchInteractionService.java
index 2df4a45..ba4c65a 100644
--- a/quickstep/src/com/android/quickstep/TouchInteractionService.java
+++ b/quickstep/src/com/android/quickstep/TouchInteractionService.java
@@ -88,6 +88,7 @@
 import com.android.launcher3.util.TraceHelper;
 import com.android.quickstep.OverviewCommandHelper.CommandType;
 import com.android.quickstep.OverviewComponentObserver.OverviewChangeListener;
+import com.android.quickstep.fallback.window.RecentsDisplayModel;
 import com.android.quickstep.fallback.window.RecentsWindowSwipeHandler;
 import com.android.quickstep.inputconsumers.BubbleBarInputConsumer;
 import com.android.quickstep.inputconsumers.OneHandedModeInputConsumer;
@@ -584,7 +585,8 @@
             initInputMonitor("onTrackpadConnected()");
         });
 
-        mTaskbarManager = new TaskbarManager(this, mAllAppsActionManager, mNavCallbacks);
+        mTaskbarManager = new TaskbarManager(this, mAllAppsActionManager, mNavCallbacks,
+                RecentsDisplayModel.getINSTANCE().get(this));
         mDesktopAppLaunchTransitionManager =
                 new DesktopAppLaunchTransitionManager(this, SystemUiProxy.INSTANCE.get(this));
         mDesktopAppLaunchTransitionManager.registerTransitions();
@@ -641,7 +643,9 @@
         mTaskAnimationManager = new TaskAnimationManager(this, mDeviceState);
         mOverviewComponentObserver = OverviewComponentObserver.INSTANCE.get(this);
         mOverviewCommandHelper = new OverviewCommandHelper(this,
-                mOverviewComponentObserver, mTaskAnimationManager);
+                mOverviewComponentObserver, mTaskAnimationManager,
+                RecentsDisplayModel.getINSTANCE().get(this),
+                SystemUiProxy.INSTANCE.get(this).getFocusState(), mTaskbarManager);
         mResetGestureInputConsumer = new ResetGestureInputConsumer(
                 mTaskAnimationManager, mTaskbarManager::getCurrentActivityContext);
         mInputConsumer.registerInputConsumer();
diff --git a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
index 31a1be8..95a3ec2 100644
--- a/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
+++ b/quickstep/src/com/android/quickstep/fallback/window/RecentsDisplayModel.kt
@@ -50,10 +50,14 @@
             DaggerSingletonObject<RecentsDisplayModel>(
                 QuickstepBaseAppComponent::getRecentsDisplayModel
             )
+
+        @JvmStatic
+        fun enableOverviewInWindow() =
+            Flags.enableFallbackOverviewInWindow() || Flags.enableLauncherOverviewInWindow()
     }
 
     init {
-        if (Flags.enableFallbackOverviewInWindow() || Flags.enableLauncherOverviewInWindow()) {
+        if (enableOverviewInWindow()) {
             displayManager.registerDisplayListener(displayListener, Executors.MAIN_EXECUTOR.handler)
             // In the scenario where displays were added before this display listener was
             // registered, we should store the RecentsDisplayResources for those displays
diff --git a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
index 002a4e8..b1a5920 100644
--- a/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
+++ b/quickstep/src/com/android/quickstep/recents/data/TasksRepository.kt
@@ -66,14 +66,10 @@
                 )
                 tasks.value = MapForStateFlow(recentTasks)
 
-                // Request data for completed tasks to prevent stale data.
-                // This will prevent thumbnail and icon from being replaced and
-                // null due to race condition.
-                taskRequests.values.forEach { (taskKey, job) ->
-                    if (job.isCompleted) {
-                        requestTaskData(taskKey.id)
-                    }
-                }
+                // Request data for tasks to prevent stale data.
+                // This will prevent thumbnail and icon from being replaced and null due to
+                // race condition. The new request will hit the cache and return immediately.
+                taskRequests.keys.forEach(::requestTaskData)
             }
         }
         return tasks.map { it.values.toList() }
diff --git a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
index 90c9553..0204b2d 100644
--- a/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/launcher3/taskbar/rules/TaskbarUnitTestRule.kt
@@ -33,6 +33,7 @@
 import com.android.launcher3.util.Executors.UI_HELPER_EXECUTOR
 import com.android.launcher3.util.TestUtil
 import com.android.quickstep.AllAppsActionManager
+import com.android.quickstep.fallback.window.RecentsDisplayModel
 import java.lang.reflect.Field
 import java.lang.reflect.ParameterizedType
 import java.util.Locale
@@ -110,6 +111,7 @@
                                     PendingIntent(IIntentSender.Default())
                                 },
                                 object : TaskbarNavButtonCallbacks {},
+                                RecentsDisplayModel.INSTANCE.get(context),
                             ) {
                             override fun recreateTaskbar() {
                                 super.recreateTaskbar()
diff --git a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
index 0ae710f..56c01f9 100644
--- a/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
+++ b/quickstep/tests/multivalentTests/src/com/android/quickstep/OverviewCommandHelperTest.kt
@@ -68,7 +68,10 @@
                     touchInteractionService = mock(),
                     overviewComponentObserver = mock(),
                     taskAnimationManager = mock(),
-                    dispatcherProvider = TestDispatcherProvider(dispatcher)
+                    dispatcherProvider = TestDispatcherProvider(dispatcher),
+                    recentsDisplayModel = mock(),
+                    focusState = mock(),
+                    taskbarManager = mock(),
                 )
             )