Merge changes Ia6722e1c,I6bef36c7

* changes:
  Better IME transition while switching app with recents (5/N)
  Better IME transition while switching app with recents (4/N)

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

@@ -113,6 +113,8 @@ import static com.android.server.wm.DisplayContentProto.OPENING_APPS;
 import static com.android.server.wm.DisplayContentProto.RESUMED_ACTIVITY;
 import static com.android.server.wm.DisplayContentProto.ROOT_DISPLAY_AREA;
 import static com.android.server.wm.DisplayContentProto.SCREEN_ROTATION_ANIMATION;
+import static com.android.server.wm.SurfaceAnimator.ANIMATION_TYPE_APP_TRANSITION;
+import static com.android.server.wm.SurfaceAnimator.ANIMATION_TYPE_RECENTS;
 import static com.android.server.wm.Task.ActivityState.RESUMED;
 import static com.android.server.wm.WindowContainer.AnimationFlags.PARENTS;
 import static com.android.server.wm.WindowContainer.AnimationFlags.TRANSITION;
@@ -163,6 +165,7 @@ import android.graphics.Rect;
 import android.graphics.RectF;
 import android.graphics.Region;
 import android.graphics.Region.Op;
+import android.hardware.HardwareBuffer;
 import android.hardware.display.DisplayManagerInternal;
 import android.metrics.LogMaker;
 import android.os.Binder;
@@ -606,6 +609,9 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     /** If {@code true} hold off on modifying the animation layer of {@link #mImeLayeringTarget} */
     boolean mImeLayeringTargetWaitingAnim;
 
+    /** The screenshot IME surface to place on the task while transitioning to the next task. */
+    SurfaceControl mImeScreenshot;
+
     private final PointerEventDispatcher mPointerEventDispatcher;
 
     private final InsetsStateController mInsetsStateController;
@@ -3530,7 +3536,7 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
             if (updateImeTarget) {
                 if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Moving IM target from "
                         + mImeLayeringTarget + " to null since mInputMethodWindow is null");
-                setImeLayeringTarget(null, mImeLayeringTargetWaitingAnim);
+                setImeLayeringTargetInner(null);
             }
             return null;
         }
@@ -3547,41 +3553,9 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mUpdateImeTarget = updateImeTarget;
         WindowState target = getWindow(mComputeImeTargetPredicate);
 
-        // Keeps the IME target with the last window while swiping up to recents to prevent
-        // flickering due to IME hide animation on top of recents.
-        // TODO(b/166736352): This logic should go away once we switch over target immediately
-        //  and do the screenshot to preserve IME on disappearing target
-        if (target != null && curTarget != null && target.isActivityTypeHome()
-                && curTarget.getInsetsState().getSource(ITYPE_IME).isVisible()) {
-            return curTarget;
-        }
-
-        // Yet more tricksyness!  If this window is a "starting" window, we do actually want
-        // to be on top of it, but it is not -really- where input will go. So look down below
-        // for a real window to target...
-        if (target != null && target.mAttrs.type == TYPE_APPLICATION_STARTING) {
-            final ActivityRecord activity = target.mActivityRecord;
-            if (activity != null) {
-                final WindowState betterTarget = activity.getImeTargetBelowWindow(target);
-                if (betterTarget != null) {
-                    target = betterTarget;
-                }
-            }
-        }
-
         if (DEBUG_INPUT_METHOD && updateImeTarget) Slog.v(TAG_WM,
                 "Proposed new IME target: " + target + " for display: " + getDisplayId());
 
-        // 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. The exception is if the
-        // current target is home since we want opening apps to become the IME target right away.
-        if (curTarget != null && !curTarget.mRemoved && curTarget.isDisplayed()
-                && curTarget.isClosing() && !curTarget.isActivityTypeHome()) {
-            if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "Not changing target till current window is"
-                    + " closing and not removed");
-            return curTarget;
-        }
-
         if (DEBUG_INPUT_METHOD) Slog.v(TAG_WM, "Desired input method target=" + target
                 + " updateImeTarget=" + updateImeTarget);
 
