Associate WindowToken object with only one display at a time

WindowTokens were global objects that contained windows that could
be on multiple displays. This model does not work with the
WindowContainer hierarachy as children (window tokens) can not have
mulitple parents (displays).
We now:
- Track the mapping of binder tokens to window tokens per display
instead of globally . So, you can have a binder token map to
individual WindowToken objects per display.
- WMS.addWindowToken is used to create a WindowToken that clients
can then later add windows to. However, when addWindowToken is called
we don't know the display the client(s) would like to add window to.
So, we track binder tokens that we are allowed to add window for in
the RootWindowContainer and create a window token for the binder on
a specific display when we try to add a window.

Bug: 30060889
Test: Manual testing and existing tests pass.
Change-Id: I81a52a32b01c33ed32169d2da0506b688ea9bc8a
diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index 508cf24..7e09c16 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -158,8 +158,10 @@
     ArrayDeque<Rect> mFrozenBounds = new ArrayDeque<>();
     ArrayDeque<Configuration> mFrozenMergedConfig = new ArrayDeque<>();
 
-    AppWindowToken(WindowManagerService service, IApplicationToken token, boolean _voiceInteraction) {
-        super(service, token != null ? token.asBinder() : null, TYPE_APPLICATION, true);
+    AppWindowToken(WindowManagerService service, IApplicationToken token, boolean _voiceInteraction,
+            DisplayContent displayContent) {
+        super(service, token != null ? token.asBinder() : null, TYPE_APPLICATION, true,
+                displayContent);
         appToken = token;
         voiceInteraction = _voiceInteraction;
         mInputApplicationHandle = new InputApplicationHandle(this);
@@ -897,7 +899,7 @@
     }
 
     boolean transferStartingWindow(IBinder transferFrom) {
-        final AppWindowToken fromToken = mService.findAppWindowToken(transferFrom);
+        final AppWindowToken fromToken = getDisplayContent().getAppWindowToken(transferFrom);
         if (fromToken == null) {
             return false;
         }
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index fb37c9f..8b720e0 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -53,6 +53,7 @@
 import android.graphics.Region.Op;
 import android.hardware.display.DisplayManagerInternal;
 import android.os.Debug;
+import android.os.IBinder;
 import android.util.DisplayMetrics;
 import android.util.Slog;
 import android.view.Display;
@@ -66,6 +67,8 @@
 import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 
 class DisplayContentList extends ArrayList<DisplayContent> {
@@ -87,6 +90,9 @@
      * from mDisplayWindows; */
     private final WindowList mWindows = new WindowList();
 
+    // Mapping from a token IBinder to a WindowToken object on this display.
+    private final HashMap<IBinder, WindowToken> mTokenMap = new HashMap();
+
     int mInitialDisplayWidth = 0;
     int mInitialDisplayHeight = 0;
     int mInitialDisplayDensity = 0;
@@ -167,6 +173,35 @@
         return mWindows;
     }
 
+    WindowToken getWindowToken(IBinder binder) {
+        return mTokenMap.get(binder);
+    }
+
+    AppWindowToken getAppWindowToken(IBinder binder) {
+        final WindowToken token = getWindowToken(binder);
+        if (token == null) {
+            return null;
+        }
+        return token.asAppWindowToken();
+    }
+
+    void setWindowToken(IBinder binder, WindowToken token) {
+        final DisplayContent dc = mService.mRoot.getWindowTokenDisplay(token);
+        if (dc != null) {
+            // We currently don't support adding a window token to the display if the display
+            // already has the binder mapped to another token. If there is a use case for supporting
+            // this moving forward we will either need to merge the WindowTokens some how or have
+            // the binder map to a list of window tokens.
+            throw new IllegalArgumentException("Can't map token=" + token + " to display=" + this
+                    + " already mapped to display=" + dc + " tokens=" + dc.mTokenMap);
+        }
+        mTokenMap.put(binder, token);
+    }
+
+    WindowToken removeWindowToken(IBinder binder) {
+        return mTokenMap.remove(binder);
+    }
+
     Display getDisplay() {
         return mDisplay;
     }
@@ -342,6 +377,7 @@
             mHomeStack = stack;
         }
         addChild(stack, onTop);
+        stack.onDisplayChanged(this);
     }
 
     void moveStack(TaskStack stack, boolean toTop) {
@@ -899,7 +935,7 @@
         // position; else we need to look some more.
         if (pos != null) {
             // Move behind any windows attached to this one.
-            final WindowToken atoken = mService.mTokenMap.get(pos.mClient.asBinder());
+            final WindowToken atoken = getWindowToken(pos.mClient.asBinder());
             if (atoken != null) {
                 tokenWindowList = getTokenWindowsOnDisplay(atoken);
                 final int NC = tokenWindowList.size();
@@ -929,7 +965,7 @@
 
         if (pos != null) {
             // Move in front of any windows attached to this one.
-            final WindowToken atoken = mService.mTokenMap.get(pos.mClient.asBinder());
+            final WindowToken atoken = getWindowToken(pos.mClient.asBinder());
             if (atoken != null) {
                 final WindowState top = atoken.getTopWindow();
                 if (top != null && top.mSubLayer >= 0) {
@@ -1238,6 +1274,25 @@
         }
     }
 
+    void dumpTokens(PrintWriter pw, boolean dumpAll) {
+        if (mTokenMap.isEmpty()) {
+            return;
+        }
+        pw.println("  Display #" + mDisplayId);
+        final Iterator<WindowToken> it = mTokenMap.values().iterator();
+        while (it.hasNext()) {
+            final WindowToken token = it.next();
+            pw.print("  ");
+            pw.print(token);
+            if (dumpAll) {
+                pw.println(':');
+                token.dump(pw, "    ");
+            } else {
+                pw.println();
+            }
+        }
+    }
+
     static final class GetWindowOnDisplaySearchResult {
         boolean reachedToken;
         WindowState foundWindow;
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index c553b52..3c13439 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -22,6 +22,7 @@
 import android.hardware.power.V1_0.PowerHint;
 import android.os.Binder;
 import android.os.Debug;
+import android.os.IBinder;
 import android.os.ParcelFileDescriptor;
 import android.os.PowerManager;
 import android.os.RemoteException;
@@ -42,9 +43,11 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.LinkedList;
 
 import static android.app.ActivityManager.StackId.DOCKED_STACK_ID;
+import static android.view.Display.DEFAULT_DISPLAY;
 import static android.view.WindowManager.LayoutParams.FLAG_KEEP_SCREEN_ON;
 import static android.view.WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_KEYGUARD;
@@ -58,8 +61,11 @@
 import static android.view.WindowManagerPolicy.FINISH_LAYOUT_REDO_LAYOUT;
 import static android.view.WindowManagerPolicy.FINISH_LAYOUT_REDO_WALLPAPER;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
+import static com.android.server.wm.AppTransition.TRANSIT_UNSET;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_ADD_REMOVE;
+import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_APP_TRANSITIONS;
 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;
 import static com.android.server.wm.WindowManagerDebugConfig.DEBUG_LAYOUT_REPEATS;
@@ -79,6 +85,7 @@
 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.LAYOUT_REPEAT_THRESHOLD;
+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;
@@ -137,6 +144,15 @@
 
     private final LinkedList<AppWindowToken> mTmpUpdateAllDrawn = new LinkedList();
 
+    private final ArrayList<WindowToken> mTmpTokensList = new ArrayList();
+
+    // Collection of binder tokens mapped to their window type we are allowed to create window
+    // tokens for but that are not current attached to any display. We need to track this here
+    // because a binder token can be added through {@link WindowManagerService#addWindowToken},
+    // but we don't know what display windows for the token will be added to until
+    // {@link WindowManagerService#addWindow} is called.
+    private final HashMap<IBinder, Integer> mUnattachedBinderTokens = new HashMap();
+
     // State for the RemoteSurfaceTrace system used in testing. If this is enabled SurfaceControl
     // instances will be replaced with an instance that writes a binary representation of all
     // commands to mSurfaceTraceFd.
@@ -179,7 +195,7 @@
         return dc;
     }
 
-    private DisplayContent getDisplayContent(int displayId) {
+    DisplayContent getDisplayContent(int displayId) {
         for (int i = mChildren.size() - 1; i >= 0; --i) {
             final DisplayContent current = mChildren.get(i);
             if (current.getDisplayId() == displayId) {
@@ -253,7 +269,6 @@
         }
 
         if (!attachedToDisplay) {
-            stack.attachDisplayContent(dc);
             dc.attachStack(stack, onTop);
         }
 
@@ -338,6 +353,202 @@
         return null;
     }
 
+    /** Return the window token associated with the input binder token on the input display */
+    WindowToken getWindowToken(IBinder binder, DisplayContent dc) {
+        final WindowToken token = dc.getWindowToken(binder);
+        if (token != null) {
+            return token;
+        }
+
+        // There is no window token mapped to the binder on the display. Create and map a window
+        // token if it is currently allowed.
+        if (!mUnattachedBinderTokens.containsKey(binder)) {
+            return null;
+        }
+
+        final int type = mUnattachedBinderTokens.get(binder);
+        return new WindowToken(mService, binder, type, true, dc);
+    }
+
+    /** Returns all window tokens mapped to the input binder. */
+    ArrayList<WindowToken> getWindowTokens(IBinder binder) {
+        mTmpTokensList.clear();
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent dc = mChildren.get(i);
+            final WindowToken token = dc.getWindowToken(binder);
+            if (token != null) {
+                mTmpTokensList.add(token);
+            }
+        }
+        return mTmpTokensList;
+    }
+
+    /**
+     * Returns the app window token for the input binder if it exist in the system.
+     * NOTE: Only one AppWindowToken is allowed to exist in the system for a binder token, since
+     * AppWindowToken represents an activity which can only exist on one display.
+     */
+    AppWindowToken getAppWindowToken(IBinder binder) {
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent dc = mChildren.get(i);
+            final AppWindowToken atoken = dc.getAppWindowToken(binder);
+            if (atoken != null) {
+                return atoken;
+            }
+        }
+        return null;
+    }
+
+    /** Returns the display object the input window token is currently mapped on. */
+    DisplayContent getWindowTokenDisplay(WindowToken token) {
+        if (token == null) {
+            return null;
+        }
+
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent dc = mChildren.get(i);
+            final WindowToken current = dc.getWindowToken(token.token);
+            if (current == token) {
+                return dc;
+            }
+        }
+
+        return null;
+    }
+
+    void addWindowToken(IBinder binder, int type) {
+        if (mUnattachedBinderTokens.containsKey(binder)) {
+            Slog.w(TAG_WM, "addWindowToken: Attempted to add existing binder token: " + binder);
+            return;
+        }
+
+        final ArrayList<WindowToken> tokens = getWindowTokens(binder);
+
+        if (!tokens.isEmpty()) {
+            Slog.w(TAG_WM, "addWindowToken: Attempted to add binder token: " + binder
+                    + " for already created window tokens: " + tokens);
+            return;
+        }
+
+        mUnattachedBinderTokens.put(binder, type);
+
+        // TODO(multi-display): By default we add this to the default display, but maybe we
+        // should provide an API for a token to be added to any display?
+        final WindowToken token = new WindowToken(mService, binder, type, true,
+                getDisplayContent(DEFAULT_DISPLAY));
+        if (type == TYPE_WALLPAPER) {
+            mService.mWallpaperControllerLocked.addWallpaperToken(token);
+        }
+    }
+
+    ArrayList<WindowToken> removeWindowToken(IBinder binder) {
+        mUnattachedBinderTokens.remove(binder);
+
+        mTmpTokensList.clear();
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent dc = mChildren.get(i);
+            final WindowToken token = dc.removeWindowToken(binder);
+            if (token != null) {
+                mTmpTokensList.add(token);
+            }
+        }
+        return mTmpTokensList;
+    }
+
+    /**
+     * Removed the mapping to the input binder for the system if it no longer as a window token
+     * associated with it on any display.
+     */
+    void removeWindowTokenIfPossible(IBinder binder) {
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final DisplayContent dc = mChildren.get(i);
+            final WindowToken token = dc.getWindowToken(binder);
+            if (token != null) {
+                return;
+            }
+        }
+
+        mUnattachedBinderTokens.remove(binder);
+    }
+
+    void removeAppToken(IBinder binder) {
+        final ArrayList<WindowToken> removedTokens = removeWindowToken(binder);
+        if (removedTokens == null || removedTokens.isEmpty()) {
+            Slog.w(TAG_WM, "removeAppToken: Attempted to remove non-existing token: " + binder);
+            return;
+        }
+
+        for (int i = removedTokens.size() - 1; i >= 0; --i) {
+            WindowToken wtoken = removedTokens.get(i);
+            AppWindowToken appToken = wtoken.asAppWindowToken();
+
+            if (appToken == null) {
+                Slog.w(TAG_WM,
+                        "Attempted to remove non-App token: " + binder + " wtoken=" + wtoken);
+                continue;
+            }
+
+            AppWindowToken startingToken = null;
+
+            if (DEBUG_APP_TRANSITIONS) Slog.v(TAG_WM, "Removing app token: " + appToken);
+
+            boolean delayed = appToken.setVisibility(null, false, TRANSIT_UNSET, true,
+                    appToken.voiceInteraction);
+
+            mService.mOpeningApps.remove(appToken);
+            appToken.waitingToShow = false;
+            if (mService.mClosingApps.contains(appToken)) {
+                delayed = true;
+            } else if (mService.mAppTransition.isTransitionSet()) {
+                mService.mClosingApps.add(appToken);
+                delayed = true;
+            }
+
+            if (DEBUG_APP_TRANSITIONS) Slog.v(TAG_WM, "Removing app " + appToken
+                    + " delayed=" + delayed
+                    + " animation=" + appToken.mAppAnimator.animation
+                    + " animating=" + appToken.mAppAnimator.animating);
+
+            if (DEBUG_ADD_REMOVE || DEBUG_TOKEN_MOVEMENT) Slog.v(TAG_WM, "removeAppToken: "
+                    + appToken + " delayed=" + delayed + " Callers=" + Debug.getCallers(4));
+
+            final TaskStack stack = appToken.mTask.mStack;
+            if (delayed && !appToken.isEmpty()) {
+                // set the token aside because it has an active animation to be finished
+                if (DEBUG_ADD_REMOVE || DEBUG_TOKEN_MOVEMENT) Slog.v(TAG_WM,
+                        "removeAppToken make exiting: " + appToken);
+                stack.mExitingAppTokens.add(appToken);
+                appToken.mIsExiting = true;
+            } else {
+                // Make sure there is no animation running on this token, so any windows associated
+                // with it will be removed as soon as their animations are complete
+                appToken.mAppAnimator.clearAnimation();
+                appToken.mAppAnimator.animating = false;
+                appToken.removeIfPossible();
+            }
+
+            appToken.removed = true;
+            if (appToken.startingData != null) {
+                startingToken = appToken;
+            }
+            appToken.stopFreezingScreen(true, true);
+            if (mService.mFocusedApp == appToken) {
+                if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + appToken);
+                mService.mFocusedApp = null;
+                mService.updateFocusedWindowLocked(
+                        UPDATE_FOCUS_NORMAL, true /*updateInputWindows*/);
+                mService.mInputMonitor.setFocusedAppLw(null);
+            }
+
+            if (!delayed) {
+                appToken.updateReportedVisibilityLocked();
+            }
+
+            // Will only remove if startingToken non null.
+            mService.scheduleRemoveStartingWindowLocked(startingToken);
+        }
+    }
+
     // TODO: Users would have their own window containers under the display container?
     void switchUser() {
         final int count = mChildren.size();
@@ -1361,6 +1572,13 @@
         }
     }
 
