Prevent IME from targeting home when closing app is still animating

There were a few issues that caused the IME target to get set to home.
1. Home is considered above the other apps since it's getting launched.
However, the visual representation is actually below the closing app.
IME tries to target the highest app, but it's not necessarily the
highest visual app
2. computeIme was called before all apps were set up. This caused the
IME target to get recomputed before adding the closing apps
3. The condition for a WS.isClosing was based on previous versions so
it was incorrect

Test: IME closes on top of current target when going home with IME open
Change-Id: Ia9b0cc913eeb76a4af20f9ac7b667c5a3d1eafc7
Fixes: 78214125
diff --git a/services/core/java/com/android/server/am/ActivityDisplay.java b/services/core/java/com/android/server/am/ActivityDisplay.java
index 27eae57..698b6f7 100644
--- a/services/core/java/com/android/server/am/ActivityDisplay.java
+++ b/services/core/java/com/android/server/am/ActivityDisplay.java
@@ -800,6 +800,20 @@
         }
     }
 
+    /**
+     * See {@link DisplayWindowController#deferUpdateImeTarget()}
+     */
+    public void deferUpdateImeTarget() {
+        mWindowContainerController.deferUpdateImeTarget();
+    }
+
+    /**
+     * See {@link DisplayWindowController#deferUpdateImeTarget()}
+     */
+    public void continueUpdateImeTarget() {
+        mWindowContainerController.continueUpdateImeTarget();
+    }
+
     public void dump(PrintWriter pw, String prefix) {
         pw.println(prefix + "displayId=" + mDisplayId + " stacks=" + mStacks.size());
         final String myPrefix = prefix + " ";
diff --git a/services/core/java/com/android/server/am/ActivityStack.java b/services/core/java/com/android/server/am/ActivityStack.java
index 87d3696..c182502 100644
--- a/services/core/java/com/android/server/am/ActivityStack.java
+++ b/services/core/java/com/android/server/am/ActivityStack.java
@@ -4595,46 +4595,58 @@
             }
         }
 
-        // Shift all activities with this task up to the top
-        // of the stack, keeping them in the same internal order.
-        insertTaskAtTop(tr, null);
+        try {
+            // Defer updating the IME target since the new IME target will try to get computed
+            // before updating all closing and opening apps, which can cause the ime target to
+            // get calculated incorrectly.
+            getDisplay().deferUpdateImeTarget();
 
-        // Don't refocus if invisible to current user
-        final ActivityRecord top = tr.getTopActivity();
-        if (top == null || !top.okToShowLocked()) {
-            if (top != null) {
-                mStackSupervisor.mRecentTasks.add(top.getTask());
+            // Shift all activities with this task up to the top
+            // of the stack, keeping them in the same internal order.
+            insertTaskAtTop(tr, null);
+
+            // Don't refocus if invisible to current user
+            final ActivityRecord top = tr.getTopActivity();
+            if (top == null || !top.okToShowLocked()) {
+                if (top != null) {
+                    mStackSupervisor.mRecentTasks.add(top.getTask());
+                }
+                ActivityOptions.abort(options);
+                return;
             }
-            ActivityOptions.abort(options);
-            return;
-        }
 
-        // Set focus to the top running activity of this stack.
-        final ActivityRecord r = topRunningActivityLocked();
-        mStackSupervisor.moveFocusableActivityStackToFrontLocked(r, reason);
+            // Set focus to the top running activity of this stack.
+            final ActivityRecord r = topRunningActivityLocked();
+            mStackSupervisor.moveFocusableActivityStackToFrontLocked(r, reason);
 
-        if (DEBUG_TRANSITION) Slog.v(TAG_TRANSITION, "Prepare to front transition: task=" + tr);
-        if (noAnimation) {
-            mWindowManager.prepareAppTransition(TRANSIT_NONE, false);
-            if (r != null) {
-                mStackSupervisor.mNoAnimActivities.add(r);
+            if (DEBUG_TRANSITION) Slog.v(TAG_TRANSITION, "Prepare to front transition: task=" + tr);
+            if (noAnimation) {
+                mWindowManager.prepareAppTransition(TRANSIT_NONE, false);
+                if (r != null) {
+                    mStackSupervisor.mNoAnimActivities.add(r);
+                }
+                ActivityOptions.abort(options);
+            } else {
+                updateTransitLocked(TRANSIT_TASK_TO_FRONT, options);
             }
-            ActivityOptions.abort(options);
-        } else {
-            updateTransitLocked(TRANSIT_TASK_TO_FRONT, options);
-        }
-        // If a new task is moved to the front, then mark the existing top activity as supporting
-        // picture-in-picture while paused only if the task would not be considered an oerlay on top
-        // of the current activity (eg. not fullscreen, or the assistant)
-        if (canEnterPipOnTaskSwitch(topActivity, tr, null /* toFrontActivity */,
-                options)) {
-            topActivity.supportsEnterPipOnTaskSwitch = true;
-        }
+            // If a new task is moved to the front, then mark the existing top activity as
+            // supporting
 
-        mStackSupervisor.resumeFocusedStackTopActivityLocked();
-        EventLog.writeEvent(EventLogTags.AM_TASK_TO_FRONT, tr.userId, tr.taskId);
+            // picture-in-picture while paused only if the task would not be considered an oerlay
+            // on top
+            // of the current activity (eg. not fullscreen, or the assistant)
+            if (canEnterPipOnTaskSwitch(topActivity, tr, null /* toFrontActivity */,
+                    options)) {
+                topActivity.supportsEnterPipOnTaskSwitch = true;
+            }
 
-        mService.mTaskChangeNotificationController.notifyTaskMovedToFront(tr.taskId);
+            mStackSupervisor.resumeFocusedStackTopActivityLocked();
+            EventLog.writeEvent(EventLogTags.AM_TASK_TO_FRONT, tr.userId, tr.taskId);
+
+            mService.mTaskChangeNotificationController.notifyTaskMovedToFront(tr.taskId);
+        } finally {
+            getDisplay().continueUpdateImeTarget();
+        }
     }
 
     /**
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index c0dc750..f3423c6 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -390,6 +390,11 @@
      */
     int mLayoutSeq = 0;
 
+    /**
+     * Specifies the count to determine whether to defer updating the IME target until ready.
+     */
+    private int mDeferUpdateImeTargetCount;
+
     /** Temporary float array to retrieve 3x3 matrix values. */
     private final float[] mTmpFloats = new float[9];
 
@@ -2454,6 +2459,12 @@
             return null;
         }
 
