Track focus changes on external displays (2/4)

Let each DisplayContent has its own focused window and focused app.
This change also moves the last tapped display to the top.

Test: atest ActivityManagerMultiDisplayTests
            ActivityStackSupervisorTests
            ActivityStackTests
            CtsWindowManagerDeviceTestCases
            DisplayContentTests
            PointerCaptureTest
Bug: 111361570
Change-Id: I776cabaeaf41ff4240f504fb1430d3e40892023d
diff --git a/services/art-profile b/services/art-profile
index 3c60eee..328f8f7 100644
--- a/services/art-profile
+++ b/services/art-profile
@@ -2254,8 +2254,8 @@
 HPLcom/android/server/wm/DisplayContent;->lambda$new$8(Lcom/android/server/wm/DisplayContent;Lcom/android/server/wm/WindowState;)V
 HPLcom/android/server/wm/DisplayContent;->prepareSurfaces()V
 HPLcom/android/server/wm/DisplayContent;->resetAnimationBackgroundAnimator()V
-HPLcom/android/server/wm/DisplayContent;->setTouchExcludeRegion(Lcom/android/server/wm/Task;)V
 HPLcom/android/server/wm/DisplayContent;->skipTraverseChild(Lcom/android/server/wm/WindowContainer;)Z
+HPLcom/android/server/wm/DisplayContent;->updateTouchExcludeRegion()V
 HPLcom/android/server/wm/DockedStackDividerController;->isResizing()Z
 HPLcom/android/server/wm/DragDropController;->dragDropActiveLocked()Z
 HPLcom/android/server/wm/InputMonitor$UpdateInputForAllWindowsConsumer;->accept(Lcom/android/server/wm/WindowState;)V
diff --git a/services/core/java/com/android/server/am/ActivityDisplay.java b/services/core/java/com/android/server/am/ActivityDisplay.java
index fab967c..9d981f8 100644
--- a/services/core/java/com/android/server/am/ActivityDisplay.java
+++ b/services/core/java/com/android/server/am/ActivityDisplay.java
@@ -920,6 +920,10 @@
                 && (mSupervisor.mService.mRunningVoice == null);
     }
 
+    void setFocusedApp(ActivityRecord r, boolean moveFocusNow) {
+        mWindowContainerController.setFocusedApp(r.appToken, moveFocusNow);
+    }
+
     /**
      * @return the stack currently above the {@param stack}.  Can be null if the {@param stack} is
      *         already top-most.
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index 29b04cc..a06ebb6c 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -5434,7 +5434,7 @@
 
         // Do not sleep activities in this stack if we're marked as focused and the keyguard
         // is in the process of going away.
-        if (mStackSupervisor.getTopDisplayFocusedStack() == this
+        if (isFocusedStackOnDisplay()
                 && mStackSupervisor.getKeyguardController().isKeyguardGoingAway()) {
             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 9688d26..98a0ea7 100644
--- a/services/core/java/com/android/server/am/ActivityStackSupervisor.java
+++ b/services/core/java/com/android/server/am/ActivityStackSupervisor.java
@@ -3588,7 +3588,7 @@
                     stack.goToSleepIfPossible(false /* shuttingDown */);
                 } else {
                     stack.awakeFromSleepingLocked();
-                    if (isTopDisplayFocusedStack(stack) && !getKeyguardController()
+                    if (stack.isFocusedStackOnDisplay() && !getKeyguardController()
                             .isKeyguardOrAodShowing(display.mDisplayId)) {
                         // If the keyguard is unlocked - resume immediately.
                         // It is possible that the display will not be awake at the time we
diff --git a/services/core/java/com/android/server/am/ActivityTaskManagerService.java b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
index 36261b5..b4f328b 100644
--- a/services/core/java/com/android/server/am/ActivityTaskManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityTaskManagerService.java
@@ -4846,8 +4846,7 @@
         updateResumedAppTrace(r);
         mLastResumedActivity = r;
 
-        // TODO(b/111361570): Support multiple focused apps in WM
-        mWindowManager.setFocusedApp(r.appToken, true);
+        r.getDisplay().setFocusedApp(r, true);
 
         applyUpdateLockStateLocked(r);
         applyUpdateVrModeLocked(r);
diff --git a/services/core/java/com/android/server/input/InputManagerService.java b/services/core/java/com/android/server/input/InputManagerService.java
index 5bd095d..4913e8b 100644
--- a/services/core/java/com/android/server/input/InputManagerService.java
+++ b/services/core/java/com/android/server/input/InputManagerService.java
@@ -209,7 +209,8 @@
     private static native void nativeSetInputDispatchMode(long ptr, boolean enabled, boolean frozen);
     private static native void nativeSetSystemUiVisibility(long ptr, int visibility);
     private static native void nativeSetFocusedApplication(long ptr,
-            InputApplicationHandle application);
+            int displayId, InputApplicationHandle application);
+    private static native void nativeSetFocusedDisplay(long ptr, int displayId);
     private static native boolean nativeTransferTouchFocus(long ptr,
             InputChannel fromChannel, InputChannel toChannel);
     private static native void nativeSetPointerSpeed(long ptr, int speed);
@@ -1431,21 +1432,27 @@
         }
     }
 