+    void dumpTokens(PrintWriter pw, boolean dumpAll) {
+        pw.println("  All tokens:");
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            mChildren.get(i).dumpTokens(pw, dumpAll);
+        }
+    }
+
     @Override
     String getName() {
         return "ROOT";
diff --git a/services/core/java/com/android/server/wm/TaskStack.java b/services/core/java/com/android/server/wm/TaskStack.java
index e98fc39..56fdc33 100644
--- a/services/core/java/com/android/server/wm/TaskStack.java
+++ b/services/core/java/com/android/server/wm/TaskStack.java
@@ -94,16 +94,16 @@
     private boolean mFillsParent = true;
 
     // Device rotation as of the last time {@link #mBounds} was set.
-    int mRotation;
+    private int mRotation;
 
     /** Density as of last time {@link #mBounds} was set. */
-    int mDensity;
+    private int mDensity;
 
     /** Support for non-zero {@link android.view.animation.Animation#getBackgroundColor()} */
-    DimLayer mAnimationBackgroundSurface;
+    private DimLayer mAnimationBackgroundSurface;
 
     /** The particular window with an Animation with non-zero background color. */
-    WindowStateAnimator mAnimationBackgroundAnimator;
+    private WindowStateAnimator mAnimationBackgroundAnimator;
 
     /** Application tokens that are exiting, but still on screen for animations. */
     final AppTokenList mExitingAppTokens = new AppTokenList();
@@ -606,12 +606,12 @@
         }
     }
 
