Merge "[CD Taskbar] Refactor TaskbarManager - Make Window Contexts also follow map pattern" into main
diff --git a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
index 17fb959..36185b1 100644
--- a/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
+++ b/quickstep/src/com/android/launcher3/taskbar/TaskbarManager.java
@@ -115,9 +115,8 @@
     private static final Uri NAV_BAR_KIDS_MODE = Settings.Secure.getUriFor(
             Settings.Secure.NAV_BAR_KIDS_MODE);
 
-    private final Context mWindowContext;
+    private final Context mParentContext;
     private final @Nullable Context mNavigationBarPanelContext;
-    private WindowManager mWindowManager;
     private final TaskbarNavButtonController mDefaultNavButtonController;
     private final ComponentCallbacks mDefaultComponentCallbacks;
 
@@ -132,6 +131,8 @@
             new NonDestroyableScopedUnfoldTransitionProgressProvider();
     /** DisplayId - {@link TaskbarActivityContext} map for Connected Display. */
     private final SparseArray<TaskbarActivityContext> mTaskbars = new SparseArray<>();
+    /** DisplayId - {@link Context} map for Connected Display. */
+    private final SparseArray<Context> mWindowContexts = new SparseArray<>();
     /** DisplayId - {@link FrameLayout} map for Connected Display. */
     private final SparseArray<FrameLayout> mRootLayouts = new SparseArray<>();
     /** DisplayId - {@link Boolean} map indicating if RootLayout was added to window. */
@@ -243,36 +244,35 @@
             Context context,
             AllAppsActionManager allAppsActionManager,
             TaskbarNavButtonCallbacks navCallbacks) {
-        Display display =
-                context.getSystemService(DisplayManager.class).getDisplay(context.getDisplayId());
-        mWindowContext = context.createWindowContext(display,
-                ENABLE_TASKBAR_NAVBAR_UNIFICATION ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL,
-                null);
+        mParentContext = context;
+        createWindowContext(context.getDisplayId());
         mAllAppsActionManager = allAppsActionManager;
+        Display display = context.getSystemService(DisplayManager.class).getDisplay(
+                getDefaultDisplayId());
         mNavigationBarPanelContext = ENABLE_TASKBAR_NAVBAR_UNIFICATION
                 ? context.createWindowContext(display, TYPE_NAVIGATION_BAR_PANEL, null)
                 : null;
         if (enableTaskbarNoRecreate()) {
-            mWindowManager = mWindowContext.getSystemService(WindowManager.class);
             createTaskbarRootLayout(getDefaultDisplayId());
         }
         mDefaultNavButtonController = createDefaultNavButtonController(context, navCallbacks);
         mDefaultComponentCallbacks = createDefaultComponentCallbacks();
-        SettingsCache.INSTANCE.get(mWindowContext)
+        SettingsCache.INSTANCE.get(getPrimaryWindowContext())
                 .register(USER_SETUP_COMPLETE_URI, mOnSettingsChangeListener);
-        SettingsCache.INSTANCE.get(mWindowContext)
+        SettingsCache.INSTANCE.get(getPrimaryWindowContext())
                 .register(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "registering component callbacks from constructor.");
-        mWindowContext.registerComponentCallbacks(mDefaultComponentCallbacks);
-        mShutdownReceiver.register(mWindowContext, Intent.ACTION_SHUTDOWN);
+        getPrimaryWindowContext().registerComponentCallbacks(mDefaultComponentCallbacks);
+        mShutdownReceiver.register(getPrimaryWindowContext(), Intent.ACTION_SHUTDOWN);
         UI_HELPER_EXECUTOR.execute(() -> {
             mSharedState.taskbarSystemActionPendingIntent = PendingIntent.getBroadcast(
-                    mWindowContext,
+                    getPrimaryWindowContext(),
                     SYSTEM_ACTION_ID_TASKBAR,
-                    new Intent(ACTION_SHOW_TASKBAR).setPackage(mWindowContext.getPackageName()),
+                    new Intent(ACTION_SHOW_TASKBAR).setPackage(
+                            getPrimaryWindowContext().getPackageName()),
                     PendingIntent.FLAG_UPDATE_CURRENT | PendingIntent.FLAG_IMMUTABLE);
             mTaskbarBroadcastReceiver.register(
-                    mWindowContext, RECEIVER_NOT_EXPORTED, ACTION_SHOW_TASKBAR);
+                    getPrimaryWindowContext(), RECEIVER_NOT_EXPORTED, ACTION_SHOW_TASKBAR);
         });
 
         debugWhyTaskbarNotDestroyed("TaskbarManager created");
