Persist PiP size to be restored on re-entry

Adds support for restoring the PiP size along with its snap fraction.
Modifies PipBoundsHandler logic to use both the previous size and
position (snap fraction) when calculaing the bounds for re-entry.
Additionally, some phone-specific logic is added to ensure that the
re-entry size that is restored is the non-expanded size (i.e. the
size before the bounds are animated to expand when pip is clicked).

Bug: 143641277
Test: enter pip, move it, touch to expand, restore. Then re-enter
pip. Pip should retain its size and position.
Test: atest PinnedStackTests
Change-Id: I21b7e4bf9360cd7da9faf2015c04d12d37c0241f
diff --git a/core/java/android/view/IPinnedStackListener.aidl b/core/java/android/view/IPinnedStackListener.aidl
index f4bee57..d01c933 100644
--- a/core/java/android/view/IPinnedStackListener.aidl
+++ b/core/java/android/view/IPinnedStackListener.aidl
@@ -67,20 +67,20 @@
     void onActionsChanged(in ParceledListSlice actions);
 
     /**
-     * Called by the window manager to notify the listener to save the reentry fraction,
+     * Called by the window manager to notify the listener to save the reentry fraction and size,
      * typically when an Activity leaves PiP (picture-in-picture) mode to fullscreen.
      * {@param componentName} represents the application component of PiP window
      * while {@param bounds} is the current PiP bounds used to calculate the
-     * reentry snap fraction.
+     * reentry snap fraction and size.
      */
-    void onSaveReentrySnapFraction(in ComponentName componentName, in Rect bounds);
+    void onSaveReentryBounds(in ComponentName componentName, in Rect bounds);
 
     /**
-     * Called by the window manager to notify the listener to reset saved reentry fraction,
+     * Called by the window manager to notify the listener to reset saved reentry fraction and size,
      * typically when an Activity enters PiP (picture-in-picture) mode from fullscreen.
      * {@param componentName} represents the application component of PiP window.
      */
-    void onResetReentrySnapFraction(in ComponentName componentName);
+    void onResetReentryBounds(in ComponentName componentName);
 
     /**
      * Called when the window manager has detected change on DisplayInfo,  or
diff --git a/core/java/com/android/internal/policy/PipSnapAlgorithm.java b/core/java/com/android/internal/policy/PipSnapAlgorithm.java
index 1afc67b..e3623c5 100644
--- a/core/java/com/android/internal/policy/PipSnapAlgorithm.java
+++ b/core/java/com/android/internal/policy/PipSnapAlgorithm.java
@@ -360,6 +360,28 @@
     }
 
     /**
+     * @return the adjusted size so that it conforms to the given aspectRatio, ensuring that the
+     * minimum edge is at least minEdgeSize.
+     */
+    public Size getSizeForAspectRatio(Size size, float aspectRatio, float minEdgeSize) {
+        final int smallestSize = Math.min(size.getWidth(), size.getHeight());
+        final int minSize = (int) Math.max(minEdgeSize, smallestSize);
+
+        final int width;
+        final int height;
+        if (aspectRatio <= 1) {
+            // Portrait, width is the minimum size.
+            width = minSize;
+            height = Math.round(width / aspectRatio);
+        } else {
+            // Landscape, height is the minimum size
+            height = minSize;
+            width = Math.round(height * aspectRatio);
+        }
+        return new Size(width, height);
+    }
+
+    /**
      * @return the closest point in {@param points} to the given {@param x} and {@param y}.
      */
     private Point findClosestPoint(int x, int y, Point[] points) {
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/system/PinnedStackListenerForwarder.java b/packages/SystemUI/shared/src/com/android/systemui/shared/system/PinnedStackListenerForwarder.java
index fe5a57a..8c0ffb8 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/system/PinnedStackListenerForwarder.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/system/PinnedStackListenerForwarder.java
@@ -83,16 +83,16 @@
     }
 
     @Override