-    void attachDisplayContent(DisplayContent displayContent) {
+    void onDisplayChanged(DisplayContent dc) {
         if (mDisplayContent != null) {
-            throw new IllegalStateException("attachDisplayContent: Already attached");
+            throw new IllegalStateException("onDisplayChanged: Already attached");
         }
 
-        mDisplayContent = displayContent;
+        mDisplayContent = dc;
         mAnimationBackgroundSurface = new DimLayer(mService, this, mDisplayContent.getDisplayId(),
                 "animation background stackId=" + mStackId);
 
@@ -625,7 +625,7 @@
             // not fullscreen. If it's fullscreen, it means that we are in the transition of
             // dismissing it, so we must not resize this stack.
             bounds = new Rect();
-            displayContent.getLogicalDisplayRect(mTmpRect);
+            dc.getLogicalDisplayRect(mTmpRect);
             mTmpRect2.setEmpty();
             if (dockedStack != null) {
                 dockedStack.getRawBounds(mTmpRect2);
@@ -638,6 +638,8 @@
         }
 
         updateDisplayInfo(bounds);
+
+        super.onDisplayChanged(dc);
     }
 
     void getStackDockedModeBoundsLocked(Rect outBounds, boolean ignoreVisibility) {
diff --git a/services/core/java/com/android/server/wm/WindowContainer.java b/services/core/java/com/android/server/wm/WindowContainer.java
index 902f4ae..56bf6cd 100644
--- a/services/core/java/com/android/server/wm/WindowContainer.java
+++ b/services/core/java/com/android/server/wm/WindowContainer.java
@@ -254,6 +254,17 @@
         }
     }
 