@@ -285,14 +285,15 @@
         return new TaskbarNavButtonController(
                 context,
                 navCallbacks,
-                SystemUiProxy.INSTANCE.get(mWindowContext),
+                SystemUiProxy.INSTANCE.get(getPrimaryWindowContext()),
                 new Handler(),
-                new ContextualSearchInvoker(mWindowContext));
+                new ContextualSearchInvoker(getPrimaryWindowContext()));
     }
 
     private ComponentCallbacks createDefaultComponentCallbacks() {
         return new ComponentCallbacks() {
-            private Configuration mOldConfig = mWindowContext.getResources().getConfiguration();
+            private Configuration mOldConfig =
+                    getPrimaryWindowContext().getResources().getConfiguration();
 
             @Override
             public void onConfigurationChanged(Configuration newConfig) {
@@ -302,7 +303,8 @@
                         "TaskbarManager#mComponentCallbacks.onConfigurationChanged: " + newConfig);
                 // TODO: adapt this logic to be specific to different displays.
                 DeviceProfile dp = mUserUnlocked
-                        ? LauncherAppState.getIDP(mWindowContext).getDeviceProfile(mWindowContext)
+                        ? LauncherAppState.getIDP(getPrimaryWindowContext()).getDeviceProfile(
+                        getPrimaryWindowContext())
                         : null;
                 int configDiff = mOldConfig.diff(newConfig) & ~SKIP_RECREATE_CONFIG_CHANGES;
 
@@ -354,6 +356,7 @@
             int displayId = mTaskbars.keyAt(i);
             destroyTaskbarForDisplay(displayId);
             removeTaskbarRootViewFromWindow(displayId);
+            removeWindowContextFromMap(displayId);
         }
     }
 
@@ -372,7 +375,8 @@
         }
         // make this display-specific
         DeviceProfile dp = mUserUnlocked ?
-                LauncherAppState.getIDP(mWindowContext).getDeviceProfile(mWindowContext) : null;
+                LauncherAppState.getIDP(getWindowContext(displayId)).getDeviceProfile(
+                        getWindowContext(displayId)) : null;
         if (dp == null || !isTaskbarEnabled(dp)) {
             removeTaskbarRootViewFromWindow(displayId);
         }
@@ -419,7 +423,8 @@
      */
     public void onUserUnlocked() {
         mUserUnlocked = true;
-        DisplayController.INSTANCE.get(mWindowContext).addChangeListener(mRecreationListener);
+        DisplayController.INSTANCE.get(getPrimaryWindowContext()).addChangeListener(
+                mRecreationListener);
         recreateTaskbar();
         addTaskbarRootViewToWindow(getDefaultDisplayId());
     }
@@ -482,7 +487,8 @@
                 return ql.getUnfoldTransitionProgressProvider();
             }
         } else {
-            return SystemUiProxy.INSTANCE.get(mWindowContext).getUnfoldTransitionProvider();
+            return SystemUiProxy.INSTANCE.get(
+                    getPrimaryWindowContext()).getUnfoldTransitionProvider();
         }
         return null;
     }
@@ -528,7 +534,8 @@
             Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "recreateTaskbarForDisplay: " + displayId);
             // TODO: make this code display specific
             DeviceProfile dp = mUserUnlocked ?
