Fix IME snapshot missing to remove by defered IME target computation

When swiping the app task to home, in Task#moveTaskToFront will call
deferUpdateImeTarget to defer updating IME target until the recents
animation finished, this makes the IME snapshot cannot be removed when
finished app transition or recents animation.

Moreover, if bring that app task to the top next and then a
configuration change makes the app's activity relaunch, IME snapshot
could not be removed by DC#removeImeSnapshotIfPossible because the
window's parent may still in app transitioning.

To fix this unexpected IME snapshot visibility, consolidate logc with:
1) Modify removeImeSnapshotIfPossible with
   removeImeSnapshotByTarget(windowContainer) that removing IME snapshot
   when the target window is the snapshot caller without checking
   isAnimating, in case IME snapsthot missed to remove when the
   activity main window is relaunching.
2) Calling removeImeSurfaceByTarget when
   WindowState#onSurfaceShownChanged(false), in case deferring IME
   target computation will not remove IME snapsot by missing the
   animation finished callback.

Fix: 235376671
Test: atest DisplayContentTests#\
       testRemoveImeScreenshot_whenTargetSurfaceWasInvisible
Test: atest DisplayContentTests#\
       testRemoveImeScreenshot_whenWindowRemoveImmediately
Change-Id: I937b549afaaa5bd22aa232adb1cbb2cbb1cd8e8e

services/core/java/com/android/server/wm/DisplayContent.java

@@ -4249,6 +4249,17 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         void detach(Transaction t) {
             removeImeSurface(t);
         }
+
+        @Override
+        public String toString() {
+            StringBuilder sb = new StringBuilder(64);
+            sb.append("ImeScreenshot{");
+            sb.append(Integer.toHexString(System.identityHashCode(this)));
+            sb.append(" imeTarget=" + mImeTarget);
+            sb.append(" surface=" + mImeSurface);
+            sb.append('}');
+            return sb.toString();
+        }
     }
 
     private void attachAndShowImeScreenshotOnTarget() {
@@ -4281,15 +4292,23 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     }
 
     /**
-     * Removes the IME screenshot when necessary.
-     *
-     * Used when app transition animation finished or obsoleted screenshot surface like size
-     * changed by rotation.
+     * Removes the IME screenshot when the caller is a part of the attached target window.
      */
-    void removeImeScreenshotIfPossible() {
-        if (mImeLayeringTarget == null
-                || mImeLayeringTarget.mAttrs.type != TYPE_APPLICATION_STARTING
-                && !mImeLayeringTarget.inTransitionSelfOrParent()) {
+    void removeImeSurfaceByTarget(WindowContainer win) {
+        if (mImeScreenshot == null || win == null) {
+            return;
+        }
+        // The starting window shouldn't be the input target to attach the IME screenshot during
+        // transitioning.
+        if (win.asWindowState() != null
+                && win.asWindowState().mAttrs.type == TYPE_APPLICATION_STARTING) {
+            return;
+        }
+
+        final WindowState screenshotTarget = mImeScreenshot.getImeTarget();
+        final boolean winIsOrContainsScreenshotTarget = (win == screenshotTarget
+                || win.getWindow(w -> w == screenshotTarget) != null);
+        if (winIsOrContainsScreenshotTarget) {
             removeImeSurfaceImmediately();
         }
     }
@@ -4637,10 +4656,8 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                     wc, SurfaceAnimator.animationTypeToString(type), mImeScreenshot,
                     mImeScreenshot.getImeTarget());
         }
-        if (mImeScreenshot != null && (wc == mImeScreenshot.getImeTarget()
-                || wc.getWindow(w -> w == mImeScreenshot.getImeTarget()) != null)
-                && (type & WindowState.EXIT_ANIMATING_TYPES) != 0) {
-            removeImeSurfaceImmediately();
+        if ((type & WindowState.EXIT_ANIMATING_TYPES) != 0) {
+            removeImeSurfaceByTarget(wc);
         }
     }
 

services/core/java/com/android/server/wm/WindowState.java