-    public void setInputWindows(InputWindowHandle[] windowHandles,
-            InputWindowHandle focusedWindowHandle, int displayId) {
-        final IWindow newFocusedWindow =
-            focusedWindowHandle != null ? focusedWindowHandle.clientWindow : null;
-        if (mFocusedWindow != newFocusedWindow) {
-            mFocusedWindow = newFocusedWindow;
-            if (mFocusedWindowHasCapture) {
-                setPointerCapture(false);
-            }
-        }
+    public void setInputWindows(InputWindowHandle[] windowHandles, int displayId) {
         nativeSetInputWindows(mPtr, windowHandles, displayId);
     }
 
-    public void setFocusedApplication(InputApplicationHandle application) {
-        nativeSetFocusedApplication(mPtr, application);
+    public void setFocusedApplication(int displayId, InputApplicationHandle application) {
+        nativeSetFocusedApplication(mPtr, displayId, application);
+    }
+
+    public void setFocusedWindow(InputWindowHandle focusedWindowHandle) {
+        final IWindow newFocusedWindow =
+            focusedWindowHandle != null ? focusedWindowHandle.clientWindow : null;
+        if (mFocusedWindow != newFocusedWindow) {
+            if (mFocusedWindowHasCapture) {
+                setPointerCapture(false);
+            }
+            mFocusedWindow = newFocusedWindow;
+        }
+    }
+
+    public void setFocusedDisplay(int displayId) {
+        nativeSetFocusedDisplay(mPtr, displayId);
     }
 
     @Override
@@ -1460,16 +1467,18 @@
             return;
         }
         setPointerCapture(enabled);
-        try {
-            mFocusedWindow.dispatchPointerCaptureChanged(enabled);
-        } catch (RemoteException ex) {
-            /* ignore */
-        }
     }
 
     private void setPointerCapture(boolean enabled) {
-        mFocusedWindowHasCapture = enabled;
-        nativeSetPointerCapture(mPtr, enabled);
+        if (mFocusedWindowHasCapture != enabled) {
+            mFocusedWindowHasCapture = enabled;
+            try {
+                mFocusedWindow.dispatchPointerCaptureChanged(enabled);
+            } catch (RemoteException ex) {
+                /* ignore */
+            }
+            nativeSetPointerCapture(mPtr, enabled);
+        }
     }
 
     public void setInputDispatchMode(boolean enabled, boolean frozen) {
diff --git a/services/core/java/com/android/server/policy/PhoneWindowManager.java b/services/core/java/com/android/server/policy/PhoneWindowManager.java
index 2557f46..b84779c 100644
--- a/services/core/java/com/android/server/policy/PhoneWindowManager.java
+++ b/services/core/java/com/android/server/policy/PhoneWindowManager.java
@@ -4777,6 +4777,7 @@
             pf.bottom = df.bottom = of.bottom = displayFrames.mUnrestricted.bottom;
             // ...with content insets above the nav bar
             cf.bottom = vf.bottom = displayFrames.mStable.bottom;
+            // TODO (b/111364446): Support showing IME on non-default displays
             if (mStatusBar != null && mFocusedWindow == mStatusBar && canReceiveInput(mStatusBar)) {
                 // The status bar forces the navigation bar while it's visible. Make sure the IME
                 // avoids the navigation bar in that case.
diff --git a/services/core/java/com/android/server/wm/AccessibilityController.java b/services/core/java/com/android/server/wm/AccessibilityController.java
index e449111..74922f6 100644
--- a/services/core/java/com/android/server/wm/AccessibilityController.java
+++ b/services/core/java/com/android/server/wm/AccessibilityController.java
@@ -1043,7 +1043,8 @@
                 // Do not send the windows if there is no current focus as
                 // the window manager is still looking for where to put it.
                 // We will do the work when we get a focus change callback.
-                if (mService.mCurrentFocus == null) {
+                // TODO(b/112273690): Support multiple displays
+                if (mService.getDefaultDisplayContentLocked().mCurrentFocus == null) {
                     return;
                 }
 
diff --git a/services/core/java/com/android/server/wm/AppWindowContainerController.java b/services/core/java/com/android/server/wm/AppWindowContainerController.java
index 36280dd..330c54c 100644
--- a/services/core/java/com/android/server/wm/AppWindowContainerController.java
+++ b/services/core/java/com/android/server/wm/AppWindowContainerController.java
@@ -407,8 +407,7 @@
                 if (mService.mAppTransition.getAppTransition()
                         == WindowManager.TRANSIT_TASK_OPEN_BEHIND) {
                     // We're launchingBehind, add the launching activity to mOpeningApps.
-                    final WindowState win =
-                            mService.getDefaultDisplayContentLocked().findFocusedWindow();
+                    final WindowState win = mContainer.getDisplayContent().findFocusedWindow();
                     if (win != null) {
                         final AppWindowToken focusedToken = win.mAppToken;
                         if (focusedToken != null) {
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index fc76102..e57fea3 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -679,11 +679,12 @@
         removed = true;
         stopFreezingScreen(true, true);
 
-        if (mService.mFocusedApp == this) {
-            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + this);
-            mService.mFocusedApp = null;
+        final DisplayContent dc = getDisplayContent();
+        if (dc.mFocusedApp == this) {
+            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + this
+                   + " displayId=" + dc.getDisplayId());
+            dc.setFocusedApp(null);
             mService.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, true /*updateInputWindows*/);
-            getDisplayContent().getInputMonitor().setFocusedAppLw(null);
         }
 
         if (!delayed) {
@@ -1064,6 +1065,18 @@
         }
     }
 
+    @Override
+    void onDisplayChanged(DisplayContent dc) {
+        DisplayContent prevDc = mDisplayContent;
+        super.onDisplayChanged(dc);
+        if (prevDc != null && prevDc.mFocusedApp == this) {
+            prevDc.setFocusedApp(null);
+            if (dc.getTopStack().getTopChild().getTopChild() == this) {
+                dc.setFocusedApp(this);
+            }
+        }
+    }
+
     /**
      * Freezes the task bounds. The size of this task reported the app will be fixed to the bounds
      * freezed by {@link Task#prepareFreezingBounds} until {@link #unfreezeBounds} gets called, even
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index a762fe9..730e422 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -70,6 +70,7 @@
 import static com.android.server.wm.DisplayContentProto.DISPLAY_INFO;
 import static com.android.server.wm.DisplayContentProto.DOCKED_STACK_DIVIDER_CONTROLLER;
 import static com.android.server.wm.DisplayContentProto.DPI;
+import static com.android.server.wm.DisplayContentProto.FOCUSED_APP;
 import static com.android.server.wm.DisplayContentProto.ID;
 import static com.android.server.wm.DisplayContentProto.IME_WINDOWS;
 import static com.android.server.wm.DisplayContentProto.PINNED_STACK_CONTROLLER;
@@ -98,12 +99,17 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowManagerService.CUSTOM_SCREEN_ROTATION;
+import static com.android.server.wm.WindowManagerService.H.REPORT_FOCUS_CHANGE;
+import static com.android.server.wm.WindowManagerService.H.REPORT_LOSING_FOCUS;
 import static com.android.server.wm.WindowManagerService.H.SEND_NEW_CONFIGURATION;
 import static com.android.server.wm.WindowManagerService.H.UPDATE_DOCKED_STACK_DIVIDER;
 import static com.android.server.wm.WindowManagerService.H.WINDOW_HIDE_TIMEOUT;
 import static com.android.server.wm.WindowManagerService.LAYOUT_REPEAT_THRESHOLD;
 import static com.android.server.wm.WindowManagerService.MAX_ANIMATION_DURATION;
 import static com.android.server.wm.WindowManagerService.SEAMLESS_ROTATION_TIMEOUT_DURATION;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_PLACING_SURFACES;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_REMOVING_FOCUS;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_WILL_ASSIGN_LAYERS;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_WILL_PLACE_SURFACES;
 import static com.android.server.wm.WindowManagerService.WINDOWS_FREEZING_SCREENS_ACTIVE;
 import static com.android.server.wm.WindowManagerService.WINDOWS_FREEZING_SCREENS_TIMEOUT;
@@ -371,6 +377,36 @@
     private final SurfaceSession mSession = new SurfaceSession();
 
     /**
+     * Window that is currently interacting with the user. This window is responsible for receiving
+     * key events and pointer events from the user.
+     */
+    WindowState mCurrentFocus = null;
+
+    /**
+     * The last focused window that we've notified the client that the focus is changed.
+     */
+    WindowState mLastFocus = null;
+
+    /**
+     * Windows that have lost input focus and are waiting for the new focus window to be displayed
+     * before they are told about this.
+     */
+    ArrayList<WindowState> mLosingFocus = new ArrayList<>();
+
+    /**
+     * The foreground app of this display. Windows below this app cannot be the focused window. If
+     * the user taps on the area outside of the task of the focused app, we will notify AM about the
+     * new task the user wants to interact with.
+     */
+    AppWindowToken mFocusedApp = null;
+
+    /** Windows added since {@link #mCurrentFocus} was set to null. Used for ANR blaming. */
+    final ArrayList<WindowState> mWinAddedSinceNullFocus = new ArrayList<>();
+
+    /** Windows removed since {@link #mCurrentFocus} was set to null. Used for ANR blaming. */
+    final ArrayList<WindowState> mWinRemovedSinceNullFocus = new ArrayList<>();
+
+    /**
      * We organize all top-level Surfaces in to the following layers.
      * mOverlayLayer contains a few Surfaces which are always on top of others
      * and omitted from Screen-Magnification, for example the strict mode flash or
@@ -466,7 +502,7 @@
     };
 
     private final ToBooleanFunction<WindowState> mFindFocusedWindow = w -> {
-        final AppWindowToken focusedApp = mService.mFocusedApp;
+        final AppWindowToken focusedApp = mFocusedApp;
         if (DEBUG_FOCUS) Slog.v(TAG_WM, "Looking for focus: " + w
                 + ", flags=" + w.mAttrs.flags + ", canReceive=" + w.canReceiveKeys());
 
@@ -625,8 +661,6 @@
         final boolean obscuredChanged = w.mObscured !=
                 mTmpApplySurfaceChangesTransactionState.obscured;
         final RootWindowContainer root = mService.mRoot;
-        // Only used if default window
-        final boolean someoneLosingFocus = !mService.mLosingFocus.isEmpty();
 
         // Update effect.
         w.mObscured = mTmpApplySurfaceChangesTransactionState.obscured;
@@ -717,9 +751,8 @@
             }
         }
 
-        if (isDefaultDisplay && someoneLosingFocus && w == mService.mCurrentFocus
-                && w.isDisplayedLw()) {
-            mTmpApplySurfaceChangesTransactionState.focusDisplayed = true;
+        if (!mLosingFocus.isEmpty() && w.isFocused() && w.isDisplayedLw()) {
+            mService.mH.obtainMessage(REPORT_LOSING_FOCUS, this).sendToTarget();
         }
 
         w.updateResizingWindowIfNeeded();
@@ -2070,22 +2103,22 @@
         return null;
     }
 
-    void setTouchExcludeRegion(Task focusedTask) {
-        // The provided task is the task on this display with focus, so if WindowManagerService's
-        // focused app is not on this display, focusedTask will be null.
+    void updateTouchExcludeRegion() {
+        final Task focusedTask = (mFocusedApp != null ? mFocusedApp.getTask() : null);
         if (focusedTask == null) {
             mTouchExcludeRegion.setEmpty();
         } else {
             mTouchExcludeRegion.set(mBaseDisplayRect);
             final int delta = dipToPixel(RESIZE_HANDLE_WIDTH_IN_DP, mDisplayMetrics);
             mTmpRect2.setEmpty();
-            for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
+            for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0;
+                    --stackNdx) {
                 final TaskStack stack = mTaskStackContainers.getChildAt(stackNdx);
                 stack.setTouchExcludeRegion(focusedTask, delta, mTouchExcludeRegion,
                         mDisplayFrames.mContent, mTmpRect2);
             }
             // If we removed the focused task above, add it back and only leave its
-            // outside touch area in the exclusion. TapDectector is not interested in
+            // outside touch area in the exclusion. TapDetector is not interested in
             // any touch inside the focused task itself.
             if (!mTmpRect2.isEmpty()) {
                 mTouchExcludeRegion.op(mTmpRect2, Region.Op.UNION);
@@ -2096,12 +2129,7 @@
             // events to be intercepted and used to change focus. This would likely cause a
             // disappearance of the input method.
             mInputMethodWindow.getTouchableRegion(mTmpRegion);
-            if (mInputMethodWindow.getDisplayId() == mDisplayId) {
-                mTouchExcludeRegion.op(mTmpRegion, Op.UNION);
-            } else {
-                // IME is on a different display, so we need to update its tap detector.
-                setTouchExcludeRegion(null /* focusedTask */);
-            }
+            mTouchExcludeRegion.op(mTmpRegion, Op.UNION);
         }
         for (int i = mTapExcludedWindows.size() - 1; i >= 0; i--) {
             final WindowState win = mTapExcludedWindows.get(i);
@@ -2372,6 +2400,9 @@
         }
         mDisplayFrames.writeToProto(proto, DISPLAY_FRAMES);
         proto.write(SURFACE_SIZE, mSurfaceSize);
+        if (mFocusedApp != null) {
+            mFocusedApp.writeNameToProto(proto, FOCUSED_APP);
+        }
         proto.end(token);
     }
 
@@ -2412,6 +2443,27 @@
         pw.print(prefix);
         pw.print("mDeferredRotationPauseCount="); pw.println(mDeferredRotationPauseCount);
 
+        pw.print("  mCurrentFocus="); pw.println(mCurrentFocus);
+        if (mLastFocus != mCurrentFocus) {
+            pw.print("  mLastFocus="); pw.println(mLastFocus);
+        }
+        if (mLosingFocus.size() > 0) {
+            pw.println();
+            pw.println("  Windows losing focus:");
+            for (int i = mLosingFocus.size() - 1; i >= 0; i--) {
+                final WindowState w = mLosingFocus.get(i);
+                pw.print("  Losing #"); pw.print(i); pw.print(' ');
+                pw.print(w);
+                if (dumpAll) {
+                    pw.println(":");
+                    w.dump(pw, "    ", true);
+                } else {
+                    pw.println();
+                }
+            }
+        }
+        pw.print("  mFocusedApp="); pw.println(mFocusedApp);
+
         pw.println();
         pw.println(prefix + "Application tokens in top down Z order:");
         for (int stackNdx = mTaskStackContainers.getChildCount() - 1; stackNdx >= 0; --stackNdx) {
@@ -2543,6 +2595,132 @@
         return mTmpWindow;
     }
 
