Merge "Fix IME flicker: move hiding the surface into the control target" into rvc-dev am: 895a2e626a

Original change: https://googleplex-android-review.googlesource.com/c/platform/frameworks/base/+/12070682

Change-Id: Ib984a4bc34f0527c809de35919340019a04b017a

17 files changed, 211 insertions, 46 deletions

diff --git a/core/java/android/inputmethodservice/InputMethodService.java b/core/java/android/inputmethodservice/InputMethodService.java
index 5647bf90d2fb..c5a11abe1136 100644
--- a/core/java/android/inputmethodservice/InputMethodService.java
+++ b/core/java/android/inputmethodservice/InputMethodService.java
@@ -605,9 +605,6 @@ public class InputMethodService extends AbstractInputMethodService {
             if (DEBUG) Log.v(TAG, "unbindInput(): binding=" + mInputBinding
                     + " ic=" + mInputConnection);
             // Unbind input is per process per display.
-            // TODO(b/150902448): free-up IME surface when target is changing.
-            //  e.g. DisplayContent#setInputMethodTarget()
-            removeImeSurface();
             onUnbindInput();
             mInputBinding = null;
             mInputConnection = null;
diff --git a/core/java/android/view/ImeInsetsSourceConsumer.java b/core/java/android/view/ImeInsetsSourceConsumer.java
index ef9d990168d2..c1998c6009cf 100644
--- a/core/java/android/view/ImeInsetsSourceConsumer.java
+++ b/core/java/android/view/ImeInsetsSourceConsumer.java
@@ -119,11 +119,11 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
         // If we had a request before to show from IME (tracked with mImeRequestedShow), reaching
         // this code here means that we now got control, so we can start the animation immediately.
         // If client window is trying to control IME and IME is already visible, it is immediate.
-        if (fromIme || mState.getSource(getType()).isVisible()) {
+        if (fromIme || mState.getSource(getType()).isVisible() && getControl() != null) {
             return ShowResult.SHOW_IMMEDIATELY;
         }
 
-        return getImm().requestImeShow(null /* resultReceiver */)
+        return getImm().requestImeShow(mController.getHost().getWindowToken())
                 ? ShowResult.IME_SHOW_DELAYED : ShowResult.IME_SHOW_FAILED;
     }
 
@@ -132,12 +132,15 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
      */
     @Override
     void notifyHidden() {
-        getImm().notifyImeHidden();
+        getImm().notifyImeHidden(mController.getHost().getWindowToken());
     }
 
     @Override
     public void removeSurface() {
-        getImm().removeImeSurface();
+        final IBinder window = mController.getHost().getWindowToken();
+        if (window != null) {
+            getImm().removeImeSurface(window);
+        }
     }
 
     @Override
@@ -146,6 +149,7 @@ public final class ImeInsetsSourceConsumer extends InsetsSourceConsumer {
         super.setControl(control, showTypes, hideTypes);
         if (control == null && !mIsRequestedVisibleAwaitingControl) {
             hide();
+            removeSurface();
         }
     }
 
diff --git a/core/java/android/view/InsetsSourceConsumer.java b/core/java/android/view/InsetsSourceConsumer.java
index 40e6f4b2fce8..700dc66fab55 100644
--- a/core/java/android/view/InsetsSourceConsumer.java
+++ b/core/java/android/view/InsetsSourceConsumer.java
@@ -153,6 +153,9 @@ public class InsetsSourceConsumer {
                 if (oldLeash == null || newLeash == null || !oldLeash.isSameSurface(newLeash)) {
                     applyHiddenToControl();
                 }
+                if (!requestedVisible && !mIsAnimationPending) {
+                    removeSurface();
+                }
             }
         }
         if (lastControl != null) {
diff --git a/core/java/android/view/ViewRootImpl.java b/core/java/android/view/ViewRootImpl.java
index f028a28f1d2b..2b7044d0c67e 100644
--- a/core/java/android/view/ViewRootImpl.java
+++ b/core/java/android/view/ViewRootImpl.java
@@ -9380,6 +9380,11 @@ public final class ViewRootImpl implements ViewParent,
         return mInputEventReceiver.getToken();
     }
 