+        final WindowState curTarget = mService.mInputMethodTarget;
+        if (!canUpdateImeTarget()) {
+            if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Defer updating IME target");
+            return curTarget;
+        }
+
         // TODO(multidisplay): Needs some serious rethought when the target and IME are not on the
         // same display. Or even when the current IME/target are not on the same screen as the next
         // IME/target. For now only look for input windows on the main screen.
@@ -2477,16 +2488,13 @@
         if (DEBUG_INPUT_METHOD && updateImeTarget) Slog.v(TAG_WM,
                 "Proposed new IME target: " + target);
 
-        // Now, a special case -- if the last target's window is in the process of exiting, and is
-        // above the new target, keep on the last target to avoid flicker. Consider for example a
-        // Dialog with the IME shown: when the Dialog is dismissed, we want to keep the IME above it
-        // until it is completely gone so it doesn't drop behind the dialog or its full-screen
-        // scrim.
-        final WindowState curTarget = mService.mInputMethodTarget;
+        // Now, a special case -- if the last target's window is in the process of exiting, and the
+        // new target is home, keep on the last target to avoid flicker. Home is a special case
+        // since its above other stacks in the ordering list, but layed out below the others.
         if (curTarget != null && curTarget.isDisplayedLw() && curTarget.isClosing()
-                && (target == null
-                    || curTarget.mWinAnimator.mAnimLayer > target.mWinAnimator.mAnimLayer)) {
-            if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "Current target higher, not changing");
+                && (target == null || target.isActivityTypeHome())) {
+            if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "New target is home while current target is"
+                    + "closing, not changing");
             return curTarget;
         }
 
@@ -3958,4 +3966,33 @@
     void assignStackOrdering() {
         mTaskStackContainers.assignStackOrdering(getPendingTransaction());
     }
+
+    /**
+     * Increment the deferral count to determine whether to update the IME target.
+     */
+    void deferUpdateImeTarget() {
+        mDeferUpdateImeTargetCount++;
+    }
+
+    /**
+     * Decrement the deferral count to determine whether to update the IME target. If the count
+     * reaches 0, a new ime target will get computed.
+     */
+    void continueUpdateImeTarget() {
+        if (mDeferUpdateImeTargetCount == 0) {
+            return;
+        }
+
+        mDeferUpdateImeTargetCount--;
+        if (mDeferUpdateImeTargetCount == 0) {
+            computeImeTarget(true /* updateImeTarget */);
+        }
+    }
+
+    /**
+     * @return Whether a new IME target should be computed.
+     */
+    private boolean canUpdateImeTarget() {
+        return mDeferUpdateImeTargetCount == 0;
+    }
 }
diff --git a/services/core/java/com/android/server/wm/DisplayWindowController.java b/services/core/java/com/android/server/wm/DisplayWindowController.java
index ba8ec69..a1639c2 100644
--- a/services/core/java/com/android/server/wm/DisplayWindowController.java
+++ b/services/core/java/com/android/server/wm/DisplayWindowController.java
@@ -94,6 +94,31 @@
         }
     }
 
+    /**
+     * Starts deferring the ability to update the IME target. This is needed when a call will
+     * attempt to update the IME target before all information about the Windows have been updated.
+     */
+    public void deferUpdateImeTarget() {
+        synchronized (mWindowMap) {
+            final DisplayContent dc = mRoot.getDisplayContent(mDisplayId);
+            if (dc != null) {
+                dc.deferUpdateImeTarget();
+            }
+        }
+    }
+
+    /**
+     * Resumes updating the IME target after deferring. See {@link #deferUpdateImeTarget()}
+     */
+    public void continueUpdateImeTarget() {
+        synchronized (mWindowMap) {
+            final DisplayContent dc = mRoot.getDisplayContent(mDisplayId);
+            if (dc != null) {
+                dc.continueUpdateImeTarget();
+            }
+        }
+    }
+
     @Override
     public String toString() {
         return "{DisplayWindowController displayId=" + mDisplayId + "}";
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index c710c97..91c449b 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2726,7 +2726,8 @@
     }
 
     boolean isClosing() {
-        return mAnimatingExit || (mService.mClosingApps.contains(mAppToken));
+        return mAnimatingExit || (mAppToken != null && mAppToken.isAnimating()
+                && mAppToken.hiddenRequested);
     }
 
     void addWinAnimatorToList(ArrayList<WindowStateAnimator> animators) {