+
+    /**
+     * Update the focused window and make some adjustments if the focus has changed.
+     *
+     * @param mode Indicates the situation we are in. Possible modes are:
+     *             {@link WindowManagerService#UPDATE_FOCUS_NORMAL},
+     *             {@link WindowManagerService#UPDATE_FOCUS_PLACING_SURFACES},
+     *             {@link WindowManagerService#UPDATE_FOCUS_WILL_PLACE_SURFACES},
+     *             {@link WindowManagerService#UPDATE_FOCUS_REMOVING_FOCUS}
+     * @param updateInputWindows Whether to sync the window information to the input module.
+     * @return {@code true} if the focused window has changed.
+     */
+    boolean updateFocusedWindowLocked(int mode, boolean updateInputWindows, boolean focusFound) {
+        final WindowState newFocus = findFocusedWindow();
+        if (mCurrentFocus == newFocus) {
+            return false;
+        }
+        mService.mH.obtainMessage(REPORT_FOCUS_CHANGE, this).sendToTarget();
+        boolean imWindowChanged = false;
+        // TODO (b/111080190): Multi-Session IME
+        if (!focusFound) {
+            final WindowState imWindow = mInputMethodWindow;
+            if (imWindow != null) {
+                final WindowState prevTarget = mService.mInputMethodTarget;
+
+                final WindowState newTarget = computeImeTarget(true /* updateImeTarget*/);
+                imWindowChanged = prevTarget != newTarget;
+
+                if (mode != UPDATE_FOCUS_WILL_ASSIGN_LAYERS
+                        && mode != UPDATE_FOCUS_WILL_PLACE_SURFACES) {
+                    assignWindowLayers(false /* setLayoutNeeded */);
+                }
+            }
+        }
+
+        if (imWindowChanged) {
+            mService.mWindowsChanged = true;
+            setLayoutNeeded();
+        }
+
+        if (DEBUG_FOCUS_LIGHT || mService.localLOGV) Slog.v(TAG_WM, "Changing focus from "
+                + mCurrentFocus + " to " + newFocus + " displayId=" + getDisplayId()
+                + " Callers=" + Debug.getCallers(4));
+        final WindowState oldFocus = mCurrentFocus;
+        mCurrentFocus = newFocus;
+        mLosingFocus.remove(newFocus);
+
+        if (newFocus != null) {
+            mWinAddedSinceNullFocus.clear();
+            mWinRemovedSinceNullFocus.clear();
+
+            if (newFocus.canReceiveKeys()) {
+                // Displaying a window implicitly causes dispatching to be unpaused.
+                // This is to protect against bugs if someone pauses dispatching but
+                // forgets to resume.
+                newFocus.mToken.paused = false;
+            }
+        }
+
+        // System UI is only shown on the default display.
+        int focusChanged = isDefaultDisplay
+                ? mService.mPolicy.focusChangedLw(oldFocus, newFocus) : 0;
+
+        if (imWindowChanged && oldFocus != mInputMethodWindow) {
+            // Focus of the input method window changed. Perform layout if needed.
+            if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
+                performLayout(true /*initial*/,  updateInputWindows);
+                focusChanged &= ~FINISH_LAYOUT_REDO_LAYOUT;
+            } else if (mode == UPDATE_FOCUS_WILL_PLACE_SURFACES) {
+                // Client will do the layout, but we need to assign layers
+                // for handleNewWindowLocked() below.
+                assignWindowLayers(false /* setLayoutNeeded */);
+            }
+        }
+
+        if ((focusChanged & FINISH_LAYOUT_REDO_LAYOUT) != 0) {
+            // The change in focus caused us to need to do a layout.  Okay.
+            setLayoutNeeded();
+            if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
+                performLayout(true /*initial*/, updateInputWindows);
+            } else if (mode == UPDATE_FOCUS_REMOVING_FOCUS) {
+                mService.mRoot.performSurfacePlacement(false);
+            }
+        }
+
+        if (mode != UPDATE_FOCUS_WILL_ASSIGN_LAYERS) {
+            // If we defer assigning layers, then the caller is responsible for doing this part.
+            getInputMonitor().setInputFocusLw(newFocus, updateInputWindows);
+        }
+
+        adjustForImeIfNeeded();
+
+        // We may need to schedule some toast windows to be removed. The toasts for an app that
+        // does not have input focus are removed within a timeout to prevent apps to redress
+        // other apps' UI.
+        scheduleToastWindowsTimeoutIfNeededLocked(oldFocus, newFocus);
+
+        if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
+            pendingLayoutChanges |= FINISH_LAYOUT_REDO_ANIM;
+        }
+        return true;
+    }
+
+    /**
+     * Set the new focused app to this display.
+     *
+     * @param newFocus the new focused AppWindowToken.
+     * @return true if the focused app is changed.
+     */
+    boolean setFocusedApp(AppWindowToken newFocus) {
+        if (newFocus != null) {
+            final DisplayContent appDisplay = newFocus.getDisplayContent();
+            if (appDisplay != this) {
+                throw new IllegalStateException(newFocus + " is not on " + getName()
+                        + " but " + ((appDisplay != null) ? appDisplay.getName() : "none"));
+            }
+        }
+        if (mFocusedApp == newFocus) {
+            return false;
+        }
+        mFocusedApp = newFocus;
+        getInputMonitor().setFocusedAppLw(newFocus);
+        updateTouchExcludeRegion();
+        return true;
+    }
+
     /** Updates the layer assignment of windows on this display. */
     void assignWindowLayers(boolean setLayoutNeeded) {
         Trace.traceBegin(Trace.TRACE_TAG_WINDOW_MANAGER, "assignWindowLayers");
@@ -2906,8 +3084,8 @@
 
         if (DEBUG_INPUT_METHOD) {
             Slog.i(TAG_WM, "Desired input method target: " + imFocus);
-            Slog.i(TAG_WM, "Current focus: " + mService.mCurrentFocus);
-            Slog.i(TAG_WM, "Last focus: " + mService.mLastFocus);
+            Slog.i(TAG_WM, "Current focus: " + mCurrentFocus + " displayId=" + mDisplayId);
+            Slog.i(TAG_WM, "Last focus: " + mLastFocus + " displayId=" + mDisplayId);
         }
 
         if (DEBUG_INPUT_METHOD) {
@@ -2975,7 +3153,7 @@
     }
 
     // TODO: Super crazy long method that should be broken down...
-    boolean applySurfaceChangesTransaction(boolean recoveringMemory) {
+    void applySurfaceChangesTransaction(boolean recoveringMemory) {
 
         final int dw = mDisplayInfo.logicalWidth;
         final int dh = mDisplayInfo.logicalHeight;
@@ -3059,8 +3237,6 @@
             // can now be shown.
             atoken.updateAllDrawn();
         }
-
-        return mTmpApplySurfaceChangesTransactionState.focusDisplayed;
     }
 
     private void updateBounds() {
@@ -3314,7 +3490,6 @@
         boolean displayHasContent;
         boolean obscured;
         boolean syswin;
-        boolean focusDisplayed;
         float preferredRefreshRate;
         int preferredModeId;
 
@@ -3322,7 +3497,6 @@
             displayHasContent = false;
             obscured = false;
             syswin = false;
-            focusDisplayed = false;
             preferredRefreshRate = 0;
             preferredModeId = 0;
         }
diff --git a/services/core/java/com/android/server/wm/DisplayWindowController.java b/services/core/java/com/android/server/wm/DisplayWindowController.java
index 76b6dbe..ab87759 100644
--- a/services/core/java/com/android/server/wm/DisplayWindowController.java
+++ b/services/core/java/com/android/server/wm/DisplayWindowController.java
@@ -17,11 +17,14 @@
 package com.android.server.wm;
 
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_DISPLAY;
+import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_FOCUS_LIGHT;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_STACK;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
 
 import android.content.res.Configuration;
 import android.os.Binder;
+import android.os.IBinder;
 import android.util.Slog;
 import android.view.Display;
 
@@ -124,6 +127,42 @@
         }
     }
 
+    /**
+     * Sets a focused app on this display.
+     *
+     * @param token Specifies which app should be focused.
+     * @param moveFocusNow Specifies if we should update the focused window immediately.
+     */
+    public void setFocusedApp(IBinder token, boolean moveFocusNow) {
+        synchronized (mWindowMap) {
+            if (mContainer == null) {
+                if (DEBUG_FOCUS_LIGHT) Slog.i(TAG_WM, "setFocusedApp: could not find displayId="
+                        + mDisplayId);
+                return;
+            }
+            final AppWindowToken newFocus;
+            if (token == null) {
+                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Clearing focused app, displayId="
+                        + mDisplayId);
+                newFocus = null;
+            } else {
+                newFocus = mRoot.getAppWindowToken(token);
+                if (newFocus == null) {
+                    Slog.w(TAG_WM, "Attempted to set focus to non-existing app token: " + token
+                            + ", displayId=" + mDisplayId);
+                }
+                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Set focused app to: " + newFocus
+                        + " moveFocusNow=" + moveFocusNow + " displayId=" + mDisplayId);
+            }
+
+            final boolean changed = mContainer.setFocusedApp(newFocus);
+            if (moveFocusNow && changed) {
+                mService.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL,
+                        true /*updateInputWindows*/);
+            }
+        }
+    }
+
     @Override
     public String toString() {
         return "{DisplayWindowController displayId=" + mDisplayId + "}";
diff --git a/services/core/java/com/android/server/wm/InputMonitor.java b/services/core/java/com/android/server/wm/InputMonitor.java
index ef3a770..15f6938 100644
--- a/services/core/java/com/android/server/wm/InputMonitor.java
+++ b/services/core/java/com/android/server/wm/InputMonitor.java
@@ -64,7 +64,6 @@
     // Array of window handles to provide to the input dispatcher.
     private InputWindowHandle[] mInputWindowHandles;
     private int mInputWindowHandleCount;
-    private InputWindowHandle mFocusedInputWindowHandle;
 
     private boolean mDisableWallpaperTouchEvents;
     private final Rect mTmpRect = new Rect();
@@ -229,16 +228,12 @@
                     + child + ", " + inputWindowHandle);
         }
         addInputWindowHandle(inputWindowHandle);
