Use snapshot in enter content-pip animation

When entering content-pip
- We take a snapshot for the host task within the
  ActivityStarter#startActivityInner function, this is to make sure that
  the snapshot is taken before an app switches to the placeholder content
- We retrieve the snapshot in WM Shell when animating the newly created
  content-pip task, and use that snapshot during the animation
- This works similar to the solid color content overlay when app enters
  PiP without specifying a valid source rect hint. Therefore, we
  abstract the functionality to a dedicated PipContentOverlay class

Video: http://recall/-/aaaaaabFQoRHlzixHdtY/dRSTf7prHc0krPhTKZ0kXx
Bug: 220191201
Test: enter content-pip from ApiDemos app, see also Video
Change-Id: I2420402341a910b44bd8f64e3ad518ebef41bc19

8 files changed, 229 insertions, 65 deletions

diff --git a/core/java/android/app/TaskInfo.java b/core/java/android/app/TaskInfo.java
index 1b87945ec985..7910f1a426c2 100644
--- a/core/java/android/app/TaskInfo.java
+++ b/core/java/android/app/TaskInfo.java
@@ -34,10 +34,7 @@ import android.graphics.Rect;
 import android.os.Build;
 import android.os.IBinder;
 import android.os.Parcel;
-import android.os.RemoteException;
-import android.util.Log;
 import android.view.DisplayCutout;
-import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 
 import java.lang.annotation.Retention;
@@ -355,20 +352,6 @@ public class TaskInfo {
         return isVisible;
     }
 
-    /**
-     * @param isLowResolution
-     * @return
-     * @hide
-     */
-    public TaskSnapshot getTaskSnapshot(boolean isLowResolution) {
-        try {
-            return ActivityTaskManager.getService().getTaskSnapshot(taskId, isLowResolution);
-        } catch (RemoteException e) {
-            Log.e(TAG, "Failed to get task snapshot, taskId=" + taskId, e);
-            return null;
-        }
-    }
-
     /** @hide */
     @NonNull
     @TestApi
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
index d357655882ff..22d1564cf2c4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipAnimationController.java
@@ -27,13 +27,11 @@ import android.animation.ValueAnimator;
 import android.annotation.IntDef;
 import android.app.TaskInfo;
 import android.content.Context;
-import android.content.res.TypedArray;
-import android.graphics.Color;
 import android.graphics.Rect;
 import android.view.Choreographer;
 import android.view.Surface;
 import android.view.SurfaceControl;
