Polish ActivityEmbedding enter/exit PiP (1/n)

1. When moveActivityToPinnedRootTask with creating a new Task for PiP,
   make sure the Task's initial bounds is the same as the activity
   parent TaskFragment so the animation starts from the correct bounds.
2. When exit PiP to previous Task, make sure we are animating the
   correct window surface. For the previous implementation. there can
   also be TRANSIT_CHANGE change for entering ActivityEmbedding split
   (from PiP) in the same transition.

Bug: 207070762
Test: atest WmTests:RootWindowContainerTests
Test: atest WmTests:TransitionTests
Change-Id: Ifba090ad9ac9fb7033d343eab1c87c1a67bb9c11
diff --git a/core/java/android/window/TransitionInfo.java b/core/java/android/window/TransitionInfo.java
index fbdd325..1bc8e6d 100644
--- a/core/java/android/window/TransitionInfo.java
+++ b/core/java/android/window/TransitionInfo.java
@@ -407,6 +407,7 @@
     public static final class Change implements Parcelable {
         private final WindowContainerToken mContainer;
         private WindowContainerToken mParent;
+        private WindowContainerToken mLastParent;
         private final SurfaceControl mLeash;
         private @TransitionMode int mMode = TRANSIT_NONE;
         private @ChangeFlags int mFlags = FLAG_NONE;
@@ -435,6 +436,7 @@
         private Change(Parcel in) {
             mContainer = in.readTypedObject(WindowContainerToken.CREATOR);
             mParent = in.readTypedObject(WindowContainerToken.CREATOR);
+            mLastParent = in.readTypedObject(WindowContainerToken.CREATOR);
             mLeash = new SurfaceControl();
             mLeash.readFromParcel(in);
             mMode = in.readInt();
@@ -458,6 +460,14 @@
             mParent = parent;
         }
 
+        /**
+         * Sets the parent of this change's container before the transition if this change's
+         * container is reparented in the transition.
+         */
+        public void setLastParent(@Nullable WindowContainerToken lastParent) {
+            mLastParent = lastParent;
+        }
+
         /** Sets the transition mode for this change */
         public void setMode(@TransitionMode int mode) {
             mMode = mode;
@@ -541,6 +551,17 @@
             return mParent;
         }
 
+        /**
+         * @return the parent of the changing container before the transition if it is reparented
+         * in the transition. The parent window may not be collected in the transition as a
+         * participant, and it may have been detached from the display. {@code null} if the changing
+         * container has not been reparented in the transition, or if the parent is not organizable.
+         */
+        @Nullable
+        public WindowContainerToken getLastParent() {
+            return mLastParent;
+        }
+
         /** @return which action this change represents. */
         public @TransitionMode int getMode() {
             return mMode;
@@ -640,6 +661,7 @@
         public void writeToParcel(@NonNull Parcel dest, int flags) {
             dest.writeTypedObject(mContainer, flags);
             dest.writeTypedObject(mParent, flags);
+            dest.writeTypedObject(mLastParent, flags);
             mLeash.writeToParcel(dest, flags);
             dest.writeInt(mMode);
             dest.writeInt(mFlags);
@@ -685,6 +707,7 @@
                     + mStartRotation + "->" + mEndRotation + ":" + mRotationAnimation
                     + " endFixedRotation=" + mEndFixedRotation;
             if (mSnapshot != null) out += " snapshot=" + mSnapshot;
+            if (mLastParent != null) out += " lastParent=" + mLastParent;
             return out + "}";
         }
     }
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 33761d2..2b36b4c 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
@@ -452,14 +452,17 @@
             @NonNull Transitions.TransitionFinishCallback finishCallback,
             @NonNull TaskInfo taskInfo, @Nullable TransitionInfo.Change pipTaskChange) {
         TransitionInfo.Change pipChange = pipTaskChange;
-        if (pipChange == null) {
+        if (mCurrentPipTaskToken == null) {
+            ProtoLog.w(ShellProtoLogGroup.WM_SHELL_PICTURE_IN_PICTURE,
+                    "%s: There is no existing PiP Task for TRANSIT_EXIT_PIP", TAG);
+        } else if (pipChange == null) {
             // The pipTaskChange is null, this can happen if we are reparenting the PIP activity
             // back to its original Task. In that case, we should animate the activity leash
-            // instead, which should be the only non-task, independent, TRANSIT_CHANGE window.
+            // instead, which should be the change whose last parent is the recorded PiP Task.
             for (int i = info.getChanges().size() - 1; i >= 0; --i) {
                 final TransitionInfo.Change change = info.getChanges().get(i);
-                if (change.getTaskInfo() == null && change.getMode() == TRANSIT_CHANGE
-                        && TransitionInfo.isIndependent(change, info)) {
+                if (mCurrentPipTaskToken.equals(change.getLastParent())) {
+                    // Find the activity that is exiting PiP.
                     pipChange = change;
                     break;
                 }
diff --git a/services/core/java/com/android/server/wm/RootWindowContainer.java b/services/core/java/com/android/server/wm/RootWindowContainer.java
index 38c7595..3c55396 100644
--- a/services/core/java/com/android/server/wm/RootWindowContainer.java
+++ b/services/core/java/com/android/server/wm/RootWindowContainer.java
@@ -2020,7 +2020,12 @@
                 // non-fullscreen bounds. Then when this new PIP task exits PIP, it can restore
                 // to its previous freeform bounds.
                 rootTask.setLastNonFullscreenBounds(task.mLastNonFullscreenBounds);
-                rootTask.setBounds(task.getBounds());
+                // When creating a new Task for PiP, set its initial bounds as the TaskFragment in
+                // case the activity is embedded, so that it can be animated to PiP window from the
+                // current bounds.
+                // Use Task#setBoundsUnchecked to skip checking windowing mode as the windowing mode
+                // will be updated later after this is collected in transition.
+                rootTask.setBoundsUnchecked(r.getTaskFragment().getBounds());
 
                 // Move the last recents animation transaction from original task to the new one.
                 if (task.mLastRecentsAnimationTransaction != null) {
diff --git a/services/core/java/com/android/server/wm/Task.java b/services/core/java/com/android/server/wm/Task.java
index 39ea7a8..155cf28 100644
--- a/services/core/java/com/android/server/wm/Task.java
+++ b/services/core/java/com/android/server/wm/Task.java
@@ -2605,6 +2605,13 @@
         return boundsChange;
     }
 
+    /** Sets the requested bounds regardless of the windowing mode. */
+    int setBoundsUnchecked(@NonNull Rect bounds) {
+        final int boundsChange = super.setBounds(bounds);
+        updateSurfaceBounds();
+        return boundsChange;
+    }
+
     @Override
     public boolean isCompatible(int windowingMode, int activityType) {
         // TODO: Should we just move this to ConfigurationContainer?
@@ -5849,10 +5856,7 @@
             return BOUNDS_CHANGE_NONE;
         }
 
-        final int result = super.setBounds(!inMultiWindowMode() ? null : bounds);
-
-        updateSurfaceBounds();
-        return result;
+        return setBoundsUnchecked(!inMultiWindowMode() ? null : bounds);
     }
 
     @Override
diff --git a/services/core/java/com/android/server/wm/Transition.java b/services/core/java/com/android/server/wm/Transition.java
index cb29e3f..ab38ed23 100644
--- a/services/core/java/com/android/server/wm/Transition.java
+++ b/services/core/java/com/android/server/wm/Transition.java
@@ -1561,6 +1561,10 @@
             if (info.mEndParent != null) {
                 change.setParent(info.mEndParent.mRemoteToken.toWindowContainerToken());
             }
+            if (info.mStartParent != null && info.mStartParent.mRemoteToken != null
+                    && target.getParent() != info.mStartParent) {
+                change.setLastParent(info.mStartParent.mRemoteToken.toWindowContainerToken());
+            }
             change.setMode(info.getTransitMode(target));
             change.setStartAbsBounds(info.mAbsoluteBounds);
             change.setFlags(info.getChangeFlags(target));
diff --git a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
index 4f03f54..e0e1d73 100644
--- a/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/RootWindowContainerTests.java
@@ -72,6 +72,7 @@
 import android.content.pm.ResolveInfo;
 import android.content.res.Configuration;
 import android.content.res.Resources;
+import android.graphics.Rect;
 import android.os.PowerManager;
 import android.os.UserHandle;
 import android.platform.test.annotations.Presubmit;
@@ -391,6 +392,33 @@
         assertEquals(WINDOWING_MODE_FULLSCREEN, fullscreenTask.getWindowingMode());
     }
 
+    @Test
+    public void testMovingEmbeddedActivityToPip() {
+        final Rect taskBounds = new Rect(0, 0, 800, 1000);
+        final Rect taskFragmentBounds = new Rect(0, 0, 400, 1000);
+        final Task task = mRootWindowContainer.getDefaultTaskDisplayArea().createRootTask(
+                WINDOWING_MODE_MULTI_WINDOW, ACTIVITY_TYPE_STANDARD, true /* onTop */);
+        task.setBounds(taskBounds);
+        assertEquals(taskBounds, task.getBounds());
+        final TaskFragment taskFragment = new TaskFragmentBuilder(mAtm)
+                .setParentTask(task)
+                .createActivityCount(2)
+                .setBounds(taskFragmentBounds)
+                .build();
+        assertEquals(taskFragmentBounds, taskFragment.getBounds());
+        final ActivityRecord topActivity = taskFragment.getTopMostActivity();
+
+        // Move the top activity to pinned root task.
+        mRootWindowContainer.moveActivityToPinnedRootTask(topActivity,
+                null /* launchIntoPipHostActivity */, "test");
+
+        final Task pinnedRootTask = task.getDisplayArea().getRootPinnedTask();
+
+        // Ensure the initial bounds of the PiP Task is the same as the TaskFragment.
+        ensureTaskPlacement(pinnedRootTask, topActivity);
+        assertEquals(taskFragmentBounds, pinnedRootTask.getBounds());
+    }
+
     private static void ensureTaskPlacement(Task task, ActivityRecord... activities) {
         final ArrayList<ActivityRecord> taskActivities = new ArrayList<>();
 
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 ec6a74c..29a514c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TransitionTests.java
@@ -1325,6 +1325,35 @@
     }
 
     @Test
+    public void testReparentChangeLastParent() {
+        final Transition transition = createTestTransition(TRANSIT_CHANGE);
+        final ArrayMap<WindowContainer, Transition.ChangeInfo> changes = transition.mChanges;
+        final ArraySet<WindowContainer> participants = transition.mParticipants;
+
+        // Reparent activity in transition.
+        final Task lastParent = createTask(mDisplayContent);
+        final Task newParent = createTask(mDisplayContent);
+        final ActivityRecord activity = createActivityRecord(lastParent);
+        activity.mVisibleRequested = true;
+        // Skip manipulate the SurfaceControl.
+        doNothing().when(activity).setDropInputMode(anyInt());
+        changes.put(activity, new Transition.ChangeInfo(activity));
+        activity.reparent(newParent, POSITION_TOP);
+        activity.mVisibleRequested = false;
+
+        participants.add(activity);
+        final ArrayList<WindowContainer> targets = Transition.calculateTargets(
+                participants, changes);
+        final TransitionInfo info = Transition.calculateTransitionInfo(
+                transition.mType, 0 /* flags */, targets, changes, mMockT);
+
+        // Change contains last parent info.
+        assertEquals(1, info.getChanges().size());
+        assertEquals(lastParent.mRemoteToken.toWindowContainerToken(),
+                info.getChanges().get(0).getLastParent());
+    }
+
+    @Test
     public void testIncludeEmbeddedActivityReparent() {
         final Transition transition = createTestTransition(TRANSIT_OPEN);
         final Task task = createTask(mDisplayContent);