@@ -3590,42 +3564,16 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                 if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Moving IM target from " + curTarget
                         + " to null." + (SHOW_STACK_CRAWLS ? " Callers="
                         + Debug.getCallers(4) : ""));
-                setImeLayeringTarget(null, mImeLayeringTargetWaitingAnim);
+                setImeLayeringTargetInner(null);
             }
 
             return null;
         }
 
         if (updateImeTarget) {
-            ActivityRecord activity = curTarget == null ? null : curTarget.mActivityRecord;
-            if (activity != null) {
-
-                // Now some fun for dealing with window animations that modify the Z order. We need
-                // to look at all windows below the current target that are in this app, finding the
-                // highest visible one in layering.
-                WindowState highestTarget = null;
-                if (activity.isAnimating(PARENTS | TRANSITION)) {
-                    highestTarget = activity.getHighestAnimLayerWindow(curTarget);
-                }
-
-                if (highestTarget != null) {
-                    if (DEBUG_INPUT_METHOD) {
-                        Slog.v(TAG_WM, mAppTransition + " " + highestTarget + " animating="
-                                + highestTarget.isAnimating(TRANSITION | PARENTS));
-                    }
-
-                    if (mAppTransition.isTransitionSet()) {
-                        // If we are currently setting up for an animation, hold everything until we
-                        // can find out what will happen.
-                        setImeLayeringTarget(highestTarget, true);
-                        return highestTarget;
-                    }
-                }
-            }
-
             if (DEBUG_INPUT_METHOD) Slog.w(TAG_WM, "Moving IM target from " + curTarget + " to "
                     + target + (SHOW_STACK_CRAWLS ? " Callers=" + Debug.getCallers(4) : ""));
-            setImeLayeringTarget(target, false);
+            setImeLayeringTargetInner(target);
         }
 
         return target;
@@ -3727,16 +3675,21 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
      * Sets the window the IME is on top of.
      * @param target window to place the IME surface on top of. If {@code null}, the IME will be
      *               placed at its parent's surface.
-     * @param targetWaitingAnim if {@code true}, hold off on modifying the animation layer of
-     *                          the target.
      */
-    private void setImeLayeringTarget(@Nullable WindowState target, boolean targetWaitingAnim) {
-        if (target == mImeLayeringTarget && mImeLayeringTargetWaitingAnim == targetWaitingAnim) {
+    private void setImeLayeringTargetInner(@Nullable WindowState target) {
+        if (target == mImeLayeringTarget) {
             return;
         }
+        // Prepare the IME screenshot for the last IME target when its task is applying app
+        // transition. This is for the better IME transition to keep IME visibility when
+        // transitioning to the next task.
+        if (mImeLayeringTarget != null && mImeLayeringTarget.isAnimating(PARENTS | TRANSITION,
+                ANIMATION_TYPE_APP_TRANSITION | ANIMATION_TYPE_RECENTS)) {
+            attachAndShowImeScreenshotOnTarget();
+        }
+
         ProtoLog.i(WM_DEBUG_IME, "setInputMethodTarget %s", target);
         mImeLayeringTarget = target;
-        mImeLayeringTargetWaitingAnim = targetWaitingAnim;
 
         // 1. Reparent the IME container window to the target root DA to get the correct bounds and
         // config. (Only happens when the target window is in a different root DA)
@@ -3748,13 +3701,12 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
                 targetRoot.placeImeContainer(mImeWindowsContainer);
             }
         }
-        // 2. Reparent the IME container surface to either the input target app, or the IME window
-        // parent.
-        updateImeParent();
-        // 3. Assign window layers based on the IME surface parent to make sure it is on top of the
+        // 2. Assign window layers based on the IME surface parent to make sure it is on top of the
         // app.
         assignWindowLayers(true /* setLayoutNeeded */);
-        // 4. Update the IME control target to apply any inset change and animation.
+        // 3. Update the IME control target to apply any inset change and animation.
+        // 4. Reparent the IME container surface to either the input target app, or the IME window
+        // parent.
         updateImeControlTarget();
     }
 
