Merge "Add UserManagerInternal#getDisplaysAssignedToUser() method." into udc-dev
diff --git a/services/core/java/com/android/server/pm/UserManagerInternal.java b/services/core/java/com/android/server/pm/UserManagerInternal.java
index eb37302..721ad88 100644
--- a/services/core/java/com/android/server/pm/UserManagerInternal.java
+++ b/services/core/java/com/android/server/pm/UserManagerInternal.java
@@ -491,8 +491,11 @@
     public abstract boolean isUserVisible(@UserIdInt int userId, int displayId);
 
     /**
-     * Returns the display id assigned to the user, or {@code Display.INVALID_DISPLAY} if the
-     * user is not assigned to any display.
+     * Returns the main display id assigned to the user, or {@code Display.INVALID_DISPLAY} if the
+     * user is not assigned to any main display.
+     *
+     * <p>In the context of multi-user multi-display, there can be multiple main displays, at most
+     * one per each zone. Main displays are where UI is launched which a user interacts with.
      *
      * <p>The current foreground user and its running profiles are associated with the
      * {@link android.view.Display#DEFAULT_DISPLAY default display}, while other users would only be
@@ -503,9 +506,20 @@
      *
      * <p>If the user is a profile and is running, it's assigned to its parent display.
      */
+    // TODO(b/272366483) rename this method to avoid confusion with getDisplaysAssignedTOUser().
     public abstract int getDisplayAssignedToUser(@UserIdInt int userId);
 
     /**
+     * Returns all display ids assigned to the user including {@link
+     * #assignUserToExtraDisplay(int, int) extra displays}, or {@code null} if there is no display
+     * assigned to the specified user.
+     *
+     * <p>Note that this method is different from {@link #getDisplayAssignedToUser(int)}, which
+     * returns a main display only.
+     */
+    public abstract @Nullable int[] getDisplaysAssignedToUser(@UserIdInt int userId);
+
+    /**
      * Returns the main user (i.e., not a profile) that is assigned to the display, or the
      * {@link android.app.ActivityManager#getCurrentUser() current foreground user} if no user is
      * associated with the display.
diff --git a/services/core/java/com/android/server/pm/UserManagerService.java b/services/core/java/com/android/server/pm/UserManagerService.java
index d1883f6..cc40363 100644
--- a/services/core/java/com/android/server/pm/UserManagerService.java
+++ b/services/core/java/com/android/server/pm/UserManagerService.java
@@ -7190,6 +7190,11 @@
         }
 
         @Override
+        public @Nullable int[] getDisplaysAssignedToUser(@UserIdInt int userId) {
+            return mUserVisibilityMediator.getDisplaysAssignedToUser(userId);
+        }
+
+        @Override
         public @UserIdInt int getUserAssignedToDisplay(int displayId) {
             return mUserVisibilityMediator.getUserAssignedToDisplay(displayId);
         }
diff --git a/services/core/java/com/android/server/pm/UserVisibilityMediator.java b/services/core/java/com/android/server/pm/UserVisibilityMediator.java
index a8615c2..b67850a 100644
--- a/services/core/java/com/android/server/pm/UserVisibilityMediator.java
+++ b/services/core/java/com/android/server/pm/UserVisibilityMediator.java
@@ -786,6 +786,49 @@
         }
     }
 
+    /** See {@link UserManagerInternal#getDisplaysAssignedToUser(int)}. */
+    @Nullable
+    public int[] getDisplaysAssignedToUser(@UserIdInt int userId) {
+        int mainDisplayId = getDisplayAssignedToUser(userId);
+        if (mainDisplayId == INVALID_DISPLAY) {
+            // The user will not have any extra displays if they have no main display.
+            // Return null if no display is assigned to the user.
+            if (DBG) {
+                Slogf.d(TAG, "getDisplaysAssignedToUser(): returning null"
+                        + " because there is no display assigned to user %d", userId);
+            }
+            return null;
+        }
+
+        synchronized (mLock) {
+            if (mExtraDisplaysAssignedToUsers == null
+                    || mExtraDisplaysAssignedToUsers.size() == 0) {
+                return new int[]{mainDisplayId};
+            }
+
+            int count = 0;
+            int[] displayIds = new int[mExtraDisplaysAssignedToUsers.size() + 1];
+            displayIds[count++] = mainDisplayId;
+            for (int i = 0; i < mExtraDisplaysAssignedToUsers.size(); ++i) {
+                if (mExtraDisplaysAssignedToUsers.valueAt(i) == userId) {
+                    displayIds[count++] = mExtraDisplaysAssignedToUsers.keyAt(i);
+                }
+            }
+            // Return the array if the array length happens to be correct.
+            if (displayIds.length == count) {
+                return displayIds;
+            }
+
+            // Copy the results to a new array with the exact length. The size of displayIds[] is
+            // initialized to `1 + mExtraDisplaysAssignedToUsers.size()`, which is usually larger
+            // than the actual length, because mExtraDisplaysAssignedToUsers contains displayIds for
+            // other users. Therefore, we need to copy to a new array with the correct length.
+            int[] results = new int[count];
+            System.arraycopy(displayIds, 0, results, 0, count);
+            return results;
+        }
+    }
+
     /**
      * See {@link UserManagerInternal#getUserAssignedToDisplay(int)}.
      */