-    public void onSaveReentrySnapFraction(ComponentName componentName, Rect bounds) {
+    public void onSaveReentryBounds(ComponentName componentName, Rect bounds) {
         for (PinnedStackListener listener : mListeners) {
-            listener.onSaveReentrySnapFraction(componentName, bounds);
+            listener.onSaveReentryBounds(componentName, bounds);
         }
     }
 
     @Override
-    public void onResetReentrySnapFraction(ComponentName componentName) {
+    public void onResetReentryBounds(ComponentName componentName) {
         for (PinnedStackListener listener : mListeners) {
-            listener.onResetReentrySnapFraction(componentName);
+            listener.onResetReentryBounds(componentName);
         }
     }
 
@@ -140,9 +140,9 @@
 
         public void onActionsChanged(ParceledListSlice actions) {}
 
-        public void onSaveReentrySnapFraction(ComponentName componentName, Rect bounds) {}
+        public void onSaveReentryBounds(ComponentName componentName, Rect bounds) {}
 
-        public void onResetReentrySnapFraction(ComponentName componentName) {}
+        public void onResetReentryBounds(ComponentName componentName) {}
 
         public void onDisplayInfoChanged(DisplayInfo displayInfo) {}
 
diff --git a/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java b/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
index 686e7db..f10274a 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
@@ -64,6 +64,7 @@
     private IPinnedStackController mPinnedStackController;
     private ComponentName mLastPipComponentName;
     private float mReentrySnapFraction = INVALID_SNAP_FRACTION;
+    private Size mReentrySize = null;
 
     private float mDefaultAspectRatio;
     private float mMinAspectRatio;
@@ -162,7 +163,7 @@
     public void onMovementBoundsChanged(Rect insetBounds, Rect normalBounds,
             Rect animatingBounds, DisplayInfo displayInfo) {
         getInsetBounds(insetBounds);
-        final Rect defaultBounds = getDefaultBounds(INVALID_SNAP_FRACTION);
+        final Rect defaultBounds = getDefaultBounds(INVALID_SNAP_FRACTION, null);
         normalBounds.set(defaultBounds);
         if (animatingBounds.isEmpty()) {
             animatingBounds.set(defaultBounds);
@@ -175,26 +176,28 @@
     }
 
     /**
-     * Responds to IPinnedStackListener on saving reentry snap fraction
+     * Responds to IPinnedStackListener on saving reentry snap fraction and size
      * for a given {@link ComponentName}.
      */
-    public void onSaveReentrySnapFraction(ComponentName componentName, Rect bounds) {
+    public void onSaveReentryBounds(ComponentName componentName, Rect bounds) {
         mReentrySnapFraction = getSnapFraction(bounds);
+        mReentrySize = new Size(bounds.width(), bounds.height());
         mLastPipComponentName = componentName;
     }
 
     /**
-     * Responds to IPinnedStackListener on resetting reentry snap fraction
+     * Responds to IPinnedStackListener on resetting reentry snap fraction and size
      * for a given {@link ComponentName}.
      */
-    public void onResetReentrySnapFraction(ComponentName componentName) {
+    public void onResetReentryBounds(ComponentName componentName) {
         if (componentName.equals(mLastPipComponentName)) {
-            onResetReentrySnapFractionUnchecked();
+            onResetReentryBoundsUnchecked();
         }
     }
 
-    private void onResetReentrySnapFractionUnchecked() {
+    private void onResetReentryBoundsUnchecked() {
         mReentrySnapFraction = INVALID_SNAP_FRACTION;
+        mReentrySize = null;
         mLastPipComponentName = null;
     }
 
@@ -233,7 +236,7 @@
     public void onPrepareAnimation(Rect sourceRectHint, float aspectRatio, Rect bounds) {
         final Rect destinationBounds;
         if (bounds == null) {
-            destinationBounds = getDefaultBounds(mReentrySnapFraction);
+            destinationBounds = getDefaultBounds(mReentrySnapFraction, mReentrySize);
         } else {
             destinationBounds = new Rect(bounds);
         }
@@ -245,7 +248,7 @@
             return;
         }
         mAspectRatio = aspectRatio;
-        onResetReentrySnapFractionUnchecked();
+        onResetReentryBoundsUnchecked();
         try {
             mPinnedStackController.startAnimation(destinationBounds, sourceRectHint,
                     -1 /* animationDuration */);
@@ -269,13 +272,14 @@
      */
     private void transformBoundsToAspectRatio(Rect stackBounds, float aspectRatio,
             boolean useCurrentMinEdgeSize) {
-        // Save the snap fraction, calculate the aspect ratio based on screen size
+
+        // Save the snap fraction and adjust the size based on the new aspect ratio.
         final float snapFraction = mSnapAlgorithm.getSnapFraction(stackBounds,
                 getMovementBounds(stackBounds));
-
         final int minEdgeSize = useCurrentMinEdgeSize ? mCurrentMinSize : mDefaultMinSize;
-        final Size size = mSnapAlgorithm.getSizeForAspectRatio(aspectRatio, minEdgeSize,
-                mDisplayInfo.logicalWidth, mDisplayInfo.logicalHeight);
+        final Size size = mSnapAlgorithm.getSizeForAspectRatio(
+                new Size(stackBounds.width(), stackBounds.height()), aspectRatio, minEdgeSize);
+
         final int left = (int) (stackBounds.centerX() - size.getWidth() / 2f);
         final int top = (int) (stackBounds.centerY() - size.getHeight() / 2f);
         stackBounds.set(left, top, left + size.getWidth(), top + size.getHeight());
@@ -286,21 +290,20 @@
     }
 
     /**
-     * @return the default bounds to show the PIP, if a {@param snapFraction} is provided, then it
-     * will apply the default bounds to the provided snap fraction.
+     * @return the default bounds to show the PIP, if a {@param snapFraction} and {@param size} are
+     * provided, then it will apply the default bounds to the provided snap fraction and size.
      */
-    private Rect getDefaultBounds(float snapFraction) {
-        final Rect insetBounds = new Rect();
-        getInsetBounds(insetBounds);
-
+    private Rect getDefaultBounds(float snapFraction, Size size) {
         final Rect defaultBounds = new Rect();
-        final Size size = mSnapAlgorithm.getSizeForAspectRatio(mDefaultAspectRatio,
-                mDefaultMinSize, mDisplayInfo.logicalWidth, mDisplayInfo.logicalHeight);
-        if (snapFraction != INVALID_SNAP_FRACTION) {
+        if (snapFraction != INVALID_SNAP_FRACTION && size != null) {
             defaultBounds.set(0, 0, size.getWidth(), size.getHeight());
             final Rect movementBounds = getMovementBounds(defaultBounds);
             mSnapAlgorithm.applySnapFraction(defaultBounds, movementBounds, snapFraction);
         } else {
+            final Rect insetBounds = new Rect();
+            getInsetBounds(insetBounds);
+            size = mSnapAlgorithm.getSizeForAspectRatio(mDefaultAspectRatio,
+                    mDefaultMinSize, mDisplayInfo.logicalWidth, mDisplayInfo.logicalHeight);
             Gravity.apply(mDefaultStackGravity, size.getWidth(), size.getHeight(), insetBounds,
                     0, Math.max(mIsImeShowing ? mImeHeight : 0,
                             mIsShelfShowing ? mShelfHeight : 0),
@@ -364,11 +367,19 @@
      * @return the default snap fraction to apply instead of the default gravity when calculating
      *         the default stack bounds when first entering PiP.
      */
-    private float getSnapFraction(Rect stackBounds) {
+    public float getSnapFraction(Rect stackBounds) {
         return mSnapAlgorithm.getSnapFraction(stackBounds, getMovementBounds(stackBounds));
     }
 
     /**
+     * Applies the given snap fraction to the given stack bounds.
+     */
+    public void applySnapFraction(Rect stackBounds, float snapFraction) {
+        final Rect movementBounds = getMovementBounds(stackBounds);
+        mSnapAlgorithm.applySnapFraction(stackBounds, movementBounds, snapFraction);
+    }
+
+    /**
      * @return the pixels for a given dp value.
      */
     private int dpToPx(float dpValue, DisplayMetrics dm) {
diff --git a/packages/SystemUI/src/com/android/systemui/pip/phone/PipManager.java b/packages/SystemUI/src/com/android/systemui/pip/phone/PipManager.java
index c33b8d9..a4707cf 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/phone/PipManager.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/phone/PipManager.java
@@ -65,6 +65,7 @@
     private final DisplayInfo mTmpDisplayInfo = new DisplayInfo();
     private final Rect mTmpInsetBounds = new Rect();
     private final Rect mTmpNormalBounds = new Rect();
+    private final Rect mReentryBounds = new Rect();
 
     private PipBoundsHandler mPipBoundsHandler;
     private InputConsumerController mInputConsumerController;
@@ -164,13 +165,25 @@
         }
 
         @Override
-        public void onSaveReentrySnapFraction(ComponentName componentName, Rect bounds) {
-            mHandler.post(() -> mPipBoundsHandler.onSaveReentrySnapFraction(componentName, bounds));
+        public void onSaveReentryBounds(ComponentName componentName, Rect bounds) {
+            mHandler.post(() -> {
+                // On phones, the expansion animation that happens on pip tap before restoring
+                // to fullscreen makes it so that the bounds received here are the expanded
+                // bounds. We want to restore to the unexpanded bounds when re-entering pip,
+                // so we save the bounds before expansion (normal) instead of the current
+                // bounds.
+                mReentryBounds.set(mTouchHandler.getNormalBounds());
+                // Apply the snap fraction of the current bounds to the normal bounds.
+                float snapFraction = mPipBoundsHandler.getSnapFraction(bounds);
+                mPipBoundsHandler.applySnapFraction(mReentryBounds, snapFraction);
+                // Save reentry bounds (normal non-expand bounds with current position applied).
+                mPipBoundsHandler.onSaveReentryBounds(componentName, mReentryBounds);
+            });
         }
 
         @Override
-        public void onResetReentrySnapFraction(ComponentName componentName) {
-            mHandler.post(() -> mPipBoundsHandler.onResetReentrySnapFraction(componentName));
+        public void onResetReentryBounds(ComponentName componentName) {
+            mHandler.post(() -> mPipBoundsHandler.onResetReentryBounds(componentName));
         }
 
         @Override
diff --git a/packages/SystemUI/src/com/android/systemui/pip/phone/PipTouchHandler.java b/packages/SystemUI/src/com/android/systemui/pip/phone/PipTouchHandler.java
index f59b372..2e90a3e 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/phone/PipTouchHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/phone/PipTouchHandler.java
@@ -592,6 +592,13 @@
     }
 
     /**
+     * @return the unexpanded bounds.
+     */
+    public Rect getNormalBounds() {
+        return mNormalBounds;
+    }
+
+    /**
      * Gesture controlling normal movement of the PIP.
      */
     private PipTouchGesture mDefaultMovementGesture = new PipTouchGesture() {
diff --git a/services/core/java/com/android/server/wm/ActivityRecord.java b/services/core/java/com/android/server/wm/ActivityRecord.java
index fcdb3b0..5bc3d22 100644
--- a/services/core/java/com/android/server/wm/ActivityRecord.java
+++ b/services/core/java/com/android/server/wm/ActivityRecord.java
@@ -3049,7 +3049,7 @@
         }
 
         // Reset the last saved PiP snap fraction on removal.
-        mDisplayContent.mPinnedStackControllerLocked.resetReentrySnapFraction(mActivityComponent);
+        mDisplayContent.mPinnedStackControllerLocked.resetReentryBounds(mActivityComponent);
 
         mRemovingFromDisplay = false;
     }
@@ -4300,7 +4300,7 @@
         ProtoLog.v(WM_DEBUG_ADD_REMOVE, "notifyAppStopped: %s", this);
         mAppStopped = true;
         // Reset the last saved PiP snap fraction on app stop.
-        mDisplayContent.mPinnedStackControllerLocked.resetReentrySnapFraction(mActivityComponent);
+        mDisplayContent.mPinnedStackControllerLocked.resetReentryBounds(mActivityComponent);
         destroySurfaces();
         // Remove any starting window that was added for this app if they are still around.
         removeStartingWindow();
@@ -6537,7 +6537,7 @@
             stackBounds = mTmpRect;
             pinnedStack.getBounds(stackBounds);
         }
-        mDisplayContent.mPinnedStackControllerLocked.saveReentrySnapFraction(
+        mDisplayContent.mPinnedStackControllerLocked.saveReentryBounds(
                 mActivityComponent, stackBounds);
     }
 
diff --git a/services/core/java/com/android/server/wm/PinnedStackController.java b/services/core/java/com/android/server/wm/PinnedStackController.java
index a5b1fda..0853f1f 100644
--- a/services/core/java/com/android/server/wm/PinnedStackController.java
+++ b/services/core/java/com/android/server/wm/PinnedStackController.java
@@ -236,10 +236,10 @@
     /**
      * Saves the current snap fraction for re-entry of the current activity into PiP.
      */
-    void saveReentrySnapFraction(final ComponentName componentName, final Rect stackBounds) {
+    void saveReentryBounds(final ComponentName componentName, final Rect stackBounds) {
         if (mPinnedStackListener == null) return;
         try {
-            mPinnedStackListener.onSaveReentrySnapFraction(componentName, stackBounds);
+            mPinnedStackListener.onSaveReentryBounds(componentName, stackBounds);
         } catch (RemoteException e) {
             Slog.e(TAG_WM, "Error delivering save reentry fraction event.", e);
         }
@@ -248,10 +248,10 @@
     /**
      * Resets the last saved snap fraction so that the default bounds will be returned.
      */
-    void resetReentrySnapFraction(ComponentName componentName) {
+    void resetReentryBounds(ComponentName componentName) {
         if (mPinnedStackListener == null) return;
         try {
-            mPinnedStackListener.onResetReentrySnapFraction(componentName);
+            mPinnedStackListener.onResetReentryBounds(componentName);
         } catch (RemoteException e) {
             Slog.e(TAG_WM, "Error delivering reset reentry fraction event.", e);
         }