Send core-created snapshots to transition player

This way, the player doesn't have to take a second
screenshot for use during animation.

This also includes luma (only for rotation) since
that requires the actual buffer and we only want to
send the snapshot surface.

Bug: 242056267
Test: atest TransitionTests
Change-Id: Ib76bb821b4aa8c01c6db1827c26b9fb498aae878
diff --git a/core/java/android/window/TransitionInfo.java b/core/java/android/window/TransitionInfo.java
index dc1f612..8ca763e 100644
--- a/core/java/android/window/TransitionInfo.java
+++ b/core/java/android/window/TransitionInfo.java
@@ -397,6 +397,8 @@
         private @Surface.Rotation int mEndFixedRotation = ROTATION_UNDEFINED;
         private int mRotationAnimation = ROTATION_ANIMATION_UNSPECIFIED;
         private @ColorInt int mBackgroundColor;
+        private SurfaceControl mSnapshot = null;
+        private float mSnapshotLuma;
 
         public Change(@Nullable WindowContainerToken container, @NonNull SurfaceControl leash) {
             mContainer = container;
@@ -420,6 +422,8 @@
             mEndFixedRotation = in.readInt();
             mRotationAnimation = in.readInt();
             mBackgroundColor = in.readInt();
+            mSnapshot = in.readTypedObject(SurfaceControl.CREATOR);
+            mSnapshotLuma = in.readFloat();
         }
 
         /** Sets the parent of this change's container. The parent must be a participant or null. */
@@ -489,6 +493,12 @@
             mBackgroundColor = backgroundColor;
         }
 
+        /** Sets a snapshot surface for the "start" state of the container. */
+        public void setSnapshot(@Nullable SurfaceControl snapshot, float luma) {
+            mSnapshot = snapshot;
+            mSnapshotLuma = luma;
+        }
+
         /** @return the container that is changing. May be null if non-remotable (eg. activity) */
         @Nullable
         public WindowContainerToken getContainer() {
@@ -587,6 +597,17 @@
             return mBackgroundColor;
         }
 
+        /** @return a snapshot surface (if applicable). */
+        @Nullable
+        public SurfaceControl getSnapshot() {
+            return mSnapshot;
+        }
+
+        /** @return the luma calculated for the snapshot surface (if applicable). */
+        public float getSnapshotLuma() {
+            return mSnapshotLuma;
+        }
+
         /** @hide */
         @Override
         public void writeToParcel(@NonNull Parcel dest, int flags) {
@@ -605,6 +626,8 @@
             dest.writeInt(mEndFixedRotation);
             dest.writeInt(mRotationAnimation);
             dest.writeInt(mBackgroundColor);
+            dest.writeTypedObject(mSnapshot, flags);
+            dest.writeFloat(mSnapshotLuma);
         }
 
         @NonNull