-        if (hasFocus) {
-            mFocusedInputWindowHandle = inputWindowHandle;
-        }
     }
 
     private void clearInputWindowHandlesLw() {
         while (mInputWindowHandleCount != 0) {
             mInputWindowHandles[--mInputWindowHandleCount] = null;
         }
-        mFocusedInputWindowHandle = null;
     }
 
     void setUpdateInputWindowsNeededLw() {
@@ -325,13 +320,13 @@
     public void setFocusedAppLw(AppWindowToken newApp) {
         // Focused app has changed.
         if (newApp == null) {
-            mService.mInputManager.setFocusedApplication(null);
+            mService.mInputManager.setFocusedApplication(mDisplayId, null);
         } else {
             final InputApplicationHandle handle = newApp.mInputApplicationHandle;
             handle.name = newApp.toString();
             handle.dispatchingTimeoutNanos = newApp.mInputDispatchingTimeoutNanos;
 
-            mService.mInputManager.setFocusedApplication(handle);
+            mService.mInputManager.setFocusedApplication(mDisplayId, handle);
         }
     }
 
@@ -370,8 +365,7 @@
     void onRemoved() {
         // If DisplayContent removed, we need find a way to remove window handles of this display
         // from InputDispatcher, so pass an empty InputWindowHandles to remove them.
-        mService.mInputManager.setInputWindows(mInputWindowHandles, mFocusedInputWindowHandle,
-                mDisplayId);
+        mService.mInputManager.setInputWindows(mInputWindowHandles, mDisplayId);
     }
 
     private final class UpdateInputForAllWindowsConsumer implements Consumer<WindowState> {
@@ -414,8 +408,7 @@
             }
 
             // Send windows to native code.
-            mService.mInputManager.setInputWindows(mInputWindowHandles, mFocusedInputWindowHandle,
-                    mDisplayId);
+            mService.mInputManager.setInputWindows(mInputWindowHandles, mDisplayId);
 
             clearInputWindowHandlesLw();
 
@@ -435,7 +428,6 @@
             final int flags = w.mAttrs.flags;
             final int privateFlags = w.mAttrs.privateFlags;
             final int type = w.mAttrs.type;
-            // TODO(b/111361570): multi-display focus, one focus window per display.
             final boolean hasFocus = w.isFocused();
             final boolean isVisible = w.isVisibleLw();
 
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index a9571be..3fef87d 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -17,19 +17,20 @@
 package com.android.server.wm;
 
 import static android.view.Display.DEFAULT_DISPLAY;
+import static android.view.Display.INVALID_DISPLAY;
 import static android.view.WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_KEYGUARD;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_SUSTAINED_PERFORMANCE_MODE;
 import static android.view.WindowManager.LayoutParams.TYPE_DREAM;
 import static android.view.WindowManager.LayoutParams.TYPE_KEYGUARD_DIALOG;
 
-import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_ANIM;
 import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_LAYOUT;
 import static com.android.server.policy.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static com.android.server.wm.RootWindowContainerProto.DISPLAYS;
 import static com.android.server.wm.RootWindowContainerProto.WINDOWS;
 import static com.android.server.wm.RootWindowContainerProto.WINDOW_CONTAINER;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_DISPLAY;
+import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_FOCUS_LIGHT;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_KEEP_SCREEN_ON;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_LAYOUT_REPEATS;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ORIENTATION;
@@ -41,9 +42,9 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_KEEP_SCREEN_ON;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
-import static com.android.server.wm.WindowManagerService.H.REPORT_LOSING_FOCUS;
 import static com.android.server.wm.WindowManagerService.H.SEND_NEW_CONFIGURATION;
 import static com.android.server.wm.WindowManagerService.H.WINDOW_FREEZE_TIMEOUT;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_PLACING_SURFACES;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_WILL_PLACE_SURFACES;
 import static com.android.server.wm.WindowManagerService.WINDOWS_FREEZING_SCREENS_NONE;
@@ -124,6 +125,10 @@
 
     private String mCloseSystemDialogsReason;
 
+    // The ID of the display which is responsible for receiving display-unspecified key and pointer
+    // events.
+    private int mTopFocusedDisplayId = INVALID_DISPLAY;
+
     // Only a seperate transaction until we seperate the apply surface changes
     // transaction from the global transaction.
     private final SurfaceControl.Transaction mDisplayTransaction = new SurfaceControl.Transaction();
@@ -153,23 +158,40 @@
         mWallpaperController = new WallpaperController(mService);
     }
 
-    WindowState computeFocusedWindow() {
-        // While the keyguard is showing, we must focus anything besides the main display.
-        // Otherwise we risk input not going to the keyguard when the user expects it to.
-        final boolean forceDefaultDisplay = mService.isKeyguardShowingAndNotOccluded();
-
+    boolean updateFocusedWindowLocked(int mode, boolean updateInputWindows) {
+        boolean changed = false;
+        int topFocusedDisplayId = INVALID_DISPLAY;
         for (int i = mChildren.size() - 1; i >= 0; i--) {
             final DisplayContent dc = mChildren.get(i);
-            final WindowState win = dc.findFocusedWindow();
-            if (win != null) {
-                if (forceDefaultDisplay && !dc.isDefaultDisplay) {
-                    EventLog.writeEvent(0x534e4554, "71786287", win.mOwnerUid, "");
-                    continue;
-                }
-                return win;
+            changed |= dc.updateFocusedWindowLocked(mode, updateInputWindows,
+                    topFocusedDisplayId != INVALID_DISPLAY /* focusFound */);
+            if (topFocusedDisplayId == INVALID_DISPLAY && dc.mCurrentFocus != null) {
+                topFocusedDisplayId = dc.getDisplayId();
             }
         }
-        return null;
+        if (topFocusedDisplayId == INVALID_DISPLAY) {
+            topFocusedDisplayId = DEFAULT_DISPLAY;
+        }
+        if (mTopFocusedDisplayId != topFocusedDisplayId) {
+            mTopFocusedDisplayId = topFocusedDisplayId;
+            mService.mInputManager.setFocusedDisplay(topFocusedDisplayId);
+            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "New topFocusedDisplayId="
+                    + topFocusedDisplayId);
+        }
+        final WindowState topFocusedWindow = getTopFocusedDisplayContent().mCurrentFocus;
+        mService.mInputManager.setFocusedWindow(
+                topFocusedWindow != null ? topFocusedWindow.mInputWindowHandle : null);
+        return changed;
+    }
+
+    DisplayContent getTopFocusedDisplayContent() {
+        return getDisplayContent(mTopFocusedDisplayId == INVALID_DISPLAY
+                ? DEFAULT_DISPLAY : mTopFocusedDisplayId);
+    }
+
+    @Override
+    void onChildPositionChanged() {
+        mService.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, false /* updateInputWindows */);
     }
 
     DisplayContent getDisplayContent(int displayId) {
@@ -636,7 +658,6 @@
             if (mService.updateFocusedWindowLocked(UPDATE_FOCUS_PLACING_SURFACES,
                     false /*updateInputWindows*/)) {
                 updateInputWindowsNeeded = true;
-                defaultDisplay.pendingLayoutChanges |= FINISH_LAYOUT_REDO_ANIM;
             }
         }
 
@@ -646,7 +667,7 @@
                     defaultDisplay.pendingLayoutChanges);
         }
 
-        final ArraySet<DisplayContent> touchExcludeRegionUpdateDisplays = handleResizingWindows();
+        handleResizingWindows();
 
         if (DEBUG_ORIENTATION && mService.mDisplayFrozen) Slog.v(TAG,
                 "With display frozen, orientationChangeComplete=" + mOrientationChangeComplete);
@@ -765,17 +786,7 @@
                 dc.getInputMonitor().updateInputWindowsLw(false /*force*/);
             });
         }
-        mService.setFocusTaskRegionLocked(null);
-        if (touchExcludeRegionUpdateDisplays != null) {
-            final DisplayContent focusedDc = mService.mFocusedApp != null
-                    ? mService.mFocusedApp.getDisplayContent() : null;
-            for (DisplayContent dc : touchExcludeRegionUpdateDisplays) {
-                // The focused DisplayContent was recalcuated in setFocusTaskRegionLocked
-                if (focusedDc != dc) {
-                    dc.setTouchExcludeRegion(null /* focusedTask */);
-                }
-            }
-        }
+        forAllDisplays(DisplayContent::updateTouchExcludeRegion);
 
         // Check to see if we are now in a state where the screen should
         // be enabled, because the window obscured flags have changed.
@@ -808,16 +819,10 @@
                     mService.getDefaultDisplayRotation());
         }
 
-        boolean focusDisplayed = false;
-
         final int count = mChildren.size();
         for (int j = 0; j < count; ++j) {
             final DisplayContent dc = mChildren.get(j);
-            focusDisplayed |= dc.applySurfaceChangesTransaction(recoveringMemory);
-        }
-
-        if (focusDisplayed) {
-            mService.mH.sendEmptyMessage(REPORT_LOSING_FOCUS);
+            dc.applySurfaceChangesTransaction(recoveringMemory);
         }
 
         // Give the display manager a chance to adjust properties like display rotation if it needs
@@ -828,12 +833,8 @@
 
     /**
      * Handles resizing windows during surface placement.
-     *
-     * @return A set of any DisplayContent whose touch exclude region needs to be recalculated due
-     *         to a tap-exclude window resizing, or null if no such DisplayContents were found.
      */