+    @NonNull
+    public IBinder getWindowToken() {
+        return mAttachInfo.mWindowToken;
+    }
+
     /**
      * Class for managing the accessibility interaction connection
      * based on the global accessibility state.
diff --git a/core/java/android/view/inputmethod/InputMethodManager.java b/core/java/android/view/inputmethod/InputMethodManager.java
index aedb59bfee42..3be0a4d39a56 100644
--- a/core/java/android/view/inputmethod/InputMethodManager.java
+++ b/core/java/android/view/inputmethod/InputMethodManager.java
@@ -2109,28 +2109,36 @@ public final class InputMethodManager {
 
     /**
      * Call showSoftInput with currently focused view.
-     * @return {@code true} if IME can be shown.
+     *
+     * @param windowToken the window from which this request originates. If this doesn't match the
+     *                    currently served view, the request is ignored and returns {@code false}.
+     *
+     * @return {@code true} if IME can (eventually) be shown, {@code false} otherwise.
      * @hide
      */
-    public boolean requestImeShow(ResultReceiver resultReceiver) {
+    public boolean requestImeShow(IBinder windowToken) {
         synchronized (mH) {
             final View servedView = getServedViewLocked();
-            if (servedView == null) {
+            if (servedView == null || servedView.getWindowToken() != windowToken) {
                 return false;
             }
-            showSoftInput(servedView, 0 /* flags */, resultReceiver);
+            showSoftInput(servedView, 0 /* flags */, null /* resultReceiver */);
             return true;
         }
     }
 
     /**
      * Notify IME directly that it is no longer visible.
+     *
+     * @param windowToken the window from which this request originates. If this doesn't match the
+     *                    currently served view, the request is ignored.
      * @hide
      */
-    public void notifyImeHidden() {
+    public void notifyImeHidden(IBinder windowToken) {
         synchronized (mH) {
             try {
-                if (mCurMethod != null) {
+                if (mCurMethod != null && mCurRootView != null
+                        && mCurRootView.getWindowToken() == windowToken) {
                     mCurMethod.notifyImeHidden();
                 }
             } catch (RemoteException re) {
@@ -2140,15 +2148,15 @@ public final class InputMethodManager {
 
     /**
      * Notify IME directly to remove surface as it is no longer visible.
+     * @param windowToken The client window token that requests the IME to remove its surface.
      * @hide
      */
-    public void removeImeSurface() {
+    public void removeImeSurface(IBinder windowToken) {
         synchronized (mH) {
             try {
-                if (mCurMethod != null) {
-                    mCurMethod.removeImeSurface();
-                }
-            } catch (RemoteException re) {
+                mService.removeImeSurfaceFromWindow(windowToken);
+            } catch (RemoteException e) {
+                throw e.rethrowFromSystemServer();
             }
         }
     }
diff --git a/core/java/com/android/internal/view/IInputMethodManager.aidl b/core/java/com/android/internal/view/IInputMethodManager.aidl
index 8ec51b89d240..a1cbd3fcae79 100644
--- a/core/java/com/android/internal/view/IInputMethodManager.aidl
+++ b/core/java/com/android/internal/view/IInputMethodManager.aidl
@@ -73,5 +73,8 @@ interface IInputMethodManager {
             in float[] matrixValues);
 
     oneway void reportPerceptible(in IBinder windowToken, boolean perceptible);
+    /** Remove the IME surface. Requires INTERNAL_SYSTEM_WINDOW permission. */
     void removeImeSurface();
+    /** Remove the IME surface. Requires passing the currently focused window. */
+    void removeImeSurfaceFromWindow(in IBinder windowToken);
 }
diff --git a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
index 1b3272572db0..7efd616c5607 100644
--- a/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
+++ b/core/tests/coretests/src/android/view/InsetsSourceConsumerTest.java
@@ -71,6 +71,9 @@ public class InsetsSourceConsumerTest {
     private SurfaceControl mLeash;
     @Mock Transaction mMockTransaction;
     private InsetsSource mSpyInsetsSource;
+    private boolean mRemoveSurfaceCalled = false;
+    private InsetsController mController;
+    private InsetsState mState;
 
     @Before
     public void setup() {
@@ -89,13 +92,19 @@ public class InsetsSourceConsumerTest {
             } catch (BadTokenException e) {
                 // activity isn't running, lets ignore BadTokenException.
             }
-            InsetsState state = new InsetsState();
+            mState = new InsetsState();
             mSpyInsetsSource = Mockito.spy(new InsetsSource(ITYPE_STATUS_BAR));
-            state.addSource(mSpyInsetsSource);
-
-            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, state,
-                    () -> mMockTransaction,
-                    new InsetsController(new ViewRootInsetsControllerHost(viewRootImpl)));
+            mState.addSource(mSpyInsetsSource);
+
+            mController = new InsetsController(new ViewRootInsetsControllerHost(viewRootImpl));
+            mConsumer = new InsetsSourceConsumer(ITYPE_STATUS_BAR, mState,
+                    () -> mMockTransaction, mController) {
+                @Override
+                public void removeSurface() {
+                    super.removeSurface();
+                    mRemoveSurfaceCalled = true;
+                }
+            };
         });
         instrumentation.waitForIdleSync();
 
@@ -171,6 +180,25 @@ public class InsetsSourceConsumerTest {
             mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
                     new int[1], hideTypes);
             assertEquals(statusBars(), hideTypes[0]);
+            assertFalse(mRemoveSurfaceCalled);
+        });
+    }
+
+    @Test
+    public void testRestore_noAnimation() {
+        InstrumentationRegistry.getInstrumentation().runOnMainSync(() -> {
+            mConsumer.hide();
+            mController.onStateChanged(mState);
+            mConsumer.setControl(null, new int[1], new int[1]);
+            reset(mMockTransaction);
+            verifyZeroInteractions(mMockTransaction);
+            mRemoveSurfaceCalled = false;
+            int[] hideTypes = new int[1];
+            mConsumer.setControl(new InsetsSourceControl(ITYPE_STATUS_BAR, mLeash, new Point()),
+                    new int[1], hideTypes);
+            assertTrue(mRemoveSurfaceCalled);
+            assertEquals(0, hideTypes[0]);
         });
+
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java b/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
index 8ba5b9951c54..c0089e53f8b6 100644
--- a/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
+++ b/packages/SystemUI/src/com/android/systemui/wm/DisplayImeController.java
@@ -226,6 +226,8 @@ public class DisplayImeController implements DisplayController.OnDisplaysChanged
                             if (!activeControl.getSurfacePosition().equals(lastSurfacePosition)
                                     && mAnimation != null) {
                                 startAnimation(mImeShowing, true /* forceRestart */);
+                            } else if (!mImeShowing) {
+                                removeImeSurface();
                             }
                         });
                     }
@@ -370,16 +372,7 @@ public class DisplayImeController implements DisplayController.OnDisplaysChanged
                     dispatchEndPositioning(mDisplayId, mCancelled, t);
                     if (mAnimationDirection == DIRECTION_HIDE && !mCancelled) {
                         t.hide(mImeSourceControl.getLeash());
-                        final IInputMethodManager imms = getImms();
-                        if (imms != null) {
-                            try {
-                                // Remove the IME surface to make the insets invisible for
-                                // non-client controlled insets.
-                                imms.removeImeSurface();
-                            } catch (RemoteException e) {
-                                Slog.e(TAG, "Failed to remove IME surface.", e);
-                            }
-                        }
+                        removeImeSurface();
                     }
                     t.apply();
                     mTransactionPool.release(t);
@@ -402,6 +395,19 @@ public class DisplayImeController implements DisplayController.OnDisplaysChanged
         }
     }
 