@@ -2449,8 +2449,8 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
 
         final DisplayContent dc = getDisplayContent();
         if (isImeLayeringTarget()) {
-            // Remove the IME screenshot surface if the layering target is not animating.
-            dc.removeImeScreenshotIfPossible();
+            // Remove the attached IME screenshot surface.
+            dc.removeImeSurfaceByTarget(this);
             // Make sure to set mImeLayeringTarget as null when the removed window is the
             // IME target, in case computeImeTarget may use the outdated target.
             dc.setImeLayeringTarget(null);
@@ -3584,6 +3584,7 @@ class WindowState extends WindowContainer<WindowState> implements WindowManagerP
         } else {
             logExclusionRestrictions(EXCLUSION_LEFT);
             logExclusionRestrictions(EXCLUSION_RIGHT);
+            getDisplayContent().removeImeSurfaceByTarget(this);
         }
         // Exclude toast because legacy apps may show toast window by themselves, so the misused
         // apps won't always be considered as foreground state.

services/tests/wmtests/src/com/android/server/wm/DisplayContentTests.java

@@ -2183,6 +2183,52 @@ public class DisplayContentTests extends WindowTestsBase {
         assertNotEquals(curSnapshot, mDisplayContent.mImeScreenshot);
     }
 
+    @UseTestDisplay(addWindows = {W_INPUT_METHOD})
+    @Test
+    public void testRemoveImeScreenshot_whenTargetSurfaceWasInvisible() {
+        final Task rootTask = createTask(mDisplayContent);
+        final Task task = createTaskInRootTask(rootTask, 0 /* userId */);
+        final ActivityRecord activity = createActivityRecord(mDisplayContent, task);
+        final WindowState win = createWindow(null, TYPE_BASE_APPLICATION, activity, "win");
+        win.onSurfaceShownChanged(true);
+        makeWindowVisible(win, mDisplayContent.mInputMethodWindow);
+        task.getDisplayContent().prepareAppTransition(TRANSIT_CLOSE);
+        doReturn(true).when(task).okToAnimate();
+        ArrayList<WindowContainer> sources = new ArrayList<>();
+        sources.add(activity);
+
+        mDisplayContent.setImeLayeringTarget(win);
+        mDisplayContent.setImeInputTarget(win);
+        mDisplayContent.getInsetsStateController().getImeSourceProvider().setImeShowing(true);
+        task.applyAnimation(null, TRANSIT_OLD_TASK_CLOSE, false /* enter */,
+                false /* isVoiceInteraction */, sources);
+        assertNotNull(mDisplayContent.mImeScreenshot);
+
+        win.onSurfaceShownChanged(false);
+        assertNull(mDisplayContent.mImeScreenshot);
+    }
+
+    @UseTestDisplay(addWindows = {W_INPUT_METHOD})
+    @Test
+    public void testRemoveImeScreenshot_whenWindowRemoveImmediately() {
+        final Task rootTask = createTask(mDisplayContent);
+        final Task task = createTaskInRootTask(rootTask, 0 /* userId */);
+        final ActivityRecord activity = createActivityRecord(mDisplayContent, task);
+        final WindowState win = createWindow(null, TYPE_BASE_APPLICATION, activity, "win");
+        makeWindowVisible(mDisplayContent.mInputMethodWindow);
+
+        mDisplayContent.setImeLayeringTarget(win);
+        mDisplayContent.setImeInputTarget(win);
+        mDisplayContent.getInsetsStateController().getImeSourceProvider().setImeShowing(true);
+        mDisplayContent.showImeScreenshot();
+        assertNotNull(mDisplayContent.mImeScreenshot);
+
+        // Expect IME snapshot will be removed when the win is IME layering target and invoked
+        // removeImeSurfaceByTarget.
+        win.removeImmediately();
+        assertNull(mDisplayContent.mImeScreenshot);
+    }
+
     @Test
     public void testRotateBounds_keepSamePhysicalPosition() {
         final DisplayContent dc =