@@ -3768,6 +3720,109 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mImeControlTarget = target;
     }
 
+    @VisibleForTesting
+    void attachAndShowImeScreenshotOnTarget() {
+        // No need to attach screenshot if the IME target not exists or screen is off.
+        if (!isImeAttachedToApp() || !mWmService.mPolicy.isScreenOn()) {
+            return;
+        }
+
+        final SurfaceControl.Transaction t = getPendingTransaction();
+        // Prepare IME screenshot for the target if it allows to attach into.
+        if (mInputMethodWindow != null && mInputMethodWindow.isVisible()) {
+            final Task task = mImeLayeringTarget.getTask();
+            // Re-new the IME screenshot when it does not exist or the size changed.
+            final boolean renewImeSurface = mImeScreenshot == null
+                    || mImeScreenshot.getWidth() != mInputMethodWindow.getFrame().width()
+                    || mImeScreenshot.getHeight() != mInputMethodWindow.getFrame().height();
+            if (task != null && !task.isHomeOrRecentsRootTask()) {
+                SurfaceControl.ScreenshotHardwareBuffer imeBuffer = renewImeSurface
+                        ? mWmService.mTaskSnapshotController.snapshotImeFromAttachedTask(task)
+                        : null;
+                if (imeBuffer != null) {
+                    // Remove the last IME surface when the surface needs to renew.
+                    removeImeSurfaceImmediately();
+                    mImeScreenshot = createImeSurface(imeBuffer, t);
+                }
+            }
+        }
+
+        final boolean isValidSnapshot = mImeScreenshot != null && mImeScreenshot.isValid();
+        // Showing the IME screenshot if the target has already in app transition stage.
+        // Note that if the current IME insets is not showing, no need to show IME screenshot
+        // to reflect the true IME insets visibility and the app task layout as possible.
+        if (isValidSnapshot && getInsetsStateController().getImeSourceProvider().isImeShowing()) {
+            if (DEBUG_INPUT_METHOD) {
+                Slog.d(TAG, "show IME snapshot, ime target=" + mImeLayeringTarget);
+            }
+            t.show(mImeScreenshot);
+        } else if (!isValidSnapshot) {
+            removeImeSurfaceImmediately();
+        }
+    }
+
+    @VisibleForTesting
+    SurfaceControl createImeSurface(SurfaceControl.ScreenshotHardwareBuffer imeBuffer,
+            Transaction t) {
+        final HardwareBuffer buffer = imeBuffer.getHardwareBuffer();
+        if (DEBUG_INPUT_METHOD) Slog.d(TAG, "create IME snapshot for "
+                + mImeLayeringTarget + ", buff width=" + buffer.getWidth()
+                + ", height=" + buffer.getHeight());
+        final ActivityRecord activity = mImeLayeringTarget.mActivityRecord;
+        final SurfaceControl imeSurface = mWmService.mSurfaceControlFactory.apply(null)
+                .setName("IME-snapshot-surface")
+                .setBufferSize(buffer.getWidth(), buffer.getHeight())
+                .setFormat(buffer.getFormat())
+                .setParent(activity.getSurfaceControl())
+                .setCallsite("DisplayContent.attachAndShowImeScreenshotOnTarget")
+                .build();
+        // Make IME snapshot as trusted overlay
+        InputMonitor.setTrustedOverlayInputInfo(imeSurface, t, getDisplayId(),
+                "IME-snapshot-surface");
+        Surface surface = mWmService.mSurfaceFactory.get();
+        surface.copyFrom(imeSurface);
+        surface.attachAndQueueBufferWithColorSpace(buffer, null);
+        surface.release();
+        t.setRelativeLayer(imeSurface, activity.getSurfaceControl(), 1);
+        t.setPosition(imeSurface, mInputMethodWindow.getDisplayFrame().left,
+                mInputMethodWindow.getDisplayFrame().top);
+        return imeSurface;
+    }
+
+    /**
+     * Shows the IME screenshot and attach to the IME target window.
+     *
+     * Used when the IME target window with IME visible is transitioning to the next target.
+     * e.g. App transitioning or swiping this the task of the IME target window to recents app.
+     */
+    void showImeScreenshot() {
+        attachAndShowImeScreenshotOnTarget();
+    }
+
+    /**
+     * Removes the IME screenshot when necessary.
+     *
+     * Used when app transition animation finished or obsoleted screenshot surface like size
+     * changed by rotation.
+     */
+    void removeImeScreenshotIfPossible() {
+        if (mImeLayeringTarget == null
+                || mImeLayeringTarget.mAttrs.type != TYPE_APPLICATION_STARTING
+                && !mImeLayeringTarget.isAnimating(PARENTS | TRANSITION,
+                ANIMATION_TYPE_APP_TRANSITION | ANIMATION_TYPE_RECENTS)) {
+            removeImeSurfaceImmediately();
+        }
+    }
+
+    /** Removes the IME screenshot immediately. */
+    void removeImeSurfaceImmediately() {
+        if (mImeScreenshot != null) {
+            if (DEBUG_INPUT_METHOD) Slog.d(TAG, "remove IME snapshot");
+            getSyncTransaction().remove(mImeScreenshot);
+            mImeScreenshot = null;
+        }
+    }
+
     /**
      * The IME input target is the window which receives input from IME. It is also a candidate
      * which controls the visibility and animation of the input method window.
@@ -3781,8 +3836,14 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
     }
 
     void updateImeControlTarget() {
+        InsetsControlTarget prevImeControlTarget = mImeControlTarget;
         mImeControlTarget = computeImeControlTarget();
         mInsetsStateController.onImeControlTargetChanged(mImeControlTarget);
+        // Update Ime parent when IME insets leash created, which is the best time that default
+        // IME visibility has been settled down after IME control target changed.
+        if (prevImeControlTarget != mImeControlTarget) {
+            updateImeParent();
+        }
 
         final WindowState win = InsetsControlTarget.asWindowOrNull(mImeControlTarget);
         final IBinder token = win != null ? win.mClient.asBinder() : null;
@@ -4037,6 +4098,16 @@ class DisplayContent extends RootDisplayArea implements WindowManagerPolicy.Disp
         mWmService.mWindowPlacerLocked.performSurfacePlacement();
     }
 
+    /**
+     * Callbacks when the given type of {@link WindowContainer} animation finished running in the
+     * hierarchy.
+     */
+    void onWindowAnimationFinished(int type) {
+        if (type == ANIMATION_TYPE_APP_TRANSITION || type == ANIMATION_TYPE_RECENTS) {
+            removeImeSurfaceImmediately();
+        }
+    }
+
     // TODO: Super unexpected long method that should be broken down...
     void applySurfaceChangesTransaction() {
         final WindowSurfacePlacer surfacePlacer = mWmService.mWindowPlacerLocked;

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

@@ -132,7 +132,7 @@ class InputConsumerImpl implements IBinder.DeathRecipient {
         t.hide(mInputSurface);
     }
 
-    void show(SurfaceControl.Transaction t, WindowState w) {
+    void show(SurfaceControl.Transaction t, WindowContainer w) {
         t.show(mInputSurface);
         t.setInputWindowInfo(mInputSurface, mWindowHandle);
         t.setRelativeLayer(mInputSurface, w.getSurfaceControl(), 1);

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

@@ -516,7 +516,7 @@ final class InputMonitor {
             if (mAddRecentsAnimationInputConsumerHandle && shouldApplyRecentsInputConsumer) {
                 if (recentsAnimationController.updateInputConsumerForApp(
                         mRecentsAnimationInputConsumer.mWindowHandle)) {
-                    mRecentsAnimationInputConsumer.show(mInputTransaction, w);
+                    mRecentsAnimationInputConsumer.show(mInputTransaction, w.mActivityRecord);
                     mAddRecentsAnimationInputConsumerHandle = false;
                 }
             }

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

@@ -282,6 +282,25 @@ public class RecentsAnimationController implements DeathRecipient {
                             task.setCanAffectSystemUiFlags(behindSystemBars);
                         }
                     }
+                    if (!behindSystemBars) {
+                        // Make sure to update the correct IME parent in case that the IME parent
+                        // may be computed as display layer when re-layout window happens during
+                        // rotation but there is intermediate state that the bounds of task and
+                        // the IME target's activity is not the same during rotating.
+                        mDisplayContent.updateImeParent();
+
+                        // Hiding IME if IME window is not attached to app.
+                        // Since some windowing mode is not proper to snapshot Task with IME window
+                        // while the app transitioning to the next task (e.g. split-screen mode)
+                        if (!mDisplayContent.isImeAttachedToApp()) {
+                            final InputMethodManagerInternal inputMethodManagerInternal =
+                                    LocalServices.getService(InputMethodManagerInternal.class);
+                            if (inputMethodManagerInternal != null) {
+                                inputMethodManagerInternal.hideCurrentInputMethod(
+                                        SoftInputShowHideReason.HIDE_RECENTS_ANIMATION);
+                            }
+                        }
+                    }
                     mService.mWindowPlacerLocked.requestTraversal();
                 }
             } finally {
@@ -299,7 +318,6 @@ public class RecentsAnimationController implements DeathRecipient {
                     if (mCanceled) {
                         return;
                     }
-
                     mInputConsumerEnabled = enabled;
                     final InputMonitor inputMonitor = mDisplayContent.getInputMonitor();
                     inputMonitor.updateInputWindowsLw(true /*force*/);
@@ -310,34 +328,9 @@ public class RecentsAnimationController implements DeathRecipient {
             }
         }
 
+        // TODO(b/166736352): Remove this method without the need to expose to launcher.
         @Override
-        public void hideCurrentInputMethod() {
-            final long token = Binder.clearCallingIdentity();
-            try {
-                synchronized (mService.getWindowManagerLock()) {
-                    // Make sure to update the correct IME parent in case that the IME parent may
-                    // be computed as display layer when re-layout window happens during rotation
-                    // but there is intermediate state that the bounds of task and the IME
-                    // target's activity is not the same during rotating.
-                    mDisplayContent.updateImeParent();
-
-                    // Ignore hiding IME if IME window is attached to app.
-                    // Since we would like to snapshot Task with IME window while transitioning
-                    // to recents.
-                    if (mDisplayContent.isImeAttachedToApp()) {
-                        return;
-                    }
-                }
-                final InputMethodManagerInternal inputMethodManagerInternal =
-                        LocalServices.getService(InputMethodManagerInternal.class);
-                if (inputMethodManagerInternal != null) {
-                    inputMethodManagerInternal.hideCurrentInputMethod(
-                            SoftInputShowHideReason.HIDE_RECENTS_ANIMATION);
-                }
-            } finally {
-                Binder.restoreCallingIdentity(token);
-            }
-        }
+        public void hideCurrentInputMethod() { }
 
         @Override
         public void setDeferCancelUntilNextTransition(boolean defer, boolean screenshot) {

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

@@ -35,6 +35,7 @@ import android.hardware.HardwareBuffer;
 import android.os.Environment;
 import android.os.Handler;
 import android.util.ArraySet;
+import android.util.Pair;
 import android.util.Slog;
 import android.view.InsetsState;
 import android.view.SurfaceControl;
@@ -275,39 +276,12 @@ class TaskSnapshotController {
      */
     @VisibleForTesting
     boolean prepareTaskSnapshot(Task task, int pixelFormat, TaskSnapshot.Builder builder) {
-        if (!mService.mPolicy.isScreenOn()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
-            }
+        final Pair<ActivityRecord, WindowState> result = checkIfReadyToSnapshot(task);
+        if (result == null) {
             return false;
         }
-        final ActivityRecord activity = findAppTokenForSnapshot(task);
-        if (activity == null) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.w(TAG_WM, "Failed to take screenshot. No visible windows for " + task);
-            }
-            return false;
-        }
-        if (activity.hasCommittedReparentToAnimationLeash()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.w(TAG_WM, "Failed to take screenshot. App is animating " + activity);
-            }
-            return false;
-        }
-
-        final WindowState mainWindow = activity.findMainWindow();
-        if (mainWindow == null) {
-            Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
-            return false;
-        }
-        if (activity.hasFixedRotationTransform()) {
-            if (DEBUG_SCREENSHOT) {
-                Slog.i(TAG_WM, "Skip taking screenshot. App has fixed rotation " + activity);
-            }
-            // The activity is in a temporal state that it has different rotation than the task.
-            return false;
-        }
-
+        final ActivityRecord activity = result.first;
+        final WindowState mainWindow = result.second;
         final Rect contentInsets = getSystemBarInsets(task.getBounds(),
                 mainWindow.getInsetsStateWithVisibilityOverride());
         InsetUtils.addInsets(contentInsets, activity.getLetterboxInsets());
@@ -339,6 +313,50 @@ class TaskSnapshotController {
         return true;
     }
 
+    /**
+     * Check if the state of the Task is appropriate to capture a snapshot, such like the task
+     * snapshot or the associated IME surface snapshot.
+     *
+     * @param task the target task to capture the snapshot
+     * @return Pair of (the top activity of the task, the main window of the task) if passed the
+     * state checking. Returns {@code null} if the task state isn't ready to snapshot.
+     */
+    Pair<ActivityRecord, WindowState> checkIfReadyToSnapshot(Task task) {
+        if (!mService.mPolicy.isScreenOn()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
+            }
+            return null;
+        }
+        final ActivityRecord activity = findAppTokenForSnapshot(task);
+        if (activity == null) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.w(TAG_WM, "Failed to take screenshot. No visible windows for " + task);
+            }
+            return null;
+        }
+        if (activity.hasCommittedReparentToAnimationLeash()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.w(TAG_WM, "Failed to take screenshot. App is animating " + activity);
+            }
+            return null;
+        }
+
+        final WindowState mainWindow = activity.findMainWindow();
+        if (mainWindow == null) {
+            Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
+            return null;
+        }
+        if (activity.hasFixedRotationTransform()) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.i(TAG_WM, "Skip taking screenshot. App has fixed rotation " + activity);
+            }
+            // The activity is in a temporal state that it has different rotation than the task.
+            return null;
+        }
+        return new Pair<>(activity, mainWindow);
+    }
+
     @Nullable
     SurfaceControl.ScreenshotHardwareBuffer createTaskSnapshot(@NonNull Task task,
             TaskSnapshot.Builder builder) {
@@ -355,6 +373,43 @@ class TaskSnapshotController {
         return createTaskSnapshot(task, scaleFraction, PixelFormat.RGBA_8888, null, builder);
     }
 
+    @Nullable
+    private SurfaceControl.ScreenshotHardwareBuffer createImeSnapshot(@NonNull Task task,
+            int pixelFormat) {
+        if (task.getSurfaceControl() == null) {
+            if (DEBUG_SCREENSHOT) {
+                Slog.w(TAG_WM, "Failed to take screenshot. No surface control for " + task);
+            }
+            return null;
+        }
+        final WindowState imeWindow = task.getDisplayContent().mInputMethodWindow;
+        SurfaceControl.ScreenshotHardwareBuffer imeBuffer = null;
+        if (imeWindow != null && imeWindow.isWinVisibleLw()) {
+            final Rect bounds = imeWindow.getContainingFrame();
+            bounds.offsetTo(0, 0);
+            imeBuffer = SurfaceControl.captureLayersExcluding(imeWindow.getSurfaceControl(),
+                    bounds, 1.0f, pixelFormat, null);
+        }
+        return imeBuffer;
+    }
+
+    /**
+     * Create the snapshot of the IME surface on the task which used for placing on the closing
+     * task to keep IME visibility while app transitioning.
+     */
+    @Nullable
+    SurfaceControl.ScreenshotHardwareBuffer snapshotImeFromAttachedTask(@NonNull Task task) {
+        // Check if the IME targets task ready to take the corresponding IME snapshot, if not,
+        // means the task is not yet visible for some reasons and no need to snapshot IME surface.
+        if (checkIfReadyToSnapshot(task) == null) {
+            return null;
+        }
+        final int pixelFormat = mPersister.use16BitFormat()
+                    ? PixelFormat.RGB_565
+                    : PixelFormat.RGBA_8888;
+        return createImeSnapshot(task, pixelFormat);
+    }
+
     @Nullable
     SurfaceControl.ScreenshotHardwareBuffer createTaskSnapshot(@NonNull Task task,
             float scaleFraction, int pixelFormat, Point outTaskSize, TaskSnapshot.Builder builder) {

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

@@ -2682,6 +2682,10 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
     protected void applyAnimationUnchecked(WindowManager.LayoutParams lp, boolean enter,
             @TransitionOldType int transit, boolean isVoiceInteraction,
             @Nullable ArrayList<WindowContainer> sources) {
+        final Task task = asTask();
+        if (task != null && !enter && !task.isHomeOrRecentsRootTask()) {
+            mDisplayContent.showImeScreenshot();
+        }
         final Pair<AnimationAdapter, AnimationAdapter> adapters = getAnimationAdapter(lp,
                 transit, enter, isVoiceInteraction);
         AnimationAdapter adapter = adapters.first;
@@ -2826,6 +2830,9 @@ class WindowContainer<E extends WindowContainer> extends ConfigurationContainer<
             mSurfaceAnimationSources.valueAt(i).onAnimationFinished(type, anim);
         }
         mSurfaceAnimationSources.clear();
+        if (mDisplayContent != null) {
+            mDisplayContent.onWindowAnimationFinished(type);
+        }
     }
 
     /**

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

@@ -2166,6 +2166,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();
             // 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);

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

@@ -65,7 +65,11 @@ import static com.android.dx.mockito.inline.extended.ExtendedMockito.times;
 import static com.android.dx.mockito.inline.extended.ExtendedMockito.verify;
 import static com.android.server.wm.DisplayContent.IME_TARGET_INPUT;
 import static com.android.server.wm.DisplayContent.IME_TARGET_LAYERING;
+import static com.android.server.wm.SurfaceAnimator.ANIMATION_TYPE_APP_TRANSITION;
 import static com.android.server.wm.SurfaceAnimator.ANIMATION_TYPE_FIXED_TRANSFORM;
+import static com.android.server.wm.SurfaceAnimator.ANIMATION_TYPE_RECENTS;
+import static com.android.server.wm.WindowContainer.AnimationFlags.PARENTS;
+import static com.android.server.wm.WindowContainer.AnimationFlags.TRANSITION;
 import static com.android.server.wm.WindowContainer.POSITION_TOP;
 import static com.android.server.wm.WindowManagerService.UPDATE_FOCUS_NORMAL;
 
@@ -81,6 +85,7 @@ import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.doAnswer;
@@ -299,22 +304,24 @@ public class DisplayContentTests extends WindowTestsBase {
         assertEquals(startingWin, imeTarget);
         startingWin.mHidden = false;
 
-        // Verify that an app window launching behind the starting window becomes the target
+        // Verify that the starting window still be an ime target even an app window launching
+        // behind it.
         final WindowState appWin = createWindow(null, TYPE_BASE_APPLICATION, activity, "appWin");
         appWin.setHasSurface(true);
         assertTrue(appWin.canBeImeTarget());
 
         imeTarget = mDisplayContent.computeImeTarget(false /* updateImeTarget */);
