Fix selecting focused stack with secondary displays

When the last activity finishes in the stack we're looking for other
stacks and making it focused. However we weren't doing that if the
stack was on a secondary display, so the focused stack records were
not updated in stack supervisor.

Now we're looking for other stacks on the same display first. If there
is nothing focusable left - shifting focus to next focusable display.

Test: android.server.cts.ActivityManagerDisplayTests
Test: #testStackFocusSwitchOnDisplayRemoved
Test: #testStackFocusSwitchOnStackEmptied
Change-Id: Ifbb893e12cbe9c4928b949a86fc8bc027de181e4
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index ba95abd..6e7b34f 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -1392,24 +1392,6 @@
         return null;
     }
 
-    ActivityStack getNextFocusableStackLocked() {
-        ArrayList<ActivityStack> stacks = mStacks;
-        final ActivityRecord parent = mActivityContainer.mParentActivity;
-        if (parent != null) {
-            stacks = parent.getStack().mStacks;
-        }
-        if (stacks != null) {
-            for (int i = stacks.size() - 1; i >= 0; --i) {
-                ActivityStack stack = stacks.get(i);
-                if (stack != this && stack.isFocusable()
-                        && stack.getStackVisibilityLocked(null) != STACK_INVISIBLE) {
-                    return stack;
-                }
-            }
-        }
-        return null;
-    }
-
     /** Returns true if the stack contains a fullscreen task. */
     private boolean hasFullscreenTask() {
         for (int i = mTaskHistory.size() - 1; i >= 0; --i) {
@@ -2104,9 +2086,15 @@
             return false;
         }
 
-        ActivityRecord parent = mActivityContainer.mParentActivity;
-        if ((parent != null && parent.state != ActivityState.RESUMED) ||
-                !mActivityContainer.isAttachedLocked()) {
+        // Find the topmost activity in this stack that is not finishing.
+        final ActivityRecord next = topRunningActivityLocked();
+
+        final boolean hasRunningActivity = next != null;
+
+        final ActivityRecord parent = mActivityContainer.mParentActivity;
+        final boolean isParentNotResumed = parent != null && parent.state != ActivityState.RESUMED;
+        if (hasRunningActivity
+                && (isParentNotResumed || !mActivityContainer.isAttachedLocked())) {
             // Do not resume this stack if its parent is not resumed.
             // TODO: If in a loop, make sure that parent stack resumeTopActivity is called 1st.
             return false;
@@ -2114,33 +2102,14 @@
 
         mStackSupervisor.cancelInitializingActivities();
 
-        // Find the first activity that is not finishing.
-        final ActivityRecord next = topRunningActivityLocked();
-
         // Remember how we'll process this pause/resume situation, and ensure
         // that the state is reset however we wind up proceeding.
         final boolean userLeaving = mStackSupervisor.mUserLeaving;
         mStackSupervisor.mUserLeaving = false;
 
-        final TaskRecord prevTask = prev != null ? prev.task : null;
-        if (next == null) {
-            // There are no more activities!
-            final String reason = "noMoreActivities";
-            if (!mFullscreen && adjustFocusToNextFocusableStackLocked(reason)) {
-                // Try to move focus to the next visible stack with a running activity if this
-                // stack is not covering the entire screen.
-                return mStackSupervisor.resumeFocusedStackTopActivityLocked(
-                        mStackSupervisor.getFocusedStack(), prev, null);
-            }
-
-            // Let's just start up the Launcher...
-            ActivityOptions.abort(options);
-            if (DEBUG_STATES) Slog.d(TAG_STATES,
-                    "resumeTopActivityLocked: No more activities go home");
-            if (DEBUG_STACK) mStackSupervisor.validateTopActivitiesLocked();
-            // Only resume home if on home display
-            return isOnHomeDisplay() &&
-                    mStackSupervisor.resumeHomeStackTask(prev, reason);
+        if (!hasRunningActivity) {
+            // There are no activities left in the stack, let's look somewhere else.
+            return resumeTopActivityInNextFocusableStack(prev, options, "noMoreActivities");
         }
 
         next.delayedResume = false;
@@ -2158,6 +2127,7 @@
         }
 
         final TaskRecord nextTask = next.task;
+        final TaskRecord prevTask = prev != null ? prev.task : null;
         if (prevTask != null && prevTask.getStack() == this &&
                 prevTask.isOverHomeStack() && prev.finishing && prev.frontOfTask) {
             if (DEBUG_STACK)  mStackSupervisor.validateTopActivitiesLocked();
@@ -2533,6 +2503,27 @@
         return true;
     }
 
+    private boolean resumeTopActivityInNextFocusableStack(ActivityRecord prev,
+            ActivityOptions options, String reason) {
+        if ((!mFullscreen || !isOnHomeDisplay())
+                && adjustFocusToNextFocusableStackLocked(reason)) {
+            // Try to move focus to the next visible stack with a running activity if this
+            // stack is not covering the entire screen or is on a secondary display (with no home
+            // stack).
+            return mStackSupervisor.resumeFocusedStackTopActivityLocked(
+                    mStackSupervisor.getFocusedStack(), prev, null);
+        }
+
+        // Let's just start up the Launcher...
+        ActivityOptions.abort(options);
+        if (DEBUG_STATES) Slog.d(TAG_STATES,
+                "resumeTopActivityInNextFocusableStack: " + reason + ", go home");
+        if (DEBUG_STACK) mStackSupervisor.validateTopActivitiesLocked();
+        // Only resume home if on home display
+        return isOnHomeDisplay() &&
+                mStackSupervisor.resumeHomeStackTask(prev, reason);
+    }
+
     private TaskRecord getNextTask(TaskRecord targetTask) {
         final int index = mTaskHistory.indexOf(targetTask);
         if (index >= 0) {
@@ -3169,7 +3160,7 @@
     }
 
     private boolean adjustFocusToNextFocusableStackLocked(String reason) {
-        final ActivityStack stack = getNextFocusableStackLocked();
+        final ActivityStack stack = mStackSupervisor.getNextFocusableStackLocked(this);
         final String myReason = reason + " adjustFocusToNextFocusableStack";
         if (stack == null) {
             return false;
diff --git a/services/core/java/com/android/server/am/ActivityStackSupervisor.java b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
index 235325b..6987f94 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -457,6 +457,13 @@
     private final FindTaskResult mTmpFindTaskResult = new FindTaskResult();
 
     /**
+     * Temp storage for display ids sorted in focus order.
+     * Maps position to id. Using {@link SparseIntArray} instead of {@link ArrayList} because
+     * it's more efficient, as the number of displays is usually small.
+     */
+    private SparseIntArray mTmpOrderedDisplayIds = new SparseIntArray();
+
+    /**
      * Used to keep track whether app visibilities got changed since the last pause. Useful to
      * determine whether to invoke the task stack change listener after pausing.
      */
@@ -639,7 +646,7 @@
     void setFocusStackUnchecked(String reason, ActivityStack focusCandidate) {
         if (!focusCandidate.isFocusable()) {
             // The focus candidate isn't focusable. Move focus to the top stack that is focusable.
-            focusCandidate = focusCandidate.getNextFocusableStackLocked();
+            focusCandidate = getNextFocusableStackLocked(focusCandidate);
         }
 
         if (focusCandidate != mFocusedStack) {
@@ -2021,6 +2028,35 @@
         return mActivityDisplays.valueAt(DEFAULT_DISPLAY).mStacks;
     }
 
+    /**
+     * Get next focusable stack in the system. This will search across displays and stacks
+     * in last-focused order for a focusable and visible stack, different from the target stack.
+     *
+     * @param currentFocus The stack that previously had focus and thus needs to be ignored when
+     *                     searching for next candidate.
+     * @return Next focusable {@link ActivityStack}, null if not found.
+     */
+    ActivityStack getNextFocusableStackLocked(ActivityStack currentFocus) {
+        mWindowManager.getDisplaysInFocusOrder(mTmpOrderedDisplayIds);
+
+        for (int i = mTmpOrderedDisplayIds.size() - 1; i >= 0; --i) {
+            final int displayId = mTmpOrderedDisplayIds.get(i);
+            final List<ActivityStack> stacks = mActivityDisplays.get(displayId).mStacks;
+            if (stacks == null) {
+                continue;
+            }
+            for (int j = stacks.size() - 1; j >= 0; --j) {
+                final ActivityStack stack = stacks.get(j);
+                if (stack != currentFocus && stack.isFocusable()
+                        && stack.getStackVisibilityLocked(null) != STACK_INVISIBLE) {
+                    return stack;
+                }
+            }
+        }
+
+        return null;
+    }
+
     ActivityRecord getHomeActivity() {
         return getHomeActivityForUser(mCurrentUser);
     }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 349740b..dc06d12 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -166,6 +166,19 @@
     }
 
     /**
+     * Get an array with display ids ordered by focus priority - last items should be given
+     * focus first. Sparse array just maps position to displayId.
+     */
+    void getDisplaysInFocusOrder(SparseIntArray displaysInFocusOrder) {
+        displaysInFocusOrder.clear();
+
+        final int size = mChildren.size();
+        for (int i = 0; i < size; ++i) {
+            displaysInFocusOrder.put(i, mChildren.get(i).getDisplayId());
+        }
+    }
+
+    /**
      * Retrieve the DisplayContent for the specified displayId. Will create a new DisplayContent if
      * there is a Display for the displayId.
      *
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 0f16750..011992f 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -7151,6 +7151,16 @@
                 displayInfo.overscanRight, displayInfo.overscanBottom);
     }
 
+    /**
+     * Get an array with display ids ordered by focus priority - last items should be given
+     * focus first. Sparse array just maps position to displayId.
+     */
+    public void getDisplaysInFocusOrder(SparseIntArray displaysInFocusOrder) {
+        synchronized(mWindowMap) {
+            mRoot.getDisplaysInFocusOrder(displaysInFocusOrder);
+        }
+    }
+
     @Override
     public void setOverscan(int displayId, int left, int top, int right, int bottom) {
         if (mContext.checkCallingOrSelfPermission(