-import android.view.SurfaceSession;
+import android.window.TaskSnapshot;
 
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.graphics.SfVsyncFrameCallbackProvider;
@@ -257,7 +255,7 @@ public class PipAnimationController {
                 mSurfaceControlTransactionFactory;
         private PipSurfaceTransactionHelper mSurfaceTransactionHelper;
         private @TransitionDirection int mTransitionDirection;
-        protected SurfaceControl mContentOverlay;
+        protected PipContentOverlay mContentOverlay;
 
         private PipTransitionAnimator(TaskInfo taskInfo, SurfaceControl leash,
                 @AnimationType int animationType,
@@ -335,43 +333,26 @@ public class PipAnimationController {
             return false;
         }
 
-        SurfaceControl getContentOverlay() {
-            return mContentOverlay;
+        SurfaceControl getContentOverlayLeash() {
+            return mContentOverlay == null ? null : mContentOverlay.mLeash;
         }
 
-        PipTransitionAnimator<T> setUseContentOverlay(Context context) {
+        void setColorContentOverlay(Context context) {
             final SurfaceControl.Transaction tx = newSurfaceControlTransaction();
             if (mContentOverlay != null) {
-                // remove existing content overlay if there is any.
-                tx.remove(mContentOverlay);
-                tx.apply();
+                mContentOverlay.detach(tx);
             }
-            mContentOverlay = new SurfaceControl.Builder(new SurfaceSession())
-                    .setCallsite("PipAnimation")
-                    .setName("PipContentOverlay")
-                    .setColorLayer()
-                    .build();
-            tx.show(mContentOverlay);
-            tx.setLayer(mContentOverlay, Integer.MAX_VALUE);
-            tx.setColor(mContentOverlay, getContentOverlayColor(context));
-            tx.setAlpha(mContentOverlay, 0f);
-            tx.reparent(mContentOverlay, mLeash);
-            tx.apply();
-            return this;
+            mContentOverlay = new PipContentOverlay.PipColorOverlay(context);
+            mContentOverlay.attach(tx, mLeash);
         }
 
-        private float[] getContentOverlayColor(Context context) {
-            final TypedArray ta = context.obtainStyledAttributes(new int[] {
-                    android.R.attr.colorBackground });
-            try {
-                int colorAccent = ta.getColor(0, 0);
-                return new float[] {
-                        Color.red(colorAccent) / 255f,
-                        Color.green(colorAccent) / 255f,
-                        Color.blue(colorAccent) / 255f };
-            } finally {
-                ta.recycle();
+        void setSnapshotContentOverlay(TaskSnapshot snapshot, Rect sourceRectHint) {
+            final SurfaceControl.Transaction tx = newSurfaceControlTransaction();
+            if (mContentOverlay != null) {
+                mContentOverlay.detach(tx);
             }
+            mContentOverlay = new PipContentOverlay.PipSnapshotOverlay(snapshot, sourceRectHint);
+            mContentOverlay.attach(tx, mLeash);
         }
 
         /**
@@ -575,7 +556,7 @@ public class PipAnimationController {
                     final Rect start = getStartValue();
                     final Rect end = getEndValue();
                     if (mContentOverlay != null) {
-                        tx.setAlpha(mContentOverlay, fraction < 0.5f ? 0 : (fraction - 0.5f) * 2);
+                        mContentOverlay.onAnimationUpdate(tx, fraction);
                     }
                     if (rotatedEndRect != null) {
                         // Animate the bounds in a different orientation. It only happens when
@@ -680,7 +661,7 @@ public class PipAnimationController {
                             .round(tx, leash, shouldApplyCornerRadius())
                             .shadow(tx, leash, shouldApplyShadowRadius());
                     // TODO(b/178632364): this is a work around for the black background when
-                    // entering PiP in buttion navigation mode.
+                    // entering PiP in button navigation mode.
                     if (isInPipDirection(direction)) {
                         tx.setWindowCrop(leash, getStartValue());
                     }
@@ -704,6 +685,9 @@ public class PipAnimationController {
                     } else {
                         getSurfaceTransactionHelper().crop(tx, leash, destBounds);
                     }
+                    if (mContentOverlay != null) {
+                        mContentOverlay.onAnimationEnd(tx, destBounds);
+                    }
                 }
 
                 @Override
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipContentOverlay.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipContentOverlay.java
new file mode 100644
index 000000000000..a032b338aae8
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipContentOverlay.java
@@ -0,0 +1,164 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.wm.shell.pip;
+
+import android.content.Context;
+import android.content.res.TypedArray;
+import android.graphics.Color;
+import android.graphics.Rect;
+import android.view.SurfaceControl;
+import android.view.SurfaceSession;
+import android.window.TaskSnapshot;
+
+/**
+ * Represents the content overlay used during the entering PiP animation.
+ */
+public abstract class PipContentOverlay {
+    protected SurfaceControl mLeash;
+
+    /** Attaches the internal {@link #mLeash} to the given parent leash. */
+    public abstract void attach(SurfaceControl.Transaction tx, SurfaceControl parentLeash);
+
+    /** Detaches the internal {@link #mLeash} from its parent by removing itself. */
+    public void detach(SurfaceControl.Transaction tx) {
+        if (mLeash != null && mLeash.isValid()) {
+            tx.remove(mLeash);
+            tx.apply();
+        }
+    }
+
+    /**
+     * Animates the internal {@link #mLeash} by a given fraction.
+     * @param atomicTx {@link SurfaceControl.Transaction} to operate, you should not explicitly
+     *                 call apply on this transaction, it should be applied on the caller side.
+     * @param fraction progress of the animation ranged from 0f to 1f.
+     */
+    public abstract void onAnimationUpdate(SurfaceControl.Transaction atomicTx, float fraction);
+
+    /**
+     * Callback when reaches the end of animation on the internal {@link #mLeash}.
+     * @param atomicTx {@link SurfaceControl.Transaction} to operate, you should not explicitly
+     *                 call apply on this transaction, it should be applied on the caller side.
+     * @param destinationBounds {@link Rect} of the final bounds.
+     */
+    public abstract void onAnimationEnd(SurfaceControl.Transaction atomicTx,
+            Rect destinationBounds);
+
+    /** A {@link PipContentOverlay} uses solid color. */
+    public static final class PipColorOverlay extends PipContentOverlay {
+        private final Context mContext;
+
+        public PipColorOverlay(Context context) {
+            mContext = context;
+            mLeash = new SurfaceControl.Builder(new SurfaceSession())
+                    .setCallsite("PipAnimation")
+                    .setName(PipColorOverlay.class.getSimpleName())
+                    .setColorLayer()
+                    .build();
+        }
+
+        @Override
+        public void attach(SurfaceControl.Transaction tx, SurfaceControl parentLeash) {
+            tx.show(mLeash);
+            tx.setLayer(mLeash, Integer.MAX_VALUE);
+            tx.setColor(mLeash, getContentOverlayColor(mContext));
+            tx.setAlpha(mLeash, 0f);
+            tx.reparent(mLeash, parentLeash);
+            tx.apply();
+        }
+
+        @Override
+        public void onAnimationUpdate(SurfaceControl.Transaction atomicTx, float fraction) {
+            atomicTx.setAlpha(mLeash, fraction < 0.5f ? 0 : (fraction - 0.5f) * 2);
+        }
+
+        @Override
+        public void onAnimationEnd(SurfaceControl.Transaction atomicTx, Rect destinationBounds) {
+            // Do nothing. Color overlay should be fully opaque by now.
+        }
+
+        private float[] getContentOverlayColor(Context context) {
+            final TypedArray ta = context.obtainStyledAttributes(new int[] {
+                    android.R.attr.colorBackground });
+            try {
+                int colorAccent = ta.getColor(0, 0);
+                return new float[] {
+                        Color.red(colorAccent) / 255f,
+                        Color.green(colorAccent) / 255f,
+                        Color.blue(colorAccent) / 255f };
+            } finally {
+                ta.recycle();
+            }
+        }
+    }
+
+    /** A {@link PipContentOverlay} uses {@link TaskSnapshot}. */
+    public static final class PipSnapshotOverlay extends PipContentOverlay {
+        private final TaskSnapshot mSnapshot;
+        private final Rect mSourceRectHint;
+
+        private float mTaskSnapshotScaleX;
+        private float mTaskSnapshotScaleY;
+
+        public PipSnapshotOverlay(TaskSnapshot snapshot, Rect sourceRectHint) {
+            mSnapshot = snapshot;
+            mSourceRectHint = new Rect(sourceRectHint);
+            mLeash = new SurfaceControl.Builder(new SurfaceSession())
+                    .setCallsite("PipAnimation")
+                    .setName(PipSnapshotOverlay.class.getSimpleName())
+                    .build();
+        }
+
+        @Override
+        public void attach(SurfaceControl.Transaction tx, SurfaceControl parentLeash) {
+            mTaskSnapshotScaleX = (float) mSnapshot.getTaskSize().x
+                    / mSnapshot.getHardwareBuffer().getWidth();
+            mTaskSnapshotScaleY = (float) mSnapshot.getTaskSize().y
+                    / mSnapshot.getHardwareBuffer().getHeight();
+            tx.show(mLeash);
+            tx.setLayer(mLeash, Integer.MAX_VALUE);
+            tx.setBuffer(mLeash, mSnapshot.getHardwareBuffer());
+            tx.setCrop(mLeash, mSourceRectHint);
+            tx.setScale(mLeash, mTaskSnapshotScaleX, mTaskSnapshotScaleY);
+            tx.reparent(mLeash, parentLeash);
+            tx.apply();
+        }
+
+        @Override
+        public void onAnimationUpdate(SurfaceControl.Transaction atomicTx, float fraction) {
+            // Do nothing. Keep the snapshot till animation ends.
+        }
+
+        @Override
+        public void onAnimationEnd(SurfaceControl.Transaction atomicTx, Rect destinationBounds) {
+            // Work around to make sure the snapshot overlay is aligned with PiP window before
+            // the atomicTx is committed along with the final WindowContainerTransaction.
+            final SurfaceControl.Transaction nonAtomicTx = new SurfaceControl.Transaction();
+            final float scaleX = (float) destinationBounds.width()
+                    / mSnapshot.getHardwareBuffer().getWidth();
+            final float scaleY = (float) destinationBounds.height()
+                    / mSnapshot.getHardwareBuffer().getHeight();
+            final float scale = Math.max(scaleX, scaleY);
+            nonAtomicTx.setScale(mLeash, scale, scale);
+            nonAtomicTx.setPosition(mLeash,
+                    -scale * mSourceRectHint.left / mTaskSnapshotScaleX,
+                    -scale * mSourceRectHint.top / mTaskSnapshotScaleY);
+            nonAtomicTx.apply();
+            atomicTx.remove(mLeash);
+        }
+    }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
index 85052148a510..170908d9d282 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTaskOrganizer.java
@@ -66,6 +66,7 @@ import android.view.Display;
 import android.view.Surface;
 import android.view.SurfaceControl;
 import android.window.TaskOrganizer;
+import android.window.TaskSnapshot;
 import android.window.WindowContainerToken;
 import android.window.WindowContainerTransaction;
 
@@ -152,8 +153,8 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
             final int direction = animator.getTransitionDirection();
             final int animationType = animator.getAnimationType();
             final Rect destinationBounds = animator.getDestinationBounds();
-            if (isInPipDirection(direction) && animator.getContentOverlay() != null) {
-                fadeOutAndRemoveOverlay(animator.getContentOverlay(),
+            if (isInPipDirection(direction) && animator.getContentOverlayLeash() != null) {
+                fadeOutAndRemoveOverlay(animator.getContentOverlayLeash(),
                         animator::clearContentOverlay, true /* withStartDelay*/);
             }
             if (mWaitForFixedRotation && animationType == ANIM_TYPE_BOUNDS
@@ -186,8 +187,8 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
         public void onPipAnimationCancel(TaskInfo taskInfo,
                 PipAnimationController.PipTransitionAnimator animator) {
             final int direction = animator.getTransitionDirection();
-            if (isInPipDirection(direction) && animator.getContentOverlay() != null) {
-                fadeOutAndRemoveOverlay(animator.getContentOverlay(),
+            if (isInPipDirection(direction) && animator.getContentOverlayLeash() != null) {
+                fadeOutAndRemoveOverlay(animator.getContentOverlayLeash(),
                         animator::clearContentOverlay, true /* withStartDelay */);
             }
             sendOnPipTransitionCancelled(direction);
@@ -803,8 +804,9 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
         final PipAnimationController.PipTransitionAnimator<?> animator =
                 mPipAnimationController.getCurrentAnimator();
         if (animator != null) {
-            if (animator.getContentOverlay() != null) {
-                removeContentOverlay(animator.getContentOverlay(), animator::clearContentOverlay);
+            if (animator.getContentOverlayLeash() != null) {
+                removeContentOverlay(animator.getContentOverlayLeash(),
+                        animator::clearContentOverlay);
             }
             animator.removeAllUpdateListeners();
             animator.removeAllListeners();
@@ -1486,7 +1488,17 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
         if (isInPipDirection(direction)) {
             // Similar to auto-enter-pip transition, we use content overlay when there is no
             // source rect hint to enter PiP use bounds animation.
-            if (sourceHintRect == null) animator.setUseContentOverlay(mContext);
+            if (sourceHintRect == null) {
+                animator.setColorContentOverlay(mContext);
+            } else {
+                final TaskSnapshot snapshot = PipUtils.getTaskSnapshot(
+                        mTaskInfo.launchIntoPipHostTaskId, false /* isLowResolution */);
+                if (snapshot != null) {
+                    // use the task snapshot during the animation, this is for
+                    // launch-into-pip aka. content-pip use case.
+                    animator.setSnapshotContentOverlay(snapshot, sourceHintRect);
+                }
+            }
             // The destination bounds are used for the end rect of animation and the final bounds
             // after animation finishes. So after the animation is started, the destination bounds
             // can be updated to new rotation (computeRotatedBounds has changed the DisplayLayout
@@ -1550,7 +1562,7 @@ public class PipTaskOrganizer implements ShellTaskOrganizer.TaskListener,
      */
     void fadeOutAndRemoveOverlay(SurfaceControl surface, Runnable callback,
             boolean withStartDelay) {
-        if (surface == null) {
+        if (surface == null || !surface.isValid()) {
             return;
         }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
index 48df28ee4cde..36e712459863 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransition.java
@@ -709,7 +709,7 @@ public class PipTransition extends PipTransitionController {
             if (sourceHintRect == null) {
                 // We use content overlay when there is no source rect hint to enter PiP use bounds
                 // animation.
-                animator.setUseContentOverlay(mContext);
+                animator.setColorContentOverlay(mContext);
             }
         } else if (mOneShotAnimationType == ANIM_TYPE_ALPHA) {
             startTransaction.setAlpha(leash, 0f);
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
index 24993c621e3c..54f46e0c9938 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipTransitionController.java
@@ -70,8 +70,8 @@ public abstract class PipTransitionController implements Transitions.TransitionH
                     if (direction == TRANSITION_DIRECTION_REMOVE_STACK) {
                         return;
                     }
-                    if (isInPipDirection(direction) && animator.getContentOverlay() != null) {
-                        mPipOrganizer.fadeOutAndRemoveOverlay(animator.getContentOverlay(),
+                    if (isInPipDirection(direction) && animator.getContentOverlayLeash() != null) {
+                        mPipOrganizer.fadeOutAndRemoveOverlay(animator.getContentOverlayLeash(),
                                 animator::clearContentOverlay, true /* withStartDelay*/);
                     }
                     onFinishResize(taskInfo, animator.getDestinationBounds(), direction, tx);
@@ -82,8 +82,8 @@ public abstract class PipTransitionController implements Transitions.TransitionH
                 public void onPipAnimationCancel(TaskInfo taskInfo,
                         PipAnimationController.PipTransitionAnimator animator) {
                     final int direction = animator.getTransitionDirection();
-                    if (isInPipDirection(direction) && animator.getContentOverlay() != null) {
-                        mPipOrganizer.fadeOutAndRemoveOverlay(animator.getContentOverlay(),
+                    if (isInPipDirection(direction) && animator.getContentOverlayLeash() != null) {
+                        mPipOrganizer.fadeOutAndRemoveOverlay(animator.getContentOverlayLeash(),
                                 animator::clearContentOverlay, true /* withStartDelay */);
                     }
                     sendOnPipTransitionCancelled(animator.getTransitionDirection());
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipUtils.java b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipUtils.java
index c6cf8b8b0566..dc60bcf742ce 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipUtils.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/pip/PipUtils.java
@@ -19,13 +19,16 @@ package com.android.wm.shell.pip;
 import static android.app.WindowConfiguration.ACTIVITY_TYPE_UNDEFINED;
 import static android.app.WindowConfiguration.WINDOWING_MODE_PINNED;
 
+import android.annotation.Nullable;
 import android.app.ActivityTaskManager;
 import android.app.ActivityTaskManager.RootTaskInfo;
 import android.app.RemoteAction;
 import android.content.ComponentName;
 import android.content.Context;
 import android.os.RemoteException;
+import android.util.Log;
 import android.util.Pair;
+import android.window.TaskSnapshot;
 
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.wm.shell.protolog.ShellProtoLogGroup;
@@ -106,4 +109,17 @@ public class PipUtils {
         }
         return false;
     }
+
+    /** @return {@link TaskSnapshot} for a given task id. */
+    @Nullable
+    public static TaskSnapshot getTaskSnapshot(int taskId, boolean isLowResolution) {
+        if (taskId <= 0) return null;
+        try {
+            return ActivityTaskManager.getService().getTaskSnapshot(
+                    taskId, isLowResolution);
+        } catch (RemoteException e) {
+            Log.e(TAG, "Failed to get task snapshot, taskId=" + taskId, e);
+            return null;
+        }
+    }
 }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 7240fd510761..ef18b50e910b 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2114,6 +2114,11 @@ class RootWindowContainer extends WindowContainer<DisplayContent>
             rootTask.setWindowingMode(WINDOWING_MODE_PINNED);
             // Set the launch bounds for launch-into-pip Activity on the root task.
             if (r.getOptions() != null && r.getOptions().isLaunchIntoPip()) {
+                // Record the snapshot now, it will be later fetched for content-pip animation.
+                // We do this early in the process to make sure the right snapshot is used for
+                // entering content-pip animation.
+                mWindowManager.mTaskSnapshotController.recordTaskSnapshot(
+                        task, false /* allowSnapshotHome */);
                 rootTask.setBounds(r.getOptions().getLaunchBounds());
             }
             rootTask.setDeferTaskAppear(false);