-    private ArraySet<DisplayContent> handleResizingWindows() {
-        ArraySet<DisplayContent> touchExcludeRegionUpdateSet = null;
+    private void handleResizingWindows() {
         for (int i = mService.mResizingWindows.size() - 1; i >= 0; i--) {
             WindowState win = mService.mResizingWindows.get(i);
             if (win.mAppFreezing) {
@@ -842,15 +843,7 @@
             }
             win.reportResized();
             mService.mResizingWindows.remove(i);
-            if (WindowManagerService.excludeWindowTypeFromTapOutTask(win.mAttrs.type)) {
-                final DisplayContent dc = win.getDisplayContent();
-                if (touchExcludeRegionUpdateSet == null) {
-                    touchExcludeRegionUpdateSet = new ArraySet<>();
-                }
-                touchExcludeRegionUpdateSet.add(dc);
-            }
         }
-        return touchExcludeRegionUpdateSet;
     }
 
     /**
@@ -1004,6 +997,10 @@
         }
     }
 
+    void dumpTopFocusedDisplayId(PrintWriter pw) {
+        pw.print("  mTopFocusedDisplayId="); pw.println(mTopFocusedDisplayId);
+    }
+
     void dumpLayoutNeededDisplayIds(PrintWriter pw) {
         if (!isLayoutNeeded()) {
             return;
diff --git a/services/core/java/com/android/server/wm/TaskPositioningController.java b/services/core/java/com/android/server/wm/TaskPositioningController.java
index 33416f6..b7e37b2 100644
--- a/services/core/java/com/android/server/wm/TaskPositioningController.java
+++ b/services/core/java/com/android/server/wm/TaskPositioningController.java
@@ -131,9 +131,9 @@
         // of the app, it may not have focus since there might be other windows
         // on top (eg. a dialog window).
         WindowState transferFocusFromWin = win;
-        if (mService.mCurrentFocus != null && mService.mCurrentFocus != win
-                && mService.mCurrentFocus.mAppToken == win.mAppToken) {
-            transferFocusFromWin = mService.mCurrentFocus;
+        if (displayContent.mCurrentFocus != null && displayContent.mCurrentFocus != win
+                && displayContent.mCurrentFocus.mAppToken == win.mAppToken) {
+            transferFocusFromWin = displayContent.mCurrentFocus;
         }
         if (!mInputManager.transferTouchFocus(
                 transferFocusFromWin.mInputChannel, mTaskPositioner.mServerChannel)) {
diff --git a/services/core/java/com/android/server/wm/TaskTapPointerEventListener.java b/services/core/java/com/android/server/wm/TaskTapPointerEventListener.java
index f1e1592..52f8510 100644
--- a/services/core/java/com/android/server/wm/TaskTapPointerEventListener.java
+++ b/services/core/java/com/android/server/wm/TaskTapPointerEventListener.java
@@ -19,12 +19,12 @@
 import android.graphics.Rect;
 import android.graphics.Region;
 import android.hardware.input.InputManager;
+import android.os.Handler;
 import android.view.MotionEvent;
 import android.view.WindowManagerPolicyConstants.PointerEventListener;
 
 import com.android.server.wm.WindowManagerService.H;
 
-import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.PointerIcon.TYPE_NOT_SPECIFIED;
 import static android.view.PointerIcon.TYPE_HORIZONTAL_DOUBLE_ARROW;
 import static android.view.PointerIcon.TYPE_VERTICAL_DOUBLE_ARROW;
@@ -36,6 +36,8 @@
     final private Region mTouchExcludeRegion = new Region();
     private final WindowManagerService mService;
     private final DisplayContent mDisplayContent;
+    private final Handler mHandler;
+    private final Runnable mMoveDisplayToTop;
     private final Rect mTmpRect = new Rect();
     private int mPointerIconType = TYPE_NOT_SPECIFIED;
 
@@ -43,6 +45,13 @@
             DisplayContent displayContent) {
         mService = service;
         mDisplayContent = displayContent;
+        mHandler = new Handler(mService.mH.getLooper());
+        mMoveDisplayToTop = () -> {
+            synchronized (mService.mWindowMap) {
+                mDisplayContent.getParent().positionChildAt(WindowContainer.POSITION_TOP,
+                        mDisplayContent, true /* includingParents */);
+            }
+        };
     }
 
     @Override
@@ -61,6 +70,7 @@
                         mService.mTaskPositioningController.handleTapOutsideTask(
                                 mDisplayContent, x, y);
                     }
+                    mHandler.post(mMoveDisplayToTop);
                 }
             }
             break;
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 4883f97..46999a2 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -273,6 +273,7 @@
             parent.mTreeWeight += child.mTreeWeight;
             parent = parent.getParent();
         }
+        onChildPositionChanged();
     }
 
     /**
@@ -298,6 +299,7 @@
             parent.mTreeWeight -= child.mTreeWeight;
             parent = parent.getParent();
         }
+        onChildPositionChanged();
     }
 
     /**
@@ -455,9 +457,15 @@
                 mChildren.remove(child);
                 mChildren.add(position, child);
         }
+        onChildPositionChanged();
     }
 
     /**
+     * Notify that a child's position has changed. Possible changes are adding or removing a child.
+     */
+    void onChildPositionChanged() { }
+
+    /**
      * Update override configuration and recalculate full config.
      * @see #mOverrideConfiguration
      * @see #mFullConfiguration
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 942e47b..fc33a4f 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -499,12 +499,6 @@
     final ArrayList<WindowState> mDestroyPreservedSurface = new ArrayList<>();
 
     /**
-     * Windows that have lost input focus and are waiting for the new
-     * focus window to be displayed before they are told about this.
-     */
-    ArrayList<WindowState> mLosingFocus = new ArrayList<>();
-
-    /**
      * This is set when we have run out of memory, and will either be an empty
      * list or contain windows that need to be force removed.
      */
@@ -639,14 +633,6 @@
      */
     final Handler mAnimationHandler = new Handler(AnimationThread.getHandler().getLooper());
 
-    WindowState mCurrentFocus = null;
-    WindowState mLastFocus = null;
-
-    /** Windows added since {@link #mCurrentFocus} was set to null. Used for ANR blaming. */
-    private final ArrayList<WindowState> mWinAddedSinceNullFocus = new ArrayList<>();
-    /** Windows removed since {@link #mCurrentFocus} was set to null. Used for ANR blaming. */
-    private final ArrayList<WindowState> mWinRemovedSinceNullFocus = new ArrayList<>();
-
     /** This just indicates the window the input method is on top of, not
      * necessarily the window its input is going to. */
     WindowState mInputMethodTarget = null;
@@ -721,9 +707,6 @@
         }
     }
 
-    // TODO: Move to RootWindowContainer
-    AppWindowToken mFocusedApp = null;
-
     PowerManager mPowerManager;
     PowerManagerInternal mPowerManagerInternal;
 
