diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index f268aa2..f4bb479 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -280,10 +280,10 @@
 
     WindowState findMainWindow() {
         WindowState candidate = null;
-        int j = windows.size();
+        int j = allAppWindows.size();
         while (j > 0) {
             j--;
-            WindowState win = windows.get(j);
+            WindowState win = allAppWindows.get(j);
             if (win.mAttrs.type == WindowManager.LayoutParams.TYPE_BASE_APPLICATION
                     || win.mAttrs.type == WindowManager.LayoutParams.TYPE_APPLICATION_STARTING) {
                 // In cases where there are multiple windows, we prefer the non-exiting window. This
@@ -542,7 +542,7 @@
         clearVisibleBeforeClientHidden();
 
         if (DEBUG_APP_TRANSITIONS || DEBUG_ANIM) Slog.d(TAG,
-                "restoreSavedSurfaces: " + appWindowToken + " allDrawn=" + allDrawn
+                "restoreSavedSurfaces: " + this + " allDrawn=" + allDrawn
                 + " numInteresting=" + numInteresting + " numDrawn=" + numDrawn);
     }
 
@@ -609,8 +609,8 @@
     }
 
     void setReplacingWindows(boolean animate) {
-        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM, "Marking app token " + appWindowToken
-                + " with replacing windows.");
+        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM,
+                "Marking app token " + this + " with replacing windows.");
 
         for (int i = allAppWindows.size() - 1; i >= 0; i--) {
             final WindowState w = allAppWindows.get(i);
@@ -627,7 +627,7 @@
     }
 
     void setReplacingChildren() {
-        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM, "Marking app token " + appWindowToken
+        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM, "Marking app token " + this
                 + " with replacing child windows.");
         for (int i = allAppWindows.size() - 1; i >= 0; i--) {
             final WindowState w = allAppWindows.get(i);
@@ -638,8 +638,8 @@
     }
 
     void resetReplacingWindows() {
-        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM, "Resetting app token " + appWindowToken
-                + " of replacing window marks.");
+        if (DEBUG_ADD_REMOVE) Slog.d(TAG_WM,
+                "Resetting app token " + this + " of replacing window marks.");
 
         for (int i = allAppWindows.size() - 1; i >= 0; i--) {
             final WindowState w = allAppWindows.get(i);
@@ -769,8 +769,8 @@
         if (!mFrozenMergedConfig.isEmpty()) {
             mFrozenMergedConfig.remove();
         }
-        for (int i = windows.size() - 1; i >= 0; i--) {
-            final WindowState win = windows.get(i);
+        for (int i = allAppWindows.size() - 1; i >= 0; i--) {
+            final WindowState win = allAppWindows.get(i);
             if (!win.mHasSurface) {
                 continue;
             }
diff --git a/services/core/java/com/android/server/wm/WindowManagerService.java b/services/core/java/com/android/server/wm/WindowManagerService.java
index 292561a..0493084 100644
--- a/services/core/java/com/android/server/wm/WindowManagerService.java
+++ b/services/core/java/com/android/server/wm/WindowManagerService.java
@@ -1621,7 +1621,6 @@
     void addInputMethodWindowToListLocked(WindowState win) {
         int pos = findDesiredInputMethodWindowIndexLocked(true);
         if (pos >= 0) {
-            win.mTargetAppToken = mInputMethodTarget.mAppToken;
             if (DEBUG_WINDOW_MOVEMENT || DEBUG_ADD_REMOVE) Slog.v(
                     TAG_WM, "Adding input method window " + win + " at " + pos);
             // TODO(multidisplay): IMEs are only supported on the default display.
@@ -1630,7 +1629,6 @@
             moveInputMethodDialogsLocked(pos + 1);
             return;
         }
-        win.mTargetAppToken = null;
         addWindowToListInOrderLocked(win, true);
         moveInputMethodDialogsLocked(pos);
     }
@@ -1674,7 +1672,6 @@
         }
 
         if (pos >= 0) {
-            final AppWindowToken targetAppToken = mInputMethodTarget.mAppToken;
             // Skip windows owned by the input method.
             if (mInputMethodWindow != null) {
                 while (pos < windows.size()) {
@@ -1689,7 +1686,6 @@
             if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "Adding " + N + " dialogs at pos=" + pos);
             for (int i=0; i<N; i++) {
                 WindowState win = dialogs.get(i);
-                win.mTargetAppToken = targetAppToken;
                 pos = win.reAddWindowLocked(pos);
             }
             if (DEBUG_INPUT_METHOD) {
@@ -1700,7 +1696,6 @@
         }
         for (int i=0; i<N; i++) {
             WindowState win = dialogs.get(i);
-            win.mTargetAppToken = null;
             reAddWindowToListInOrderLocked(win);
             if (DEBUG_INPUT_METHOD) {
                 Slog.v(TAG_WM, "No IM target, final list:");
@@ -1756,11 +1751,6 @@
                     pos++;
                 }
                 if (pos >= N) {
-                    // Z order is good.
-                    // The IM target window may be changed, so update the mTargetAppToken.
-                    if (imWin != null) {
-                        imWin.mTargetAppToken = mInputMethodTarget.mAppToken;
-                    }
                     return false;
                 }
             }
@@ -1775,7 +1765,6 @@
                     Slog.v(TAG_WM, "List after removing with new pos " + imPos + ":");
                     logWindowList(windows, "  ");
                 }
-                imWin.mTargetAppToken = mInputMethodTarget.mAppToken;
                 imWin.reAddWindowLocked(imPos);
                 if (DEBUG_INPUT_METHOD) {
                     Slog.v(TAG_WM, "List after moving IM to " + imPos + ":");
@@ -1793,7 +1782,6 @@
             if (imWin != null) {
                 if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "Moving IM from " + imPos);
                 imWin.removeFromWindowList(0);
-                imWin.mTargetAppToken = null;
                 reAddWindowToListInOrderLocked(imWin);
                 if (DEBUG_INPUT_METHOD) {
                     Slog.v(TAG_WM, "List with no IM target:");
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index f0054f8..d20529d 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -158,7 +158,6 @@
     WindowToken mToken;
     WindowToken mRootToken;
     AppWindowToken mAppToken;
-    AppWindowToken mTargetAppToken;
 
     // mAttrs.flags is tested in animation without being locked. If the bits tested are ever
     // modified they will need to be locked.
@@ -619,17 +618,17 @@
         }
         mIsFloatingLayer = mIsImWindow || mIsWallpaper;
 
-        final WindowState appWin = getTopParentWindow();
-        WindowToken appToken = appWin.mToken;
-        while (appToken.appWindowToken == null) {
-            WindowToken parent = mService.mTokenMap.get(appToken.token);
-            if (parent == null || appToken == parent) {
+        final WindowState topParentWindow = getTopParentWindow();
+        WindowToken rootToken = topParentWindow.mToken;
+        while (rootToken.appWindowToken == null) {
+            WindowToken parent = mService.mTokenMap.get(rootToken.token);
+            if (parent == null || rootToken == parent) {
                 break;
             }
-            appToken = parent;
+            rootToken = parent;
         }
-        mRootToken = appToken;
-        mAppToken = appToken.appWindowToken;
+        mRootToken = rootToken;
+        mAppToken = rootToken.appWindowToken;
         if (mAppToken != null) {
             final DisplayContent appDisplay = getDisplayContent();
             mNotOnAppsDisplay = displayContent != appDisplay;
@@ -1495,14 +1494,21 @@
     }
 
     int getAnimLayerAdjustment() {
-        if (mTargetAppToken != null) {
-            return mTargetAppToken.mAppAnimator.animLayerAdjustment;
-        } else if (mAppToken != null) {
-            return mAppToken.mAppAnimator.animLayerAdjustment;
-        } else {
-            // Nothing is animating, so there is no animation adjustment.
-            return 0;
+        final boolean isImeType =
+                mAttrs.type == TYPE_INPUT_METHOD || mAttrs.type == TYPE_INPUT_METHOD_DIALOG;
+        if (isImeType && mService.mInputMethodTarget != null) {
+            final AppWindowToken appToken = mService.mInputMethodTarget.mAppToken;
+            if (appToken != null) {
+                return appToken.mAppAnimator.animLayerAdjustment;
+            }
         }
+
+        if (mAppToken != null) {
+            return mAppToken.mAppAnimator.animLayerAdjustment;
+        }
+
+        // Nothing is animating, so there is no animation adjustment.
+        return 0;
     }
 
     void scheduleAnimationIfDimming() {
@@ -2580,9 +2586,7 @@
             pw.print(prefix); pw.print("mBaseLayer="); pw.print(mBaseLayer);
                     pw.print(" mSubLayer="); pw.print(mSubLayer);
                     pw.print(" mAnimLayer="); pw.print(mLayer); pw.print("+");
-                    pw.print((mTargetAppToken != null ?
-                            mTargetAppToken.mAppAnimator.animLayerAdjustment
-                          : (mAppToken != null ? mAppToken.mAppAnimator.animLayerAdjustment : 0)));
+                    pw.print(getAnimLayerAdjustment());
                     pw.print("="); pw.print(mWinAnimator.mAnimLayer);
                     pw.print(" mLastLayer="); pw.println(mWinAnimator.mLastLayer);
         }
@@ -2595,9 +2599,6 @@
                 pw.print(isAnimatingWithSavedSurface());
                 pw.print(" mAppDied=");pw.println(mAppDied);
             }
-            if (mTargetAppToken != null) {
-                pw.print(prefix); pw.print("mTargetAppToken="); pw.println(mTargetAppToken);
-            }
             pw.print(prefix); pw.print("mViewVisibility=0x");
             pw.print(Integer.toHexString(mViewVisibility));
             pw.print(" mHaveFrame="); pw.print(mHaveFrame);