+    void removeImeSurface() {
+        final IInputMethodManager imms = getImms();
+        if (imms != null) {
+            try {
+                // Remove the IME surface to make the insets invisible for
+                // non-client controlled insets.
+                imms.removeImeSurface();
+            } catch (RemoteException e) {
+                Slog.e(TAG, "Failed to remove IME surface.", e);
+            }
+        }
+    }
+
     /**
      * Allows other things to synchronize with the ime position
      */
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
index 70f0399d1070..05cf40a091b6 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerInternal.java
@@ -120,6 +120,11 @@ public abstract class InputMethodManagerInternal {
     public abstract void reportImeControl(@Nullable IBinder windowToken);
 
     /**
+     * Destroys the IME surface.
+     */
+    public abstract void removeImeSurface();
+
+    /**
      * Fake implementation of {@link InputMethodManagerInternal}.  All the methods do nothing.
      */
     private static final InputMethodManagerInternal NOP =
@@ -166,6 +171,10 @@ public abstract class InputMethodManagerInternal {
                 @Override
                 public void reportImeControl(@Nullable IBinder windowToken) {
                 }
+
+                @Override
+                public void removeImeSurface() {
+                }
             };
 
     /**
diff --git a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
index c027ebcfd568..d8ee32e7bd74 100644
--- a/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/InputMethodManagerService.java
@@ -211,6 +211,7 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
     static final int MSG_INITIALIZE_IME = 1040;
     static final int MSG_CREATE_SESSION = 1050;
     static final int MSG_REMOVE_IME_SURFACE = 1060;
+    static final int MSG_REMOVE_IME_SURFACE_FROM_WINDOW = 1061;
 
     static final int MSG_START_INPUT = 2000;
 
@@ -4005,6 +4006,13 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         mHandler.sendMessage(mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE));
     }
 
+    @Override
+    public void removeImeSurfaceFromWindow(IBinder windowToken) {
+        // No permission check, because we'll only execute the request if the calling window is
+        // also the current IME client.
+        mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE_FROM_WINDOW, windowToken).sendToTarget();
+    }
+
     @BinderThread
     private void notifyUserAction(@NonNull IBinder token) {
         if (DEBUG) {
@@ -4278,11 +4286,27 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
                 return true;
             }
             case MSG_REMOVE_IME_SURFACE: {
-                try {
-                    if (mEnabledSession != null && mEnabledSession.session != null) {
-                        mEnabledSession.session.removeImeSurface();
+                synchronized (mMethodMap) {
+                    try {
+                        if (mEnabledSession != null && mEnabledSession.session != null
+                                && !mShowRequested) {
+                            mEnabledSession.session.removeImeSurface();
+                        }
+                    } catch (RemoteException e) {
+                    }
+                }
+                return true;
+            }
+            case MSG_REMOVE_IME_SURFACE_FROM_WINDOW: {
+                IBinder windowToken = (IBinder) msg.obj;
+                synchronized (mMethodMap) {
+                    try {
+                        if (windowToken == mCurFocusedWindow
+                                && mEnabledSession != null && mEnabledSession.session != null) {
+                            mEnabledSession.session.removeImeSurface();
+                        }
+                    } catch (RemoteException e) {
                     }
-                } catch (RemoteException e) {
                 }
                 return true;
             }
@@ -5116,6 +5140,11 @@ public class InputMethodManagerService extends IInputMethodManager.Stub
         public void reportImeControl(@Nullable IBinder windowToken) {
             mService.reportImeControl(windowToken);
         }
+
+        @Override
+        public void removeImeSurface() {
+            mService.mHandler.sendMessage(mService.mHandler.obtainMessage(MSG_REMOVE_IME_SURFACE));
+        }
     }
 
     @BinderThread
diff --git a/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java b/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
index 2e3d3963c9df..19dff9807075 100644
--- a/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
+++ b/services/core/java/com/android/server/inputmethod/MultiClientInputMethodManagerService.java
@@ -225,6 +225,11 @@ public final class MultiClientInputMethodManagerService {
                         @Override
                         public void reportImeControl(@Nullable IBinder windowToken) {
                         }
+
+                        @Override
+                        public void removeImeSurface() {
+                            reportNotSupported();
+                        }
                     });
         }
 
@@ -1473,6 +1478,12 @@ public final class MultiClientInputMethodManagerService {
 
         @BinderThread
         @Override
+        public void removeImeSurfaceFromWindow(IBinder windowToken) {
+            reportNotSupported();
+        }
+
+        @BinderThread
+        @Override
         public boolean showSoftInput(
                 IInputMethodClient client, IBinder token, int flags,
                 ResultReceiver resultReceiver) {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index e8a4234f44a7..b94fb0471af4 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -3574,12 +3574,11 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         }
     }
 
-    private void updateImeControlTarget() {
+    void updateImeControlTarget() {
         mInputMethodControlTarget = computeImeControlTarget();
         mInsetsStateController.onImeControlTargetChanged(mInputMethodControlTarget);
 
-        final WindowState win = mInputMethodControlTarget != null
-                ? mInputMethodControlTarget.getWindow() : null;
+        final WindowState win = InsetsControlTarget.asWindowOrNull(mInputMethodControlTarget);
         final IBinder token = win != null ? win.mClient.asBinder() : null;
         // Note: not allowed to call into IMMS with the WM lock held, hence the post.
         mWmService.mH.post(() ->
@@ -3603,6 +3602,17 @@ class DisplayContent extends WindowContainer<DisplayContent.DisplayChildWindowCo
         if (!isImeControlledByApp() && mRemoteInsetsControlTarget != null) {
             return mRemoteInsetsControlTarget;
         } else {
+            // Now, a special case -- if the last target's window is in the process of exiting, but
+            // not removed, keep on the last target to avoid IME flicker.
+            final WindowState cur = InsetsControlTarget.asWindowOrNull(mInputMethodControlTarget);
+            if (cur != null && !cur.mRemoved && cur.isDisplayedLw() && cur.isClosing()
+                    && !cur.isActivityTypeHome()) {
+                if (DEBUG_INPUT_METHOD) {
+                    Slog.v(TAG_WM, "Not changing control while current window"
+                            + " is closing and not removed");
+                }
+                return cur;
+            }
             // Otherwise, we just use the ime target as received from IME.
             return mInputMethodInputTarget;
         }
diff --git a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
index 8298763c1392..99ee5e121b7a 100644
--- a/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
+++ b/services/core/java/com/android/server/wm/ImeInsetsSourceProvider.java
@@ -18,13 +18,14 @@ package com.android.server.wm;
 
 import static com.android.server.wm.ProtoLogGroup.WM_DEBUG_IME;
 
-import android.graphics.PixelFormat;
 import android.view.InsetsSource;
 import android.view.WindowInsets;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.protolog.common.ProtoLog;
 
+import java.io.PrintWriter;
+
 /**
  * Controller for IME inset source on the server. It's called provider as it provides the
  * {@link InsetsSource} to the client that uses it in {@link InsetsSourceConsumer}.
@@ -132,8 +133,17 @@ class ImeInsetsSourceProvider extends InsetsSourceProvider {
                 || (mImeTargetFromIme != null && dcTarget.getParentWindow() == mImeTargetFromIme
                         && dcTarget.mSubLayer > mImeTargetFromIme.mSubLayer)
                 || mImeTargetFromIme == mDisplayContent.getImeFallback()
-                // If IME target is transparent but control target matches requesting window.
-                || (controlTarget == mImeTargetFromIme
-                        && PixelFormat.formatHasAlpha(dcTarget.mAttrs.format));
+                || (!mImeTargetFromIme.isClosing() && controlTarget == mImeTargetFromIme);
+    }
+
+    @Override
+    public void dump(PrintWriter pw, String prefix) {
+        super.dump(pw, prefix);
+        if (mImeTargetFromIme != null) {
+            pw.print(prefix);
+            pw.print("showImePostLayout pending for mImeTargetFromIme=");
+            pw.print(mImeTargetFromIme);
+            pw.println();
+        }
     }
 }
diff --git a/services/core/java/com/android/server/wm/InsetsControlTarget.java b/services/core/java/com/android/server/wm/InsetsControlTarget.java
index c50f296504fc..3ffc26a7a8ad 100644
--- a/services/core/java/com/android/server/wm/InsetsControlTarget.java
+++ b/services/core/java/com/android/server/wm/InsetsControlTarget.java
@@ -62,4 +62,8 @@ interface InsetsControlTarget {
         return false;
     }
 
+    /** Returns {@code target.getWindow()}, or null if {@code target} is {@code null}. */
+    static WindowState asWindowOrNull(InsetsControlTarget target) {
+        return target != null ? target.getWindow() : null;
+    }
 }