-        assertEquals(appWin, imeTarget);
+        assertEquals(startingWin, imeTarget);
         appWin.mHidden = false;
 
-        // Verify that an child window can be an ime target even behind a launching app window
+        // Verify that the starting window still be an ime target even the child window behind a
+        // launching app window
         final WindowState childWin = createWindow(appWin,
                 TYPE_APPLICATION_ATTACHED_DIALOG, "childWin");
         childWin.setHasSurface(true);
         assertTrue(childWin.canBeImeTarget());
         imeTarget = mDisplayContent.computeImeTarget(false /* updateImeTarget */);
-        assertEquals(childWin, imeTarget);
+        assertEquals(startingWin, imeTarget);
     }
 
     @UseTestDisplay(addAllCommonWindows = true)
@@ -1732,6 +1739,63 @@ public class DisplayContentTests extends WindowTestsBase {
         verify(child1, never()).needsRelativeLayeringToIme();
     }
 
+    @UseTestDisplay(addWindows = {W_INPUT_METHOD}, addAllCommonWindows = true)
+    @Test
+    public void testAttachAndShowImeScreenshotOnTarget() {
+        // Preparation: Simulate screen state is on.
+        spyOn(mWm.mPolicy);
+        doReturn(true).when(mWm.mPolicy).isScreenOn();
+
+        // Preparation: Simulate snapshot IME surface.
+        spyOn(mWm.mTaskSnapshotController);
+        doReturn(mock(SurfaceControl.ScreenshotHardwareBuffer.class)).when(
+                mWm.mTaskSnapshotController).snapshotImeFromAttachedTask(any());
+        final SurfaceControl imeSurface = mock(SurfaceControl.class);
+        spyOn(imeSurface);
+        doReturn(true).when(imeSurface).isValid();
+        doReturn(imeSurface).when(mDisplayContent).createImeSurface(any(), any());
+
+        // Preparation: Simulate snapshot Task.
+        ActivityRecord act1 = createActivityRecord(mDisplayContent);
+        final WindowState appWin1 = createWindow(null, TYPE_BASE_APPLICATION, act1, "appWin1");
+        spyOn(appWin1);
+        spyOn(appWin1.mWinAnimator);
+        appWin1.setHasSurface(true);
+        assertTrue(appWin1.canBeImeTarget());
+        doReturn(true).when(appWin1.mWinAnimator).getShown();
+        doReturn(true).when(appWin1.mActivityRecord).isSurfaceShowing();
+        appWin1.mWinAnimator.mLastAlpha = 1f;
+
+        // Test step 1: appWin1 is the current IME target and soft-keyboard is visible.
+        mDisplayContent.computeImeTarget(true);
+        assertEquals(appWin1, mDisplayContent.getImeTarget(IME_TARGET_LAYERING));
+        spyOn(mDisplayContent.mInputMethodWindow);
+        doReturn(true).when(mDisplayContent.mInputMethodWindow).isVisible();
+        mDisplayContent.getInsetsStateController().getImeSourceProvider().setImeShowing(true);
+
+        // Test step 2: Simulate launching appWin2 and appWin1 is in app transition.
+        ActivityRecord act2 = createActivityRecord(mDisplayContent);
+        final WindowState appWin2 = createWindow(null, TYPE_BASE_APPLICATION, act2, "appWin2");
+        appWin2.setHasSurface(true);
+        assertTrue(appWin2.canBeImeTarget());
+        doReturn(true).when(appWin1).isAnimating(PARENTS | TRANSITION,
+                ANIMATION_TYPE_APP_TRANSITION | ANIMATION_TYPE_RECENTS);
+
+        // Test step 3: Verify appWin2 will be the next IME target and the IME snapshot surface will
+        // be shown at this time.
+        final Transaction t = mDisplayContent.getPendingTransaction();
+        spyOn(t);
+        mDisplayContent.setImeInputTarget(appWin2);
+        mDisplayContent.computeImeTarget(true);
+        assertEquals(appWin2, mDisplayContent.getImeTarget(IME_TARGET_LAYERING));
+        assertTrue(mDisplayContent.isImeAttachedToApp());
+
+        verify(mDisplayContent, atLeast(1)).attachAndShowImeScreenshotOnTarget();
+        verify(mWm.mTaskSnapshotController).snapshotImeFromAttachedTask(appWin1.getTask());
+        assertNotNull(mDisplayContent.mImeScreenshot);
+        verify(t).show(mDisplayContent.mImeScreenshot);
+    }
+
     private boolean isOptionsPanelAtRight(int displayId) {
         return (mWm.getPreferredOptionsPanelGravity(displayId) & Gravity.RIGHT) == Gravity.RIGHT;
     }