+    /**
+     * Notify that the display this container is on has changed.
+     * @param dc The new display this container is on.
+     */
+    void onDisplayChanged(DisplayContent dc) {
+        for (int i = mChildren.size() - 1; i >= 0; --i) {
+            final WindowContainer child = mChildren.get(i);
+            child.onDisplayChanged(dc);
+        }
+    }
+
     void setWaitingForDrawnIfResizingChanged() {
         for (int i = mChildren.size() - 1; i >= 0; --i) {
             final WindowContainer wc = mChildren.get(i);
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 12ad09c..b02d1f7 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -163,7 +163,6 @@
 import java.util.ArrayList;
 import java.util.Date;
 import java.util.HashMap;
-import java.util.Iterator;
 import java.util.List;
 
 import static android.Manifest.permission.MANAGE_APP_TOKENS;
@@ -403,11 +402,6 @@
     final HashMap<IBinder, WindowState> mWindowMap = new HashMap<>();
 
     /**
-     * Mapping from a token IBinder to a WindowToken object.
-     */
-    final HashMap<IBinder, WindowToken> mTokenMap = new HashMap<>();
-
-    /**
      * List of window tokens that have finished starting their application,
      * and now need to have the policy remove their windows.
      */
@@ -912,7 +906,7 @@
         @Override
         public void onAppTransitionFinishedLocked(IBinder token) {
             mH.sendEmptyMessage(H.NOTIFY_APP_TRANSITION_FINISHED);
-            AppWindowToken atoken = findAppWindowToken(token);
+            final AppWindowToken atoken = mRoot.getAppWindowToken(token);
             if (atoken == null) {
                 return;
             }
@@ -1535,7 +1529,8 @@
             final boolean hasParent = parentWindow != null;
             // Use existing parent window token for child windows since they go in the same token
             // as there parent window so we can apply the same policy on them.
-            WindowToken token = mTokenMap.get(hasParent ? parentWindow.mAttrs.token : attrs.token);
+            WindowToken token = mRoot.getWindowToken(
+                    hasParent ? parentWindow.mAttrs.token : attrs.token, displayContent);
             // If this is a child window, we want to apply the same type checking rules as the
             // parent window type.
             final int rootType = hasParent ? parentWindow.mAttrs.type : type;
@@ -1587,7 +1582,7 @@
                         return WindowManagerGlobal.ADD_BAD_APP_TOKEN;
                     }
                 }
-                token = new WindowToken(this, attrs.token, -1, false);
+                token = new WindowToken(this, attrs.token, -1, false, displayContent);
             } else if (rootType >= FIRST_APPLICATION_WINDOW && rootType <= LAST_APPLICATION_WINDOW) {
                 atoken = token.asAppWindowToken();
                 if (atoken == null) {
@@ -1655,7 +1650,7 @@
                 // It is not valid to use an app token with other system types; we will
                 // instead make a new token for it (as if null had been passed in for the token).
                 attrs.token = null;
-                token = new WindowToken(this, null, -1, false);
+                token = new WindowToken(this, null, -1, false, displayContent);
             }
 
             WindowState win = new WindowState(this, session, client, token, parentWindow,
@@ -2025,7 +2020,7 @@
         // Window will already be removed from token before this post clean-up method is called.
         if (token.isEmpty()) {
             if (!token.explicit) {
-                mTokenMap.remove(token.token);
+                token.removeImmediately();
             } else if (atoken != null) {
                 // TODO: Should this be moved into AppWindowToken.removeWindow? Might go away after
                 // re-factor.
@@ -2750,14 +2745,6 @@
         return !mDisplayFrozen && mDisplayEnabled && mPolicy.isScreenOn();
     }
 
-    AppWindowToken findAppWindowToken(IBinder token) {
-        WindowToken wtoken = mTokenMap.get(token);
-        if (wtoken == null) {
-            return null;
-        }
-        return wtoken.asAppWindowToken();
-    }
-
     @Override
     public void addWindowToken(IBinder token, int type) {
         if (!checkCallingPermission(MANAGE_APP_TOKENS, "addWindowToken()")) {
@@ -2765,15 +2752,7 @@
         }
 
         synchronized(mWindowMap) {
-            WindowToken wtoken = mTokenMap.get(token);
-            if (wtoken != null) {
-                Slog.w(TAG_WM, "Attempted to add existing input method token: " + token);
-                return;
-            }
-            wtoken = new WindowToken(this, token, type, true);
-            if (type == TYPE_WALLPAPER) {
-                mWallpaperControllerLocked.addWallpaperToken(wtoken);
-            }
+            mRoot.addWindowToken(token, type);
         }
     }
 
@@ -2784,20 +2763,28 @@
         }
 
         final long origId = Binder.clearCallingIdentity();
-        synchronized(mWindowMap) {
-            final WindowToken wtoken = mTokenMap.remove(token);
-            if (wtoken != null) {
-                wtoken.setExiting();
-                if (wtoken.windowType == TYPE_WALLPAPER) {
-                    mWallpaperControllerLocked.removeWallpaperToken(wtoken);
+        try {
+            synchronized (mWindowMap) {
+                final ArrayList<WindowToken> removedTokens = mRoot.removeWindowToken(token);
+                if (removedTokens == null || removedTokens.isEmpty()) {
+                    Slog.w(TAG_WM,
+                            "removeWindowToken: Attempted to remove non-existing token: " + token);
+                    return;
                 }
 
-                mInputMonitor.updateInputWindowsLw(true /*force*/);
-            } else {
-                Slog.w(TAG_WM, "Attempted to remove non-existing token: " + token);
+                for (int i = removedTokens.size() - 1; i >= 0; --i) {
+                    final WindowToken wtoken = removedTokens.get(i);
+                    wtoken.setExiting();
+                    if (wtoken.windowType == TYPE_WALLPAPER) {
+                        mWallpaperControllerLocked.removeWallpaperToken(wtoken);
+                    }
+
+                    mInputMonitor.updateInputWindowsLw(true /*force*/);
+                }
             }
+        } finally {
+            Binder.restoreCallingIdentity(origId);
         }
-        Binder.restoreCallingIdentity(origId);
     }
 
     private Task createTaskLocked(int taskId, int stackId, int userId, AppWindowToken atoken,
@@ -2841,12 +2828,18 @@
         }
 
         synchronized(mWindowMap) {
-            AppWindowToken atoken = findAppWindowToken(token.asBinder());
+            AppWindowToken atoken = mRoot.getAppWindowToken(token.asBinder());
             if (atoken != null) {
                 Slog.w(TAG_WM, "Attempted to add existing app token: " + token);
                 return;
             }
-            atoken = new AppWindowToken(this, token, voiceInteraction);
+
+            final TaskStack stack = mStackIdToStack.get(stackId);
+            if (stack == null) {
+                throw new IllegalArgumentException("addAppToken: invalid stackId=" + stackId);
+            }
+
+            atoken = new AppWindowToken(this, token, voiceInteraction, stack.getDisplayContent());
             atoken.inputDispatchingTimeoutNanos = inputDispatchingTimeoutNanos;
             atoken.setFillsParent(fullscreen);
             atoken.showForAllUsers = showForAllUsers;
@@ -2882,7 +2875,7 @@
         }
 
         synchronized(mWindowMap) {
-            final AppWindowToken atoken = findAppWindowToken(token);
+            final AppWindowToken atoken = mRoot.getAppWindowToken(token);
             if (atoken == null) {
                 Slog.w(TAG_WM, "Attempted to set task id of non-existing app token: " + token);
                 return;
@@ -2998,12 +2991,10 @@
         Configuration config = null;
 
         if (updateOrientationFromAppTokensLocked(false)) {
-            // If we changed the orientation but mOrientationChangeComplete is
-            // already true, we used seamless rotation, and we don't need
-            // to freeze the screen.
-            if (freezeThisOneIfNeeded != null &&
-                    !mRoot.mOrientationChangeComplete) {
-                final AppWindowToken atoken = findAppWindowToken(freezeThisOneIfNeeded);
+            // If we changed the orientation but mOrientationChangeComplete is already true,
+            // we used seamless rotation, and we don't need to freeze the screen.
+            if (freezeThisOneIfNeeded != null && !mRoot.mOrientationChangeComplete) {
+                final AppWindowToken atoken = mRoot.getAppWindowToken(freezeThisOneIfNeeded);
                 if (atoken != null) {
                     atoken.startFreezingScreen();
                 }
@@ -3117,7 +3108,7 @@
         }
 
         synchronized(mWindowMap) {
-            final AppWindowToken atoken = findAppWindowToken(token.asBinder());
+            final AppWindowToken atoken = mRoot.getAppWindowToken(token.asBinder());
             if (atoken == null) {
                 Slog.w(TAG_WM, "Attempted to set orientation of non-existing app token: " + token);
                 return;
@@ -3130,9 +3121,9 @@
     @Override
     public int getAppOrientation(IApplicationToken token) {
         synchronized(mWindowMap) {
-            AppWindowToken wtoken = findAppWindowToken(token.asBinder());
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token.asBinder());
             if (wtoken == null) {
-                return ActivityInfo.SCREEN_ORIENTATION_UNSPECIFIED;
+                return SCREEN_ORIENTATION_UNSPECIFIED;
             }
 
             return wtoken.getOrientation();
@@ -3161,7 +3152,7 @@
                 if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Clearing focused app, was " + mFocusedApp);
                 newFocus = null;
             } else {
-                newFocus = findAppWindowToken(token);
+                newFocus = mRoot.getAppWindowToken(token);
                 if (newFocus == null) {
                     Slog.w(TAG_WM, "Attempted to set focus to non-existing app token: " + token);
                 }
@@ -3353,7 +3344,7 @@
                     TAG_WM, "setAppStartingWindow: token=" + token + " pkg=" + pkg
                     + " transferFrom=" + transferFrom);
 
-            AppWindowToken wtoken = findAppWindowToken(token);
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null) {
                 Slog.w(TAG_WM, "Attempted to set icon of non-existing app token: " + token);
                 return false;
@@ -3440,14 +3431,14 @@
 
     public void removeAppStartingWindow(IBinder token) {
         synchronized (mWindowMap) {
-            final AppWindowToken wtoken = mTokenMap.get(token).asAppWindowToken();
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
             scheduleRemoveStartingWindowLocked(wtoken);
         }
     }
 
     public void setAppFullscreen(IBinder token, boolean toOpaque) {
         synchronized (mWindowMap) {
-            final AppWindowToken atoken = findAppWindowToken(token);
+            final AppWindowToken atoken = mRoot.getAppWindowToken(token);
             if (atoken != null) {
                 atoken.setFillsParent(toOpaque);
                 setWindowOpaqueLocked(token, toOpaque);
@@ -3463,7 +3454,7 @@
     }
 
     private void setWindowOpaqueLocked(IBinder token, boolean isOpaque) {
-        final AppWindowToken wtoken = findAppWindowToken(token);
+        final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
         if (wtoken != null) {
             final WindowState win = wtoken.findMainWindow();
             if (win != null) {
@@ -3488,8 +3479,7 @@
         }
 
         synchronized(mWindowMap) {
-            final AppWindowToken wtoken;
-            wtoken = findAppWindowToken(token);
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null) {
                 Slog.w(TAG_WM, "Attempted to notify resumed of non-existing app token: " + token);
                 return;
@@ -3506,7 +3496,7 @@
 
         synchronized(mWindowMap) {
             final AppWindowToken wtoken;
-            wtoken = findAppWindowToken(token);
+            wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null) {
                 Slog.w(TAG_WM, "Attempted to notify stopped of non-existing app token: " + token);
                 return;
@@ -3524,7 +3514,7 @@
         AppWindowToken wtoken;
 
         synchronized(mWindowMap) {
-            wtoken = findAppWindowToken(token);
+            wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null) {
                 Slog.w(TAG_WM, "Attempted to set visibility of non-existing app token: " + token);
                 return;
@@ -3644,7 +3634,7 @@
                 return;
             }
 
-            final AppWindowToken wtoken = findAppWindowToken(token);
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null || wtoken.appToken == null) {
                 Slog.w(TAG_WM, "Attempted to freeze screen with non-existing app token: " + wtoken);
                 return;
@@ -3662,7 +3652,7 @@
         }
 
         synchronized(mWindowMap) {
-            final AppWindowToken wtoken = findAppWindowToken(token);
+            final AppWindowToken wtoken = mRoot.getAppWindowToken(token);
             if (wtoken == null || wtoken.appToken == null) {
                 return;
             }
@@ -3680,71 +3670,14 @@
             throw new SecurityException("Requires MANAGE_APP_TOKENS permission");
         }
 
-        AppWindowToken wtoken = null;
-        AppWindowToken startingToken = null;
-        boolean delayed = false;
-
         final long origId = Binder.clearCallingIdentity();
-        synchronized(mWindowMap) {
-            WindowToken basewtoken = mTokenMap.remove(token);
-            if (basewtoken != null && (wtoken = basewtoken.asAppWindowToken()) != null) {
-                if (DEBUG_APP_TRANSITIONS) Slog.v(TAG_WM, "Removing app token: " + wtoken);
-                delayed = wtoken.setVisibility(null, false,
-                        TRANSIT_UNSET, true, wtoken.voiceInteraction);
-                mOpeningApps.remove(wtoken);
-                wtoken.waitingToShow = false;
-                if (mClosingApps.contains(wtoken)) {
-                    delayed = true;
-                } else if (mAppTransition.isTransitionSet()) {
-                    mClosingApps.add(wtoken);
-                    delayed = true;
-                }
-                if (DEBUG_APP_TRANSITIONS) Slog.v(
-                        TAG_WM, "Removing app " + wtoken + " delayed=" + delayed
-                        + " animation=" + wtoken.mAppAnimator.animation
-                        + " animating=" + wtoken.mAppAnimator.animating);
-                if (DEBUG_ADD_REMOVE || DEBUG_TOKEN_MOVEMENT) Slog.v(TAG_WM, "removeAppToken: "
-                        + wtoken + " delayed=" + delayed + " Callers=" + Debug.getCallers(4));
-                final TaskStack stack = wtoken.mTask.mStack;
-                if (delayed && !wtoken.isEmpty()) {
-                    // set the token aside because it has an active animation to be finished
-                    if (DEBUG_ADD_REMOVE || DEBUG_TOKEN_MOVEMENT) Slog.v(TAG_WM,
-                            "removeAppToken make exiting: " + wtoken);
-                    stack.mExitingAppTokens.add(wtoken);
-                    wtoken.mIsExiting = true;
-                } else {
-                    // Make sure there is no animation running on this token,
-                    // so any windows associated with it will be removed as
-                    // soon as their animations are complete
-                    wtoken.mAppAnimator.clearAnimation();
-                    wtoken.mAppAnimator.animating = false;
-                    wtoken.removeIfPossible();
-                }
-
-                wtoken.removed = true;
-                if (wtoken.startingData != null) {
-                    startingToken = wtoken;
-                }
-                wtoken.stopFreezingScreen(true, true);
-                if (mFocusedApp == wtoken) {
-                    if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + wtoken);
-                    mFocusedApp = null;
-                    updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, true /*updateInputWindows*/);
-                    mInputMonitor.setFocusedAppLw(null);
-                }
-            } else {
-                Slog.w(TAG_WM, "Attempted to remove non-existing app token: " + token);
+        try {
+            synchronized(mWindowMap) {
+                mRoot.removeAppToken(token);
             }
-
-            if (!delayed && wtoken != null) {
-                wtoken.updateReportedVisibilityLocked();
-            }
-
-            // Will only remove if startingToken non null.
-            scheduleRemoveStartingWindowLocked(startingToken);
+        } finally {
+            Binder.restoreCallingIdentity(origId);
         }
-        Binder.restoreCallingIdentity(origId);
-
     }
 
     void scheduleRemoveStartingWindowLocked(AppWindowToken wtoken) {
@@ -5542,7 +5475,7 @@
 
         if (mRotation == rotation && mAltOrientation == altOrientation) {
             // No change.
-             return false;
+            return false;
         }
 
         if (DEBUG_ORIENTATION) {
@@ -6616,9 +6549,11 @@
         }
 
         synchronized (mWindowMap) {
-            WindowToken token = mTokenMap.get(_token);
-            if (token != null) {
-                mInputMonitor.pauseDispatchingLw(token);
+            final ArrayList<WindowToken> tokens = mRoot.getWindowTokens(_token);
+            if (tokens != null && !tokens.isEmpty()) {
+                for (int i = tokens.size() - 1; i >= 0; --i) {
+                    mInputMonitor.pauseDispatchingLw(tokens.get(i));
+                }
             }
         }
     }
@@ -6630,9 +6565,11 @@
         }
 
         synchronized (mWindowMap) {
-            WindowToken token = mTokenMap.get(_token);
-            if (token != null) {
-                mInputMonitor.resumeDispatchingLw(token);
+            final ArrayList<WindowToken> tokens = mRoot.getWindowTokens(_token);
+            if (tokens != null && !tokens.isEmpty()) {
+                for (int i = tokens.size() - 1; i >= 0; --i) {
+                    mInputMonitor.resumeDispatchingLw(tokens.get(i));
+                }
             }
         }
     }
@@ -8661,7 +8598,7 @@
 
     public void notifyAppRelaunching(IBinder token) {
         synchronized (mWindowMap) {
-            AppWindowToken appWindow = findAppWindowToken(token);
+            final AppWindowToken appWindow = mRoot.getAppWindowToken(token);
             if (appWindow != null) {
                 appWindow.startRelaunching();
             }
@@ -8670,7 +8607,7 @@
 
     public void notifyAppRelaunchingFinished(IBinder token) {
         synchronized (mWindowMap) {
-            AppWindowToken appWindow = findAppWindowToken(token);
+            final AppWindowToken appWindow = mRoot.getAppWindowToken(token);
             if (appWindow != null) {
                 appWindow.finishRelaunching();
             }
@@ -8679,7 +8616,7 @@
 
     public void notifyAppRelaunchesCleared(IBinder token) {
         synchronized (mWindowMap) {
-            final AppWindowToken appWindow = findAppWindowToken(token);
+            final AppWindowToken appWindow = mRoot.getAppWindowToken(token);
             if (appWindow != null) {
                 appWindow.clearRelaunching();
             }
@@ -8703,20 +8640,7 @@
 
     private void dumpTokensLocked(PrintWriter pw, boolean dumpAll) {
         pw.println("WINDOW MANAGER TOKENS (dumpsys window tokens)");
-        if (!mTokenMap.isEmpty()) {
-            pw.println("  All tokens:");
-            Iterator<WindowToken> it = mTokenMap.values().iterator();
-            while (it.hasNext()) {
-                WindowToken token = it.next();
-                pw.print("  "); pw.print(token);
-                if (dumpAll) {
-                    pw.println(':');
-                    token.dump(pw, "    ");
-                } else {
-                    pw.println();
-                }
-            }
-        }
+        mRoot.dumpTokens(pw, dumpAll);
         mWallpaperControllerLocked.dumpTokens(pw, "  ", dumpAll);
         if (!mFinishedStarting.isEmpty()) {
             pw.println();
@@ -9248,7 +9172,7 @@
      */
     public void setWillReplaceWindow(IBinder token, boolean animate) {
         synchronized (mWindowMap) {
-            final AppWindowToken appWindowToken = findAppWindowToken(token);
+            final AppWindowToken appWindowToken = mRoot.getAppWindowToken(token);
             if (appWindowToken == null || !appWindowToken.hasContentToDisplay()) {
                 Slog.w(TAG_WM, "Attempted to set replacing window on non-existing app token "
                         + token);
@@ -9270,9 +9194,9 @@
      */
     // TODO: The s at the end of the method name is the only difference with the name of the method
     // above. We should combine them or find better names.
-    public void setWillReplaceWindows(IBinder token, boolean childrenOnly) {
+    void setWillReplaceWindows(IBinder token, boolean childrenOnly) {
         synchronized (mWindowMap) {
-            final AppWindowToken appWindowToken = findAppWindowToken(token);
+            final AppWindowToken appWindowToken = mRoot.getAppWindowToken(token);
             if (appWindowToken == null || !appWindowToken.hasContentToDisplay()) {
                 Slog.w(TAG_WM, "Attempted to set replacing window on non-existing app token "
                         + token);
@@ -9300,7 +9224,7 @@
      */
     public void scheduleClearWillReplaceWindows(IBinder token, boolean replacing) {
         synchronized (mWindowMap) {
-            final AppWindowToken appWindowToken = findAppWindowToken(token);
+            final AppWindowToken appWindowToken = mRoot.getAppWindowToken(token);
             if (appWindowToken == null) {
                 Slog.w(TAG_WM, "Attempted to reset replacing window on non-existing app token "
                         + token);
@@ -9765,9 +9689,12 @@
         public void removeWindowToken(IBinder token, boolean removeWindows) {
             synchronized(mWindowMap) {
                 if (removeWindows) {
-                    final WindowToken wtoken = mTokenMap.remove(token);
-                    if (wtoken != null) {
-                        wtoken.removeAllWindows();
+                    final ArrayList<WindowToken> removedTokens = mRoot.removeWindowToken(token);
+                    if (removedTokens != null && !removedTokens.isEmpty()) {
+                        for (int i = removedTokens.size() - 1; i >= 0; --i) {
+                            final WindowToken wtoken = removedTokens.get(i);
+                            wtoken.removeAllWindows();
+                        }
                     }
                 }
                 WindowManagerService.this.removeWindowToken(token);
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 184b6cd..a784a8f 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -469,6 +469,7 @@
     boolean mHasSurface = false;
 
     private boolean mNotOnAppsDisplay = false;
+    // TODO: Get this from WindowToken? Done in ag/1522894
     DisplayContent  mDisplayContent;
 
     /** When true this window can be displayed on screens owther than mOwnerUid's */
diff --git a/services/core/java/com/android/server/wm/WindowToken.java b/services/core/java/com/android/server/wm/WindowToken.java
index 166107c..28c9e5e 100644
--- a/services/core/java/com/android/server/wm/WindowToken.java
+++ b/services/core/java/com/android/server/wm/WindowToken.java
@@ -78,12 +78,16 @@
     // windows will be put to the bottom of the list.
     boolean sendingToBottom;
 
-    WindowToken(WindowManagerService service, IBinder _token, int type, boolean _explicit) {
+    // The display this token is on.
+    private DisplayContent mDisplayContent;
+
+    WindowToken(WindowManagerService service, IBinder _token, int type, boolean _explicit,
+            DisplayContent dc) {
         mService = service;
         token = _token;
         windowType = type;
         explicit = _explicit;
-        mService.mTokenMap.put(token, this);
+        onDisplayChanged(dc);
     }
 
     void removeAllWindows() {
@@ -400,6 +404,33 @@
         return null;
     }
 
+    DisplayContent getDisplayContent() {
+        return mDisplayContent;
+    }
+
+    @Override
+    void removeImmediately() {
+        super.removeImmediately();
+        if (mDisplayContent != null) {
+            mDisplayContent.removeWindowToken(token);
+            mService.mRoot.removeWindowTokenIfPossible(token);
+        }
+    }
+
+    void onDisplayChanged(DisplayContent dc) {
+        if (mDisplayContent == dc) {
+            return;
+        }
+
+        if (mDisplayContent != null) {
+            mDisplayContent.removeWindowToken(token);
+        }
+        mDisplayContent = dc;
+        mDisplayContent.setWindowToken(token, this);
+
+        super.onDisplayChanged(dc);
+    }
+
     void dump(PrintWriter pw, String prefix) {
         pw.print(prefix); pw.print("windows="); pw.println(mChildren);
         pw.print(prefix); pw.print("windowType="); pw.print(windowType);