diff --git a/services/tests/mockingservicestests/src/com/android/server/pm/UserVisibilityMediatorTestCase.java b/services/tests/mockingservicestests/src/com/android/server/pm/UserVisibilityMediatorTestCase.java
index 239b6fd..70b5ac0 100644
--- a/services/tests/mockingservicestests/src/com/android/server/pm/UserVisibilityMediatorTestCase.java
+++ b/services/tests/mockingservicestests/src/com/android/server/pm/UserVisibilityMediatorTestCase.java
@@ -596,6 +596,7 @@
                 .that(mMediator.assignUserToExtraDisplay(userId, displayId))
                 .isTrue();
         expectUserIsVisibleOnDisplay(userId, displayId);
+        expectDisplaysAssignedToUserContainsDisplayId(userId, displayId);
 
         if (unassign) {
             Log.d(TAG, "Calling unassignUserFromExtraDisplay(" + userId + ", " + displayId + ")");
@@ -603,6 +604,7 @@
                     .that(mMediator.unassignUserFromExtraDisplay(userId, displayId))
                     .isTrue();
             expectUserIsNotVisibleOnDisplay(userId, displayId);
+            expectDisplaysAssignedToUserDoesNotContainDisplayId(userId, displayId);
         }
     }
 
@@ -668,6 +670,7 @@
         expectUserIsNotVisibleOnDisplay(userId, INVALID_DISPLAY);
         expectUserIsNotVisibleOnDisplay(userId, SECONDARY_DISPLAY_ID);
         expectUserIsNotVisibleOnDisplay(userId, OTHER_SECONDARY_DISPLAY_ID);
+        expectDisplaysAssignedToUserIsEmpty(userId);
     }
 
     protected void expectDisplayAssignedToUser(@UserIdInt int userId, int displayId) {
@@ -680,6 +683,24 @@
                 .that(mMediator.getDisplayAssignedToUser(userId)).isEqualTo(INVALID_DISPLAY);
     }
 
+    protected void expectDisplaysAssignedToUserContainsDisplayId(
+            @UserIdInt int userId, int displayId) {
+        expectWithMessage("getDisplaysAssignedToUser(%s)", userId)
+                .that(mMediator.getDisplaysAssignedToUser(userId)).asList().contains(displayId);
+    }
+
+    protected void expectDisplaysAssignedToUserDoesNotContainDisplayId(
+            @UserIdInt int userId, int displayId) {
+        expectWithMessage("getDisplaysAssignedToUser(%s)", userId)
+                .that(mMediator.getDisplaysAssignedToUser(userId)).asList()
+                .doesNotContain(displayId);
+    }
+
+    protected void expectDisplaysAssignedToUserIsEmpty(@UserIdInt int userId) {
+        expectWithMessage("getDisplaysAssignedToUser(%s)", userId)
+                .that(mMediator.getDisplaysAssignedToUser(userId)).isNull();
+    }
+
     protected void expectUserCannotBeUnassignedFromDisplay(@UserIdInt int userId, int displayId) {
         expectWithMessage("unassignUserFromExtraDisplay(%s, %s)", userId, displayId)
                 .that(mMediator.unassignUserFromExtraDisplay(userId, displayId)).isFalse();