-                    LauncherAppState.getIDP(mWindowContext).getDeviceProfile(mWindowContext) : null;
+                    LauncherAppState.getIDP(getWindowContext(displayId)).getDeviceProfile(
+                            getWindowContext(displayId)) : null;
 
             // All Apps action is unrelated to navbar unification, so we only need to check DP.
             final boolean isLargeScreenTaskbar = dp != null && dp.isTaskbarPresent;
@@ -542,7 +549,7 @@
                 + " FLAG_HIDE_NAVBAR_WINDOW=" + ENABLE_TASKBAR_NAVBAR_UNIFICATION
                 + " dp.isTaskbarPresent=" + (dp == null ? "null" : dp.isTaskbarPresent));
             if (!isTaskbarEnabled || !isLargeScreenTaskbar) {
-                SystemUiProxy.INSTANCE.get(mWindowContext)
+                SystemUiProxy.INSTANCE.get(getPrimaryWindowContext())
                     .notifyTaskbarStatus(/* visible */ false, /* stashed */ false);
                 if (!isTaskbarEnabled) {
                     return;
@@ -754,23 +761,24 @@
         mRecentsViewContainer = null;
         debugWhyTaskbarNotDestroyed("TaskbarManager#destroy()");
         removeActivityCallbacksAndListeners();
-        mTaskbarBroadcastReceiver.unregisterReceiverSafely(mWindowContext);
-        destroyAllTaskbars();
+        mTaskbarBroadcastReceiver.unregisterReceiverSafely(getPrimaryWindowContext());
+
         if (mUserUnlocked) {
-            DisplayController.INSTANCE.get(mWindowContext).removeChangeListener(
+            DisplayController.INSTANCE.get(getPrimaryWindowContext()).removeChangeListener(
                     mRecreationListener);
         }
-        SettingsCache.INSTANCE.get(mWindowContext)
+        SettingsCache.INSTANCE.get(getPrimaryWindowContext())
                 .unregister(USER_SETUP_COMPLETE_URI, mOnSettingsChangeListener);
-        SettingsCache.INSTANCE.get(mWindowContext)
+        SettingsCache.INSTANCE.get(getPrimaryWindowContext())
                 .unregister(NAV_BAR_KIDS_MODE, mOnSettingsChangeListener);
         Log.d(TASKBAR_NOT_DESTROYED_TAG, "unregistering component callbacks from destroy().");
-        mWindowContext.unregisterComponentCallbacks(mDefaultComponentCallbacks);
-        mShutdownReceiver.unregisterReceiverSafely(mWindowContext);
+        getPrimaryWindowContext().unregisterComponentCallbacks(mDefaultComponentCallbacks);
+        mShutdownReceiver.unregisterReceiverSafely(getPrimaryWindowContext());
+        destroyAllTaskbars();
     }
 
     public @Nullable TaskbarActivityContext getCurrentActivityContext() {
-        return getTaskbarForDisplay(mWindowContext.getDisplayId());
+        return getTaskbarForDisplay(getDefaultDisplayId());
     }
 
     public void dumpLogs(String prefix, PrintWriter pw) {
@@ -786,7 +794,6 @@
                 taskbar.dumpLogs(prefix + "\t\t", pw);
             }
         }
-
     }
 
     private void addTaskbarRootViewToWindow(int displayId) {
@@ -800,8 +807,7 @@
         if (!isTaskbarRootLayoutAddedForDisplay(displayId)) {
             FrameLayout rootLayout = getTaskbarRootLayoutForDisplay(displayId);
             if (rootLayout != null) {
-                mWindowManager.addView(getTaskbarRootLayoutForDisplay(displayId),
-                        taskbar.getWindowLayoutParams());
+                getWindowManager(displayId).addView(rootLayout, taskbar.getWindowLayoutParams());
                 mAddedRootLayouts.put(displayId, true);
             } else {
                 Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW,
@@ -822,7 +828,7 @@
         }
 
         if (isTaskbarRootLayoutAddedForDisplay(displayId)) {
-            mWindowManager.removeViewImmediate(rootLayout);
+            getWindowManager(displayId).removeViewImmediate(rootLayout);
             mAddedRootLayouts.put(displayId, false);
             removeTaskbarRootLayoutFromMap(displayId);
         }
@@ -855,10 +861,10 @@
      * Creates a {@link TaskbarActivityContext} for the given display and adds it to the map.
      */
     private TaskbarActivityContext createTaskbarActivityContext(DeviceProfile dp, int displayId) {
-        TaskbarActivityContext newTaskbar = new TaskbarActivityContext(mWindowContext,
+        TaskbarActivityContext newTaskbar = new TaskbarActivityContext(getWindowContext(displayId),
                 mNavigationBarPanelContext, dp, mDefaultNavButtonController,
                 mUnfoldProgressProvider, isDefaultDisplay(displayId),
-                SystemUiProxy.INSTANCE.get(mWindowContext));
+                SystemUiProxy.INSTANCE.get(getPrimaryWindowContext()));
 
         addTaskbarToMap(displayId, newTaskbar);
         return newTaskbar;
@@ -892,7 +898,7 @@
      */
     private void createTaskbarRootLayout(int displayId) {
         Log.d(ILLEGAL_ARGUMENT_WM_ADD_VIEW, "createTaskbarRootLayout: " + displayId);
-        FrameLayout newTaskbarRootLayout = new FrameLayout(mWindowContext) {
+        FrameLayout newTaskbarRootLayout = new FrameLayout(getWindowContext(displayId)) {
             @Override
             public boolean dispatchTouchEvent(MotionEvent ev) {
                 // The motion events can be outside the view bounds of task bar, and hence
@@ -958,8 +964,75 @@
         Log.d(NULL_TASKBAR_ROOT_LAYOUT_TAG, "mRootLayouts.size()=" + mRootLayouts.size());
     }
 
+    /**
+     * Creates {@link Context} for the taskbar on the specified display and›› adds it to map.
+     * @param displayId The ID of the display for which to create the window context.
+     */
+    private void createWindowContext(int displayId) {
+        DisplayManager displayManager = mParentContext.getSystemService(DisplayManager.class);
+        if (displayManager == null) {
+            return;
+        }
+
+        Display display = displayManager.getDisplay(displayId);
+        if (display != null) {
+            int windowType = (ENABLE_TASKBAR_NAVBAR_UNIFICATION && isDefaultDisplay(displayId))
+                    ? TYPE_NAVIGATION_BAR : TYPE_NAVIGATION_BAR_PANEL;
+            Context newContext = mParentContext.createWindowContext(display, windowType, null);
+            addWindowContextToMap(displayId, newContext);
+        }
+    }
+
+    /**
+     * Retrieves the window context of the taskbar for the specified display.
+     *
+     * @param displayId The ID of the display for which to retrieve the window context.
+     * @return The Window Context {@link Context} for a given display or {@code null}.
+     */
+    private Context getWindowContext(int displayId) {
+        return mWindowContexts.get(displayId);
+    }
+
+    @VisibleForTesting
+    public Context getPrimaryWindowContext() {
+        return getWindowContext(getDefaultDisplayId());
+    }
+
+    /**
+     * Retrieves the window manager {@link WindowManager} of the taskbar for the specified display.
+     *
+     * @param displayId The ID of the display for which to retrieve the window manager.
+     * @return The window manager {@link WindowManager} for a given display or {@code null}.
+     */
+    private WindowManager getWindowManager(int displayId) {
+        return getWindowContext(displayId).getSystemService(WindowManager.class);
+    }
+
+    /**
+     * Adds the window context {@link Context} to taskbar map, mapped to display ID.
+     *
+     * @param displayId The ID of the display to associate with the taskbar root layout.
+     * @param windowContext The window context {@link Context} to add to the map.
+     */
+    private void addWindowContextToMap(int displayId, @NonNull Context windowContext) {
+        if (!mWindowContexts.contains(displayId)) {
+            mWindowContexts.put(displayId, windowContext);
+        }
+    }
+
+    /**
+     * Removes the window context {@link Context} for given display ID from the taskbar map.
+     *
+     * @param displayId The ID of the display for which to remove the taskbar root layout.
+     */
+    private void removeWindowContextFromMap(int displayId) {
+        if (mWindowContexts.contains(displayId)) {
+            mWindowContexts.delete(displayId);
+        }
+    }
+
     private int getDefaultDisplayId() {
-        return mWindowContext.getDisplayId();
+        return mParentContext.getDisplayId();
     }
 
     /** Temp logs for b/254119092. */
@@ -974,9 +1047,14 @@
 
         boolean activityTaskbarPresent = mActivity != null
                 && mActivity.getDeviceProfile().isTaskbarPresent;
-        // TODO: make this display specific
-        boolean contextTaskbarPresent = mUserUnlocked && LauncherAppState.getIDP(mWindowContext)
-                .getDeviceProfile(mWindowContext).isTaskbarPresent;
+        Context windowContext = getWindowContext(displayId);
+        if (windowContext == null) {
+            log.add("window context for displayId" + displayId);
+            return;
+        }
+
+        boolean contextTaskbarPresent = mUserUnlocked && LauncherAppState.getIDP(windowContext)
+                .getDeviceProfile(windowContext).isTaskbarPresent;
         if (activityTaskbarPresent == contextTaskbarPresent) {
             log.add("mActivity and mWindowContext agree taskbarIsPresent=" + contextTaskbarPresent);
             Log.d(TASKBAR_NOT_DESTROYED_TAG, log.toString());
@@ -993,12 +1071,12 @@
             log.add("\t\tmActivity.getDeviceProfile().isTaskbarPresent="
                     + activityTaskbarPresent);
         }
-        log.add("\tmWindowContext logs:");
-        log.add("\t\tmWindowContext=" + mWindowContext);
-        log.add("\t\tmWindowContext.getResources().getConfiguration()="
-                + mWindowContext.getResources().getConfiguration());
+        log.add("\tWindowContext logs:");
+        log.add("\t\tWindowContext=" + windowContext);
+        log.add("\t\tWindowContext.getResources().getConfiguration()="
+                + windowContext.getResources().getConfiguration());
         if (mUserUnlocked) {
-            log.add("\t\tLauncherAppState.getIDP().getDeviceProfile(mWindowContext)"
+            log.add("\t\tLauncherAppState.getIDP().getDeviceProfile(getPrimaryWindowContext())"
                     + ".isTaskbarPresent=" + contextTaskbarPresent);
         } else {
             log.add("\t\tCouldn't get DeviceProfile because !mUserUnlocked");
@@ -1010,8 +1088,4 @@
     private final DeviceProfile.OnDeviceProfileChangeListener mDebugActivityDeviceProfileChanged =
             dp -> debugWhyTaskbarNotDestroyed("mActivity onDeviceProfileChanged");
 
-    @VisibleForTesting
-    public Context getWindowContext() {
-        return mWindowContext;
-    }
 }
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 e150568..90c9553 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
@@ -125,7 +125,8 @@
                     // Needs to be set on window context instead of sandbox context, because it does
                     // does not propagate between them. However, this change will impact created
                     // TaskbarActivityContext instances, since they wrap the window context.
-                    taskbarManager.windowContext.resources.configuration.setLayoutDirection(
+                    // TODO: iterate through all window contexts and do this.
+                    taskbarManager.primaryWindowContext.resources.configuration.setLayoutDirection(
                         RTL_LOCALE
                     )
                 }