diff --git a/services/core/java/com/android/server/wm/InsetsStateController.java b/services/core/java/com/android/server/wm/InsetsStateController.java
index 63083faaddb1..77bd4a47a884 100644
--- a/services/core/java/com/android/server/wm/InsetsStateController.java
+++ b/services/core/java/com/android/server/wm/InsetsStateController.java
@@ -44,6 +44,7 @@ import android.view.InsetsState;
 import android.view.InsetsState.InternalInsetsType;
 import android.view.WindowManager;
 
+import com.android.server.inputmethod.InputMethodManagerInternal;
 import com.android.server.protolog.common.ProtoLog;
 
 import java.io.PrintWriter;
@@ -74,7 +75,21 @@ class InsetsStateController {
             w.notifyInsetsChanged();
         }
     };
-    private final InsetsControlTarget mEmptyImeControlTarget = new InsetsControlTarget() { };
+    private final InsetsControlTarget mEmptyImeControlTarget = new InsetsControlTarget() {
+        @Override
+        public void notifyInsetsControlChanged() {
+            InsetsSourceControl[] controls = getControlsForDispatch(this);
+            if (controls == null) {
+                return;
+            }
+            for (InsetsSourceControl control : controls) {
+                if (control.getType() == ITYPE_IME) {
+                    mDisplayContent.mWmService.mH.post(() ->
+                            InputMethodManagerInternal.get().removeImeSurface());
+                }
+            }
+        }
+    };
 
     InsetsStateController(DisplayContent displayContent) {
         mDisplayContent = displayContent;
diff --git a/services/core/java/com/android/server/wm/WindowState.java b/services/core/java/com/android/server/wm/WindowState.java
index 837fafec1634..00be75fddf04 100644
--- a/services/core/java/com/android/server/wm/WindowState.java
+++ b/services/core/java/com/android/server/wm/WindowState.java
@@ -2170,6 +2170,9 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
         if (isInputMethodTarget()) {
             dc.computeImeTarget(true /* updateImeTarget */);
         }
+        if (dc.mInputMethodControlTarget == this) {
+            dc.updateImeControlTarget();
+        }
 
         final int type = mAttrs.type;
         if (WindowManagerService.excludeWindowTypeFromTapOutTask(type)) {
diff --git a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
index 77e3c597311c..d64fdb81107c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java
@@ -898,6 +898,26 @@ public class DisplayContentTests extends WindowTestsBase {
     }
 
     @Test
+    public void testComputeImeControlTarget_exitingApp() throws Exception {
+        final DisplayContent dc = createNewDisplay();
+
+        WindowState exitingWin = createWindow(null, TYPE_BASE_APPLICATION, "exiting app");
+        makeWindowVisible(exitingWin);
+        exitingWin.mWinAnimator.mDrawState = WindowStateAnimator.HAS_DRAWN;
+        exitingWin.mAnimatingExit = true;
+
+        dc.mInputMethodControlTarget = exitingWin;
+        dc.mInputMethodTarget = dc.mInputMethodInputTarget =
+                createWindow(null, TYPE_BASE_APPLICATION, "starting app");
+
+        assertEquals(exitingWin, dc.computeImeControlTarget());
+
+        exitingWin.removeImmediately();
+
+        assertEquals(dc.mInputMethodInputTarget, dc.computeImeControlTarget());
+    }
+
+    @Test
     public void testComputeImeControlTarget_splitscreen() throws Exception {
         final DisplayContent dc = createNewDisplay();
         dc.mInputMethodInputTarget = createWindow(null, TYPE_BASE_APPLICATION, "app");