@@ -629,11 +652,13 @@
 
         @Override
         public String toString() {
-            return "{" + mContainer + "(" + mParent + ") leash=" + mLeash
+            String out = "{" + mContainer + "(" + mParent + ") leash=" + mLeash
                     + " m=" + modeToString(mMode) + " f=" + flagsToString(mFlags) + " sb="
                     + mStartAbsBounds + " eb=" + mEndAbsBounds + " eo=" + mEndRelOffset + " r="
                     + mStartRotation + "->" + mEndRotation + ":" + mRotationAnimation
-                    + " endFixedRotation=" + mEndFixedRotation + "}";
+                    + " endFixedRotation=" + mEndFixedRotation;
+            if (mSnapshot != null) out += " snapshot=" + mSnapshot;
+            return out + "}";
         }
     }
 
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/ScreenRotationAnimation.java b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/ScreenRotationAnimation.java
index 6388ca1..b647f43 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/transition/ScreenRotationAnimation.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/transition/ScreenRotationAnimation.java
@@ -139,39 +139,48 @@
                 .build();
 
         try {
-            SurfaceControl.LayerCaptureArgs args =
-                    new SurfaceControl.LayerCaptureArgs.Builder(mSurfaceControl)
-                            .setCaptureSecureLayers(true)
-                            .setAllowProtected(true)
-                            .setSourceCrop(new Rect(0, 0, mStartWidth, mStartHeight))
-                            .build();
-            SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
-                    SurfaceControl.captureLayers(args);
-            if (screenshotBuffer == null) {
-                Slog.w(TAG, "Unable to take screenshot of display");
-                return;
-            }
+            if (change.getSnapshot() != null) {
+                mScreenshotLayer = change.getSnapshot();
+                t.reparent(mScreenshotLayer, mAnimLeash);
+                mStartLuma = change.getSnapshotLuma();
+            } else {
+                SurfaceControl.LayerCaptureArgs args =
+                        new SurfaceControl.LayerCaptureArgs.Builder(mSurfaceControl)
+                                .setCaptureSecureLayers(true)
+                                .setAllowProtected(true)
+                                .setSourceCrop(new Rect(0, 0, mStartWidth, mStartHeight))
+                                .build();
+                SurfaceControl.ScreenshotHardwareBuffer screenshotBuffer =
+                        SurfaceControl.captureLayers(args);
+                if (screenshotBuffer == null) {
+                    Slog.w(TAG, "Unable to take screenshot of display");
+                    return;
+                }
 
-            mScreenshotLayer = new SurfaceControl.Builder(session)
-                    .setParent(mAnimLeash)
-                    .setBLASTLayer()
-                    .setSecure(screenshotBuffer.containsSecureLayers())
-                    .setOpaque(true)
-                    .setCallsite("ShellRotationAnimation")
-                    .setName("RotationLayer")
-                    .build();
+                mScreenshotLayer = new SurfaceControl.Builder(session)
+                        .setParent(mAnimLeash)
+                        .setBLASTLayer()
+                        .setSecure(screenshotBuffer.containsSecureLayers())
+                        .setOpaque(true)
+                        .setCallsite("ShellRotationAnimation")
+                        .setName("RotationLayer")
+                        .build();
+
+                final ColorSpace colorSpace = screenshotBuffer.getColorSpace();
+                final HardwareBuffer hardwareBuffer = screenshotBuffer.getHardwareBuffer();
+                t.setDataSpace(mScreenshotLayer, colorSpace.getDataSpace());
+                t.setBuffer(mScreenshotLayer, hardwareBuffer);
+                t.show(mScreenshotLayer);
+                if (!isCustomRotate()) {
+                    mStartLuma = getMedianBorderLuma(hardwareBuffer, colorSpace);
+                }
+            }
 
             t.setLayer(mAnimLeash, SCREEN_FREEZE_LAYER_BASE);
             t.show(mAnimLeash);
             // Crop the real content in case it contains a larger child layer, e.g. wallpaper.
             t.setCrop(mSurfaceControl, new Rect(0, 0, mEndWidth, mEndHeight));
 
-            final ColorSpace colorSpace = screenshotBuffer.getColorSpace();
-            final HardwareBuffer hardwareBuffer = screenshotBuffer.getHardwareBuffer();
-            t.setDataSpace(mScreenshotLayer, colorSpace.getDataSpace());
-            t.setBuffer(mScreenshotLayer, hardwareBuffer);
-            t.show(mScreenshotLayer);
-
             if (!isCustomRotate()) {
                 mBackColorSurface = new SurfaceControl.Builder(session)
                         .setParent(rootLeash)
@@ -181,8 +190,6 @@
                         .setName("BackColorSurface")
                         .build();
 
-                mStartLuma = getMedianBorderLuma(hardwareBuffer, colorSpace);
-
                 t.setLayer(mBackColorSurface, -1);
                 t.setColor(mBackColorSurface, new float[]{mStartLuma, mStartLuma, mStartLuma});
                 t.show(mBackColorSurface);
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index 8e389d3..acec9b6 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -92,11 +92,13 @@
 import com.android.internal.protolog.common.ProtoLog;
 import com.android.internal.util.function.pooled.PooledLambda;
 import com.android.server.inputmethod.InputMethodManagerInternal;
+import com.android.server.wm.utils.RotationAnimationUtils;
 
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.function.Predicate;
 
 /**
@@ -288,9 +290,9 @@
         if (mContainerFreezer == null) {
             mContainerFreezer = new ScreenshotFreezer();
         }
-        mIsSeamlessRotation = true;
         final WindowState top = dc.getDisplayPolicy().getTopFullscreenOpaqueWindow();
         if (top != null) {
+            mIsSeamlessRotation = true;
             top.mSyncMethodOverride = BLASTSyncEngine.METHOD_BLAST;
             ProtoLog.v(ProtoLogGroup.WM_DEBUG_WINDOW_TRANSITIONS, "Override sync-method for %s "
                     + "because seamless rotating", top.getName());
@@ -1630,6 +1632,9 @@
                 change.setEndAbsBounds(bounds);
             }
             change.setRotation(info.mRotation, endRotation);
+            if (info.mSnapshot != null) {
+                change.setSnapshot(info.mSnapshot, info.mSnapshotLuma);
+            }
 
             out.addChange(change);
         }
@@ -1785,6 +1790,10 @@
         /** These are just extra info. They aren't used for change-detection. */
         @Flag int mFlags = FLAG_NONE;
 
+        /** Snapshot surface and luma, if relevant. */
+        SurfaceControl mSnapshot;
+        float mSnapshotLuma;
+
         ChangeInfo(@NonNull WindowContainer origState) {
             mVisible = origState.isVisibleRequested();
             mWindowingMode = origState.getWindowingMode();
@@ -2093,8 +2102,8 @@
      */
     @VisibleForTesting
     private class ScreenshotFreezer implements IContainerFreezer {
-        /** Values are the screenshot "surfaces" or null if it was frozen via BLAST override. */
-        private final ArrayMap<WindowContainer, SurfaceControl> mSnapshots = new ArrayMap<>();
+        /** Keeps track of which windows are frozen. Not all frozen windows have snapshots. */
+        private final ArraySet<WindowContainer> mFrozen = new ArraySet<>();
 
         /** Takes a screenshot and puts it at the top of the container's surface. */
         @Override
@@ -2104,7 +2113,7 @@
             // Check if any parents have already been "frozen". If so, `wc` is already part of that
             // snapshot, so just skip it.
             for (WindowContainer p = wc; p != null; p = p.getParent()) {
-                if (mSnapshots.containsKey(p)) return false;
+                if (mFrozen.contains(p)) return false;
             }
 
             if (mIsSeamlessRotation) {
@@ -2113,7 +2122,7 @@
                 if (top != null && (top == wc || top.isDescendantOf(wc))) {
                     // Don't use screenshots for seamless windows: these will use BLAST even if not
                     // BLAST mode.
-                    mSnapshots.put(wc, null);
+                    mFrozen.add(wc);
                     return true;
                 }
             }
@@ -2146,7 +2155,15 @@
                     .setCallsite("Transition.ScreenshotSync")
                     .setBLASTLayer()
                     .build();
-            mSnapshots.put(wc, snapshotSurface);
+            mFrozen.add(wc);
+            final ChangeInfo changeInfo = Objects.requireNonNull(mChanges.get(wc));
+            changeInfo.mSnapshot = snapshotSurface;
+            if (wc.asDisplayContent() != null) {
+                // This isn't cheap, so only do it for rotations: assume display-level is rotate
+                // since most of the time it is.
+                changeInfo.mSnapshotLuma = RotationAnimationUtils.getMedianBorderLuma(
+                        screenshotBuffer.getHardwareBuffer(), screenshotBuffer.getColorSpace());
+            }
             SurfaceControl.Transaction t = wc.mWmService.mTransactionFactory.get();
 
             t.setBuffer(snapshotSurface, buffer);
@@ -2167,11 +2184,12 @@
 
         @Override
         public void cleanUp(SurfaceControl.Transaction t) {
-            for (int i = 0; i < mSnapshots.size(); ++i) {
-                SurfaceControl snap = mSnapshots.valueAt(i);
+            for (int i = 0; i < mFrozen.size(); ++i) {
+                SurfaceControl snap =
+                        Objects.requireNonNull(mChanges.get(mFrozen.valueAt(i))).mSnapshot;
                 // May be null if it was frozen via BLAST override.
                 if (snap == null) continue;
-                t.remove(snap);
+                t.reparent(snap, null /* newParent */);
             }
         }
     }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
index 4f68e98..c9438bb 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -27,6 +27,7 @@
 import static android.view.WindowManager.LayoutParams.TYPE_NAVIGATION_BAR_PANEL;
 import static android.view.WindowManager.LayoutParams.TYPE_STATUS_BAR;
 import static android.view.WindowManager.LayoutParams.TYPE_WALLPAPER;
+import static android.view.WindowManager.TRANSIT_CHANGE;
 import static android.view.WindowManager.TRANSIT_CLOSE;
 import static android.view.WindowManager.TRANSIT_OPEN;
 import static android.view.WindowManager.TRANSIT_TO_BACK;
@@ -82,6 +83,7 @@
 
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Objects;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
@@ -1167,6 +1169,42 @@
         assertTrue(mSyncEngine.isReady(transition.getSyncId()));
     }
 
+    @Test
+    public void testVisibleChange_snapshot() {
+        registerTestTransitionPlayer();
+        final ActivityRecord app = createActivityRecord(mDisplayContent);
+        final Transition transition = new Transition(TRANSIT_CHANGE, 0 /* flags */,
+                app.mTransitionController, mWm.mSyncEngine);
+        app.mTransitionController.moveToCollecting(transition, BLASTSyncEngine.METHOD_NONE);
+        final SurfaceControl mockSnapshot = mock(SurfaceControl.class);
+        transition.setContainerFreezer(new Transition.IContainerFreezer() {
+            @Override
+            public boolean freeze(@NonNull WindowContainer wc, @NonNull Rect bounds) {
+                Objects.requireNonNull(transition.mChanges.get(wc)).mSnapshot = mockSnapshot;
+                return true;
+            }
+
+            @Override
+            public void cleanUp(SurfaceControl.Transaction t) {
+            }
+        });
+        final Task task = app.getTask();
+        transition.collect(task);
+        final Rect bounds = new Rect(task.getBounds());
+        Configuration c = new Configuration(task.getRequestedOverrideConfiguration());
+        bounds.inset(10, 10);
+        c.windowConfiguration.setBounds(bounds);
+        task.onRequestedOverrideConfigurationChanged(c);
+
+        ArrayList<WindowContainer> targets = Transition.calculateTargets(
+                transition.mParticipants, transition.mChanges);
+        TransitionInfo info = Transition.calculateTransitionInfo(
+                TRANSIT_CHANGE, 0, targets, transition.mChanges, mMockT);
+        assertEquals(mockSnapshot,
+                info.getChange(task.mRemoteToken.toWindowContainerToken()).getSnapshot());
+        transition.abort();
+    }
+
     private static void makeTaskOrganized(Task... tasks) {
         final ITaskOrganizer organizer = mock(ITaskOrganizer.class);
         for (Task t : tasks) {