@@ -1371,8 +1354,8 @@
                 // the screen after the activity goes away.
                 if (addToastWindowRequiresToken
                         || (attrs.flags & LayoutParams.FLAG_NOT_FOCUSABLE) == 0
-                        || mCurrentFocus == null
-                        || mCurrentFocus.mOwnerUid != callingUid) {
+                        || displayContent.mCurrentFocus == null
+                        || displayContent.mCurrentFocus.mOwnerUid != callingUid) {
                     mH.sendMessageDelayed(
                             mH.obtainMessage(H.WINDOW_HIDE_TIMEOUT, win),
                             win.mAttrs.hideTimeoutMilliseconds);
@@ -1382,8 +1365,8 @@
             // From now on, no exceptions or errors allowed!
 
             res = WindowManagerGlobal.ADD_OKAY;
-            if (mCurrentFocus == null) {
-                mWinAddedSinceNullFocus.add(win);
+            if (displayContent.mCurrentFocus == null) {
+                displayContent.mWinAddedSinceNullFocus.add(win);
             }
 
             if (excludeWindowTypeFromTapOutTask(type)) {
@@ -1504,7 +1487,7 @@
             win.getParent().assignChildLayers();
 
             if (focusChanged) {
-                displayContent.getInputMonitor().setInputFocusLw(mCurrentFocus,
+                displayContent.getInputMonitor().setInputFocusLw(displayContent.mCurrentFocus,
                         false /*updateInputWindows*/);
             }
             displayContent.getInputMonitor().updateInputWindowsLw(false /*force*/);
@@ -1679,8 +1662,9 @@
 
         win.resetAppOpsState();
 
-        if (mCurrentFocus == null) {
-            mWinRemovedSinceNullFocus.add(win);
+        final DisplayContent dc = win.getDisplayContent();
+        if (dc.mCurrentFocus == null) {
+            dc.mWinRemovedSinceNullFocus.add(win);
         }
         mPendingRemove.remove(win);
         mResizingWindows.remove(win);
@@ -1716,7 +1700,6 @@
             atoken.postWindowRemoveStartingWindowCleanup(win);
         }
 
-        final DisplayContent dc = win.getDisplayContent();
         if (win.mAttrs.type == TYPE_WALLPAPER) {
             dc.mWallpaperController.clearLastWallpaperTimeoutTime();
             dc.pendingLayoutChanges |= FINISH_LAYOUT_REDO_WALLPAPER;
@@ -1978,9 +1961,9 @@
             boolean imMayMove = (flagChanges & (FLAG_ALT_FOCUSABLE_IM | FLAG_NOT_FOCUSABLE)) != 0
                     || becameVisible;
             final boolean isDefaultDisplay = win.isDefaultDisplay();
-            boolean focusMayChange = isDefaultDisplay && (win.mViewVisibility != viewVisibility
+            boolean focusMayChange = win.mViewVisibility != viewVisibility
                     || ((flagChanges & FLAG_NOT_FOCUSABLE) != 0)
-                    || (!win.mRelayoutCalled));
+                    || (!win.mRelayoutCalled);
 
             boolean wallpaperMayMove = win.mViewVisibility != viewVisibility
                     && (win.mAttrs.flags & FLAG_SHOW_WALLPAPER) != 0;
@@ -2025,8 +2008,7 @@
                 }
                 result |= RELAYOUT_RES_SURFACE_CHANGED;
                 if (!win.mWillReplaceWindow) {
-                    focusMayChange = tryStartExitingAnimation(win, winAnimator, isDefaultDisplay,
-                            focusMayChange);
+                    focusMayChange = tryStartExitingAnimation(win, winAnimator, focusMayChange);
                 }
             }
 
@@ -2051,7 +2033,7 @@
                     return 0;
                 }
                 if ((result & WindowManagerGlobal.RELAYOUT_RES_FIRST_TIME) != 0) {
-                    focusMayChange = isDefaultDisplay;
+                    focusMayChange = true;
                 }
                 final DisplayContent displayContent = win.getDisplayContent();
                 if (win.mAttrs.type == TYPE_INPUT_METHOD
@@ -2199,7 +2181,7 @@
     }
 
     private boolean tryStartExitingAnimation(WindowState win, WindowStateAnimator winAnimator,
-            boolean isDefaultDisplay, boolean focusMayChange) {
+            boolean focusMayChange) {
         // Try starting an animation; if there isn't one, we
         // can destroy the surface right away.
         int transit = WindowManagerPolicy.TRANSIT_EXIT;
@@ -2207,7 +2189,7 @@
             transit = WindowManagerPolicy.TRANSIT_PREVIEW_DONE;
         }
         if (win.isWinVisibleLw() && winAnimator.applyAnimationLocked(transit, false)) {
-            focusMayChange = isDefaultDisplay;
+            focusMayChange = true;
             win.mAnimatingExit = true;
         } else if (win.isAnimating()) {
             // Currently in a hide animation... turn this into
@@ -2504,57 +2486,6 @@
         }
     }
 
-    void setFocusTaskRegionLocked(AppWindowToken previousFocus) {
-        final Task focusedTask = mFocusedApp != null ? mFocusedApp.getTask() : null;
-        final Task previousTask = previousFocus != null ? previousFocus.getTask() : null;
-        final DisplayContent focusedDisplayContent =
-                focusedTask != null ? focusedTask.getDisplayContent() : null;
-        final DisplayContent previousDisplayContent =
-                previousTask != null ? previousTask.getDisplayContent() : null;
-        if (previousDisplayContent != null && previousDisplayContent != focusedDisplayContent) {
-            previousDisplayContent.setTouchExcludeRegion(null);
-        }
-        if (focusedDisplayContent != null) {
-            focusedDisplayContent.setTouchExcludeRegion(focusedTask);
-        }
-    }
-
-    @Override
-    public void setFocusedApp(IBinder token, boolean moveFocusNow) {
-        if (!checkCallingPermission(MANAGE_APP_TOKENS, "setFocusedApp()")) {
-            throw new SecurityException("Requires MANAGE_APP_TOKENS permission");
-        }
-
-        synchronized(mWindowMap) {
-            final AppWindowToken newFocus;
-            if (token == null) {
-                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Clearing focused app, was " + mFocusedApp);
-                newFocus = null;
-            } else {
-                newFocus = mRoot.getAppWindowToken(token);
-                if (newFocus == null) {
-                    Slog.w(TAG_WM, "Attempted to set focus to non-existing app token: " + token);
-                }
-                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Set focused app to: " + newFocus
-                        + " old focus=" + mFocusedApp + " moveFocusNow=" + moveFocusNow);
-            }
-
-            final boolean changed = mFocusedApp != newFocus;
-            if (changed) {
-                AppWindowToken prev = mFocusedApp;
-                mFocusedApp = newFocus;
-                mFocusedApp.getDisplayContent().getInputMonitor().setFocusedAppLw(newFocus);
-                setFocusTaskRegionLocked(prev);
-            }
-
-            if (moveFocusNow && changed) {
-                final long origId = Binder.clearCallingIdentity();
-                updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, true /*updateInputWindows*/);
-                Binder.restoreCallingIdentity(origId);
-            }
-        }
-    }
-
     @Override
     public void prepareAppTransition(@TransitionType int transit, boolean alwaysKeepCurrent) {
         prepareAppTransition(transit, alwaysKeepCurrent, 0 /* flags */, false /* forceOverride */);
@@ -4395,7 +4326,8 @@
     }
 
     private WindowState getFocusedWindowLocked() {
-        return mCurrentFocus;
+        // Return the focused window in the focused display.
+        return mRoot.getTopFocusedDisplayContent().mCurrentFocus;
     }
 
     TaskStack getImeFocusStackLocked() {
@@ -4403,8 +4335,11 @@
         // Also don't use mInputMethodTarget's stack, because some window with FLAG_NOT_FOCUSABLE
         // and FLAG_ALT_FOCUSABLE_IM flags both set might be set to IME target so they're moved
         // to make room for IME, but the window is not the focused window that's taking input.
-        return (mFocusedApp != null && mFocusedApp.getTask() != null) ?
-                mFocusedApp.getTask().mStack : null;
+        // TODO (b/111080190): Consider the case of multiple IMEs on multi-display.
+        final DisplayContent topFocusedDisplay = mRoot.getTopFocusedDisplayContent();
+        final AppWindowToken focusedApp = topFocusedDisplay.mFocusedApp;
+        return (focusedApp != null && focusedApp.getTask() != null)
+                ? focusedApp.getTask().mStack : null;
     }
 
     public boolean detectSafeMode() {
@@ -4601,6 +4536,7 @@
             }
             switch (msg.what) {
                 case REPORT_FOCUS_CHANGE: {
+                    final DisplayContent displayContent = (DisplayContent) msg.obj;
                     WindowState lastFocus;
                     WindowState newFocus;
 
@@ -4608,24 +4544,22 @@
 
                     synchronized(mWindowMap) {
                         // TODO(multidisplay): Accessibility supported only of default desiplay.
-                        if (mAccessibilityController != null && getDefaultDisplayContentLocked()
-                                .getDisplayId() == DEFAULT_DISPLAY) {
+                        if (mAccessibilityController != null && displayContent.isDefaultDisplay) {
                             accessibilityController = mAccessibilityController;
                         }
 
-                        lastFocus = mLastFocus;
-                        newFocus = mCurrentFocus;
+                        lastFocus = displayContent.mLastFocus;
+                        newFocus = displayContent.mCurrentFocus;
                         if (lastFocus == newFocus) {
                             // Focus is not changing, so nothing to do.
                             return;
                         }
-                        mLastFocus = newFocus;
+                        displayContent.mLastFocus = newFocus;
                         if (DEBUG_FOCUS_LIGHT) Slog.i(TAG_WM, "Focus moving from " + lastFocus +
-                                " to " + newFocus);
-                        if (newFocus != null && lastFocus != null
-                                && !newFocus.isDisplayedLw()) {
-                            //Slog.i(TAG_WM, "Delaying loss of focus...");
-                            mLosingFocus.add(lastFocus);
+                                " to " + newFocus + " displayId=" + displayContent.getDisplayId());
+                        if (newFocus != null && lastFocus != null && !newFocus.isDisplayedLw()) {
+                            if (DEBUG_FOCUS_LIGHT) Slog.i(TAG_WM, "Delaying loss of focus...");
+                            displayContent.mLosingFocus.add(lastFocus);
                             lastFocus = null;
                         }
                     }
@@ -4636,8 +4570,6 @@
                         accessibilityController.onWindowFocusChangedNotLocked();
                     }
 
-                    //System.out.println("Changing focus from " + lastFocus
-                    //                   + " to " + newFocus);
                     if (newFocus != null) {
                         if (DEBUG_FOCUS_LIGHT) Slog.i(TAG_WM, "Gaining focus: " + newFocus);
                         newFocus.reportFocusChangedSerialized(true, mInTouchMode);
@@ -4651,15 +4583,16 @@
                 } break;
 
                 case REPORT_LOSING_FOCUS: {
+                    final DisplayContent displayContent = (DisplayContent) msg.obj;
                     ArrayList<WindowState> losers;
 
                     synchronized(mWindowMap) {
-                        losers = mLosingFocus;
-                        mLosingFocus = new ArrayList<WindowState>();
+                        losers = displayContent.mLosingFocus;
+                        displayContent.mLosingFocus = new ArrayList<>();
                     }
 
                     final int N = losers.size();
-                    for (int i=0; i<N; i++) {
+                    for (int i = 0; i < N; i++) {
                         if (DEBUG_FOCUS_LIGHT) Slog.i(TAG_WM, "Losing delayed focus: " +
                                 losers.get(i));
                         losers.get(i).reportFocusChangedSerialized(false, mInTouchMode);
@@ -5553,89 +5486,11 @@
         }
     }
 
-    // TODO: Move to DisplayContent
     boolean updateFocusedWindowLocked(int mode, boolean updateInputWindows) {
-        WindowState newFocus = mRoot.computeFocusedWindow();
-        if (mCurrentFocus != newFocus) {
-            Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "wmUpdateFocus");
-            // This check makes sure that we don't already have the focus
-            // change message pending.
-            mH.removeMessages(H.REPORT_FOCUS_CHANGE);
-            mH.sendEmptyMessage(H.REPORT_FOCUS_CHANGE);
-            final DisplayContent displayContent = (newFocus != null) ? newFocus.getDisplayContent()
-                    : getDefaultDisplayContentLocked();
-            boolean imWindowChanged = false;
-            if (displayContent.mInputMethodWindow != null) {
-                final WindowState prevTarget = mInputMethodTarget;
-
-                final WindowState newTarget =
-                        displayContent.computeImeTarget(true /* updateImeTarget*/);
-                imWindowChanged = prevTarget != newTarget;
-
-                if (mode != UPDATE_FOCUS_WILL_ASSIGN_LAYERS
-                        && mode != UPDATE_FOCUS_WILL_PLACE_SURFACES) {
-                    displayContent.assignWindowLayers(false /* setLayoutNeeded */);
-                }
-            }
-
-            if (imWindowChanged) {
-                mWindowsChanged = true;
-                displayContent.setLayoutNeeded();
-                newFocus = mRoot.computeFocusedWindow();
-            }
-
-            if (DEBUG_FOCUS_LIGHT || localLOGV) Slog.v(TAG_WM, "Changing focus from " +
-                    mCurrentFocus + " to " + newFocus + " Callers=" + Debug.getCallers(4));
-            final WindowState oldFocus = mCurrentFocus;
-            mCurrentFocus = newFocus;
-            mLosingFocus.remove(newFocus);
-
-            if (mCurrentFocus != null) {
-                mWinAddedSinceNullFocus.clear();
-                mWinRemovedSinceNullFocus.clear();
-            }
-
-            int focusChanged = mPolicy.focusChangedLw(oldFocus, newFocus);
-
-            if (imWindowChanged && oldFocus != displayContent.mInputMethodWindow) {
-                // Focus of the input method window changed. Perform layout if needed.
-                if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
-                    displayContent.performLayout(true /*initial*/,  updateInputWindows);
-                    focusChanged &= ~FINISH_LAYOUT_REDO_LAYOUT;
-                } else if (mode == UPDATE_FOCUS_WILL_PLACE_SURFACES) {
-                    // Client will do the layout, but we need to assign layers
-                    // for handleNewWindowLocked() below.
-                    displayContent.assignWindowLayers(false /* setLayoutNeeded */);
-                }
-            }
-
-            if ((focusChanged & FINISH_LAYOUT_REDO_LAYOUT) != 0) {
-                // The change in focus caused us to need to do a layout.  Okay.
-                displayContent.setLayoutNeeded();
-                if (mode == UPDATE_FOCUS_PLACING_SURFACES) {
-                    displayContent.performLayout(true /*initial*/, updateInputWindows);
-                } else if (mode == UPDATE_FOCUS_REMOVING_FOCUS) {
-                    mRoot.performSurfacePlacement(false);
-                }
-            }
-
-            if (mode != UPDATE_FOCUS_WILL_ASSIGN_LAYERS) {
-                // If we defer assigning layers, then the caller is responsible for
-                // doing this part.
-                displayContent.getInputMonitor().setInputFocusLw(mCurrentFocus, updateInputWindows);
-            }
-
-            displayContent.adjustForImeIfNeeded();
-
-            // We may need to schedule some toast windows to be removed. The toasts for an app that
-            // does not have input focus are removed within a timeout to prevent apps to redress
-            // other apps' UI.
-            displayContent.scheduleToastWindowsTimeoutIfNeededLocked(oldFocus, newFocus);
-
-            Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
-            return true;
-        }
-        return false;
+        Trace.traceBegin(TRACE_TAG_WINDOW_MANAGER, "wmUpdateFocus");
+        boolean changed = mRoot.updateFocusedWindowLocked(mode, updateInputWindows);
+        Trace.traceEnd(TRACE_TAG_WINDOW_MANAGER);
+        return changed;
     }
 
     void startFreezingDisplayLocked(int exitAnim, int enterAnim) {
@@ -6197,11 +6052,12 @@
     void writeToProtoLocked(ProtoOutputStream proto, boolean trim) {
         mPolicy.writeToProto(proto, POLICY);
         mRoot.writeToProto(proto, ROOT_WINDOW_CONTAINER, trim);
-        if (mCurrentFocus != null) {
-            mCurrentFocus.writeIdentifierToProto(proto, FOCUSED_WINDOW);
+        final DisplayContent topFocusedDisplayContent = mRoot.getTopFocusedDisplayContent();
+        if (topFocusedDisplayContent.mCurrentFocus != null) {
+            topFocusedDisplayContent.mCurrentFocus.writeIdentifierToProto(proto, FOCUSED_WINDOW);
         }
-        if (mFocusedApp != null) {
-            mFocusedApp.writeNameToProto(proto, FOCUSED_APP);
+        if (topFocusedDisplayContent.mFocusedApp != null) {
+            topFocusedDisplayContent.mFocusedApp.writeNameToProto(proto, FOCUSED_APP);
         }
         final WindowState imeWindow = mRoot.getCurrentInputMethodWindow();
         if (imeWindow != null) {
@@ -6299,23 +6155,6 @@
                 }
             }
         }
-        if (mLosingFocus.size() > 0) {
-            pw.println();
-            pw.println("  Windows losing focus:");
-            for (int i=mLosingFocus.size()-1; i>=0; i--) {
-                WindowState w = mLosingFocus.get(i);
-                if (windows == null || windows.contains(w)) {
-                    pw.print("  Losing #"); pw.print(i); pw.print(' ');
-                            pw.print(w);
-                    if (dumpAll) {
-                        pw.println(":");
-                        w.dump(pw, "    ", true);
-                    } else {
-                        pw.println();
-                    }
-                }
-            }
-        }
         if (mResizingWindows.size() > 0) {
             pw.println();
             pw.println("  Windows waiting to resize:");
@@ -6344,11 +6183,7 @@
         pw.println();
         pw.print("  mGlobalConfiguration="); pw.println(mRoot.getConfiguration());
         pw.print("  mHasPermanentDpad="); pw.println(mHasPermanentDpad);
-        pw.print("  mCurrentFocus="); pw.println(mCurrentFocus);
-        if (mLastFocus != mCurrentFocus) {
-            pw.print("  mLastFocus="); pw.println(mLastFocus);
-        }
-        pw.print("  mFocusedApp="); pw.println(mFocusedApp);
+        mRoot.dumpTopFocusedDisplayId(pw);
         if (mInputMethodTarget != null) {
             pw.print("  mInputMethodTarget="); pw.println(mInputMethodTarget);
         }
@@ -6479,11 +6314,17 @@
         if (reason != null) {
             pw.println("  Reason: " + reason);
         }
-        if (!mWinAddedSinceNullFocus.isEmpty()) {
-            pw.println("  Windows added since null focus: " + mWinAddedSinceNullFocus);
-        }
-        if (!mWinRemovedSinceNullFocus.isEmpty()) {
-            pw.println("  Windows removed since null focus: " + mWinRemovedSinceNullFocus);
+        for (int i = mRoot.getChildCount() - 1; i >= 0; i--) {
+            final DisplayContent dc = mRoot.getChildAt(i);
+            final int displayId = dc.getDisplayId();
+            if (!dc.mWinAddedSinceNullFocus.isEmpty()) {
+                pw.println("  Windows added in display #" + displayId + " since null focus: "
+                        + dc.mWinAddedSinceNullFocus);
+            }
+            if (!dc.mWinRemovedSinceNullFocus.isEmpty()) {
+                pw.println("  Windows removed in display #" + displayId + " since null focus: "
+                        + dc.mWinRemovedSinceNullFocus);
+            }
         }
         pw.println();
         dumpWindowsNoHeaderLocked(pw, true, null);
@@ -6818,9 +6659,9 @@
     @Override
     public void setDockedStackDividerTouchRegion(Rect touchRegion) {
         synchronized (mWindowMap) {
-            getDefaultDisplayContentLocked().getDockedDividerController()
-                    .setTouchRegion(touchRegion);
-            setFocusTaskRegionLocked(null);
+            final DisplayContent dc = getDefaultDisplayContentLocked();
+            dc.getDockedDividerController().setTouchRegion(touchRegion);
+            dc.updateTouchExcludeRegion();
         }
     }
 
@@ -7373,7 +7214,14 @@
         @Override
         public boolean isUidFocused(int uid) {
             synchronized (mWindowMap) {
-                return mCurrentFocus != null ? uid == mCurrentFocus.getOwningUid() : false;
+                for (int i = mRoot.getChildCount() - 1; i >= 0; i--) {
+                    final DisplayContent displayContent = mRoot.getChildAt(i);
+                    if (displayContent.mCurrentFocus != null
+                            && uid == displayContent.mCurrentFocus.getOwningUid()) {
+                        return true;
+                    }
+                }
+                return false;
             }
         }
 
@@ -7396,8 +7244,9 @@
                 // press home.  Sometimes the IME won't go down.)
                 // Would be nice to fix this more correctly, but it's
                 // way at the end of a release, and this should be good enough.
-                if (mCurrentFocus != null && mCurrentFocus.mSession.mUid == uid
-                        && mCurrentFocus.mSession.mPid == pid) {
+                final WindowState currentFocus = mRoot.getTopFocusedDisplayContent().mCurrentFocus;
+                if (currentFocus != null && currentFocus.mSession.mUid == uid
+                        && currentFocus.mSession.mPid == pid) {
                     return true;
                 }
             }
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index a4bac31..8276952 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -1833,7 +1833,7 @@
         if (startingWindow && DEBUG_STARTING_WINDOW) Slog.d(TAG_WM,
                 "Starting window removed " + this);
 
-        if (localLOGV || DEBUG_FOCUS || DEBUG_FOCUS_LIGHT && this == mService.mCurrentFocus)
+        if (localLOGV || DEBUG_FOCUS || DEBUG_FOCUS_LIGHT && isFocused())
             Slog.v(TAG_WM, "Remove " + this + " client="
                         + Integer.toHexString(System.identityHashCode(mClient.asBinder()))
                         + ", surfaceController=" + mWinAnimator.mSurfaceController + " Callers="
@@ -1945,7 +1945,7 @@
             if (wasVisible && mService.updateOrientationFromAppTokensLocked(displayId)) {
                 mService.mH.obtainMessage(SEND_NEW_CONFIGURATION, displayId).sendToTarget();
             }
-            mService.updateFocusedWindowLocked(mService.mCurrentFocus == this
+            mService.updateFocusedWindowLocked(isFocused()
                             ? UPDATE_FOCUS_REMOVING_FOCUS
                             : UPDATE_FOCUS_NORMAL,
                     true /*updateInputWindows*/);
@@ -2180,7 +2180,7 @@
             mPolicyVisibility = mPolicyVisibilityAfterAnim;
             if (!mPolicyVisibility) {
                 mWinAnimator.hide("checkPolicyVisibilityChange");
-                if (mService.mCurrentFocus == this) {
+                if (isFocused()) {
                     if (DEBUG_FOCUS_LIGHT) Slog.i(TAG,
                             "setAnimationLocked: setting mFocusMayChange true");
                     mService.mFocusMayChange = true;
@@ -2482,6 +2482,7 @@
             }
         }
         mPolicyVisibilityAfterAnim = false;
+        final boolean isFocused = isFocused();
         if (!doAnimation) {
             if (DEBUG_VISIBILITY) Slog.v(TAG, "Policy visibility false: " + this);
             mPolicyVisibility = false;
@@ -2489,7 +2490,7 @@
             // for it to be displayed before enabling the display, that
             // we allow the display to be enabled now.
             mService.enableScreenIfNeededLocked();
-            if (mService.mCurrentFocus == this) {
+            if (isFocused) {
                 if (DEBUG_FOCUS_LIGHT) Slog.i(TAG,
                         "WindowState.hideLw: setting mFocusMayChange true");
                 mService.mFocusMayChange = true;
@@ -2498,7 +2499,7 @@
         if (requestAnim) {
             mService.scheduleAnimationLocked();
         }
-        if (mService.mCurrentFocus == this) {
+        if (isFocused) {
             mService.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, false /* updateImWindows */);
         }
         return true;
@@ -2994,10 +2995,8 @@
         }
     }
 
-    public boolean isFocused() {
-        synchronized(mService.mWindowMap) {
-            return mService.mCurrentFocus == this;
-        }
+    boolean isFocused() {
+        return getDisplayContent().mCurrentFocus == this;
     }
 
     @Override
@@ -4435,7 +4434,12 @@
         @Override
         public boolean isFocused() {
             final WindowState outer = mOuter.get();
-            return outer != null && outer.isFocused();
+            if (outer != null) {
+                synchronized (outer.mService.mWindowMap) {
+                    return outer.isFocused();
+                }
+            }
+            return false;
         }
     }
 
@@ -4663,10 +4667,7 @@
 
         mTapExcludeRegionHolder.updateRegion(regionId, left, top, width, height);
         // Trigger touch exclude region update on current display.
-        final boolean isAppFocusedOnDisplay = mService.mFocusedApp != null
-                && mService.mFocusedApp.getDisplayContent() == currentDisplay;
-        currentDisplay.setTouchExcludeRegion(isAppFocusedOnDisplay ? mService.mFocusedApp.getTask()
-                : null);
+        currentDisplay.updateTouchExcludeRegion();
     }
 
     /** Union the region with current tap exclude region that this window provides. */
diff --git a/services/core/java/com/android/server/wm/WindowSurfacePlacer.java b/services/core/java/com/android/server/wm/WindowSurfacePlacer.java
index 080a3a2..24ac07a 100644
--- a/services/core/java/com/android/server/wm/WindowSurfacePlacer.java
+++ b/services/core/java/com/android/server/wm/WindowSurfacePlacer.java
@@ -763,6 +763,7 @@
 
     private void processApplicationsAnimatingInPlace(int transit) {
         if (transit == TRANSIT_TASK_IN_PLACE) {
+            // TODO (b/111362605): non-default-display transition.
             // Find the focused window
             final WindowState win = mService.getDefaultDisplayContentLocked().findFocusedWindow();
             if (win != null) {
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index fefd305..0cf79b6 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -255,6 +255,7 @@
         super.removeImmediately();
     }
 
+    @Override
     void onDisplayChanged(DisplayContent dc) {
         dc.reParentWindowToken(this);
         mDisplayContent = dc;
diff --git a/services/core/jni/com_android_server_input_InputManagerService.cpp b/services/core/jni/com_android_server_input_InputManagerService.cpp
index a754d2a..1f2eef4 100644
--- a/services/core/jni/com_android_server_input_InputManagerService.cpp
+++ b/services/core/jni/com_android_server_input_InputManagerService.cpp
@@ -221,7 +221,8 @@
     status_t unregisterInputChannel(JNIEnv* env, const sp<InputChannel>& inputChannel);
 
     void setInputWindows(JNIEnv* env, jobjectArray windowHandleObjArray, int32_t displayId);
-    void setFocusedApplication(JNIEnv* env, jobject applicationHandleObj);
+    void setFocusedApplication(JNIEnv* env, int32_t displayId, jobject applicationHandleObj);
+    void setFocusedDisplay(JNIEnv* env, int32_t displayId);
     void setInputDispatchMode(bool enabled, bool frozen);
     void setSystemUiVisibility(int32_t visibility);
     void setPointerSpeed(int32_t speed);
@@ -771,10 +772,15 @@
     }
 }
 
-void NativeInputManager::setFocusedApplication(JNIEnv* env, jobject applicationHandleObj) {
+void NativeInputManager::setFocusedApplication(JNIEnv* env, int32_t displayId,
+        jobject applicationHandleObj) {
     sp<InputApplicationHandle> applicationHandle =
             android_server_InputApplicationHandle_getHandle(env, applicationHandleObj);
-    mInputManager->getDispatcher()->setFocusedApplication(applicationHandle);
+    mInputManager->getDispatcher()->setFocusedApplication(displayId, applicationHandle);
+}
+
+void NativeInputManager::setFocusedDisplay(JNIEnv* env, int32_t displayId) {
+    mInputManager->getDispatcher()->setFocusedDisplay(displayId);
 }
 
 void NativeInputManager::setInputDispatchMode(bool enabled, bool frozen) {
@@ -1413,10 +1419,17 @@
 }
 
 static void nativeSetFocusedApplication(JNIEnv* env, jclass /* clazz */,
-        jlong ptr, jobject applicationHandleObj) {
+        jlong ptr, jint displayId, jobject applicationHandleObj) {
     NativeInputManager* im = reinterpret_cast<NativeInputManager*>(ptr);
 
-    im->setFocusedApplication(env, applicationHandleObj);
+    im->setFocusedApplication(env, displayId, applicationHandleObj);
+}
+
+static void nativeSetFocusedDisplay(JNIEnv* env, jclass /* clazz */,
+        jlong ptr, jint displayId) {
+    NativeInputManager* im = reinterpret_cast<NativeInputManager*>(ptr);
+
+    im->setFocusedDisplay(env, displayId);
 }
 
 static void nativeSetPointerCapture(JNIEnv* env, jclass /* clazz */, jlong ptr,
@@ -1638,8 +1651,10 @@
             (void*) nativeToggleCapsLock },
     { "nativeSetInputWindows", "(J[Lcom/android/server/input/InputWindowHandle;I)V",
             (void*) nativeSetInputWindows },
-    { "nativeSetFocusedApplication", "(JLcom/android/server/input/InputApplicationHandle;)V",
+    { "nativeSetFocusedApplication", "(JILcom/android/server/input/InputApplicationHandle;)V",
             (void*) nativeSetFocusedApplication },
+    { "nativeSetFocusedDisplay", "(JI)V",
+            (void*) nativeSetFocusedDisplay },
     { "nativeSetPointerCapture", "(JZ)V",
             (void*) nativeSetPointerCapture },
     { "nativeSetInputDispatchMode", "(JZZ)V",
diff --git a/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java b/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
index 0345a81..3d0bdc9 100644
--- a/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
+++ b/services/tests/servicestests/src/com/android/server/am/ActivityStackSupervisorTests.java
@@ -232,6 +232,7 @@
         doReturn(displaySleeping).when(display).isSleeping();
         doReturn(keyguardShowing).when(keyguard).isKeyguardOrAodShowing(anyInt());
 
+        doReturn(isFocusedStack).when(stack).isFocusedStackOnDisplay();
         doReturn(isFocusedStack ? stack : null).when(display).getFocusedStack();
         mSupervisor.applySleepTokensLocked(true);
         verify(stack, times(expectWakeFromSleep ? 1 : 0)).awakeFromSleepingLocked();
diff --git a/services/tests/servicestests/src/com/android/server/am/ActivityStackTests.java b/services/tests/servicestests/src/com/android/server/am/ActivityStackTests.java
index ab814ee..5fcd2aa 100644
--- a/services/tests/servicestests/src/com/android/server/am/ActivityStackTests.java
+++ b/services/tests/servicestests/src/com/android/server/am/ActivityStackTests.java
@@ -38,6 +38,7 @@
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.spy;
 
 import android.content.pm.ActivityInfo;
 import android.os.UserHandle;
@@ -71,8 +72,8 @@
 
         setupActivityTaskManagerService();
         mDefaultDisplay = mSupervisor.getDefaultDisplay();
-        mStack = mDefaultDisplay.createStack(WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_STANDARD,
-                true /* onTop */);
+        mStack = spy(mDefaultDisplay.createStack(WINDOWING_MODE_UNDEFINED, ACTIVITY_TYPE_STANDARD,
+                true /* onTop */));
         mTask = new TaskBuilder(mSupervisor).setStack(mStack).build();
     }
 
@@ -720,7 +721,7 @@
         doReturn(display).when(mSupervisor).getActivityDisplay(anyInt());
         doReturn(keyguardGoingAway).when(keyguardController).isKeyguardGoingAway();
         doReturn(displaySleeping).when(display).isSleeping();
-        doReturn(focusedStack ? mStack : null).when(mSupervisor).getTopDisplayFocusedStack();
+        doReturn(focusedStack).when(mStack).isFocusedStackOnDisplay();
 
         assertEquals(expected, mStack.shouldSleepActivities());
     }
diff --git a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
index b330304..96c5fde 100644
--- a/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/servicestests/src/com/android/server/wm/DisplayContentTests.java
@@ -32,6 +32,7 @@
 import static android.view.WindowManager.LayoutParams.TYPE_VOICE_INTERACTION;
 
 import static com.android.server.wm.WindowContainer.POSITION_TOP;
+import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
 
 import static org.hamcrest.Matchers.is;
 import static org.junit.Assert.assertEquals;
@@ -303,7 +304,8 @@
                 createTapEvent(dm0.widthPixels / 2, dm0.heightPixels / 2, false));
 
         // Check focus is on primary display.
-        assertEquals(sWm.mCurrentFocus, dc0.findFocusedWindow());
+        assertEquals(sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus,
+                dc0.findFocusedWindow());
 
         // Tap on secondary display
         DisplayMetrics dm1 = dc1.getDisplayMetrics();
@@ -313,7 +315,8 @@
                 createTapEvent(dm1.widthPixels / 2, dm1.heightPixels / 2, false));
 
         // Check focus is on secondary.
-        assertEquals(sWm.mCurrentFocus, dc1.findFocusedWindow());
+        assertEquals(sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus,
+                dc1.findFocusedWindow());
     }
 
     @Test
@@ -321,34 +324,29 @@
         // Create a focusable window and check that focus is calculated correctly
         final WindowState window1 =
                 createWindow(null, TYPE_BASE_APPLICATION, mDisplayContent, "window1");
-        assertEquals(window1, sWm.mRoot.computeFocusedWindow());
+        updateFocusedWindow();
+        assertTrue(window1.isFocused());
+        assertEquals(window1, sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus);
 
         // Check that a new display doesn't affect focus
         final DisplayContent dc = createNewDisplay();
-        assertEquals(window1, sWm.mRoot.computeFocusedWindow());
+        updateFocusedWindow();
+        assertTrue(window1.isFocused());
+        assertEquals(window1, sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus);
 
         // Add a window to the second display, and it should be focused
         final WindowState window2 = createWindow(null, TYPE_BASE_APPLICATION, dc, "window2");
-        assertEquals(window2, sWm.mRoot.computeFocusedWindow());
+        updateFocusedWindow();
+        assertTrue(window1.isFocused());
+        assertTrue(window2.isFocused());
+        assertEquals(window2, sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus);
 
         // Move the first window to the to including parents, and make sure focus is updated
         window1.getParent().positionChildAt(POSITION_TOP, window1, true);
-        assertEquals(window1, sWm.mRoot.computeFocusedWindow());
-    }
-
-    @Test
-    public void testKeyguard_preventsSecondaryDisplayFocus() throws Exception {
-        final WindowState keyguard = createWindow(null, TYPE_STATUS_BAR,
-                sWm.getDefaultDisplayContentLocked(), "keyguard");
-        assertEquals(keyguard, sWm.mRoot.computeFocusedWindow());
-
-        // Add a window to a second display, and it should be focused
-        final DisplayContent dc = createNewDisplay();
-        final WindowState win = createWindow(null, TYPE_BASE_APPLICATION, dc, "win");
-        assertEquals(win, sWm.mRoot.computeFocusedWindow());
-
-        mWmRule.getWindowManagerPolicy().keyguardShowingAndNotOccluded = true;
-        assertEquals(keyguard, sWm.mRoot.computeFocusedWindow());
+        updateFocusedWindow();
+        assertTrue(window1.isFocused());
+        assertTrue(window2.isFocused());
+        assertEquals(window1, sWm.mRoot.getTopFocusedDisplayContent().mCurrentFocus);
     }
 
     /**
@@ -590,6 +588,12 @@
         assertEquals(displayContent.mBaseDisplayDensity, expectedBaseDensity);
     }
 
+    private void updateFocusedWindow() {
+        synchronized (sWm.mWindowMap) {
+            sWm.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, false);
+        }
+    }
+
     /**
      * Create DisplayContent that does not update display base/initial values from device to keep
      * the values set by test.