2 files changed, 40 insertions, 8 deletions

diff --git a/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java b/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
index 8bbd15babf19..583953ce34af 100644
--- a/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
+++ b/packages/SystemUI/src/com/android/systemui/pip/PipBoundsHandler.java
@@ -177,7 +177,7 @@ public class PipBoundsHandler {
         }
         if (isValidPictureInPictureAspectRatio(mAspectRatio)) {
             transformBoundsToAspectRatio(normalBounds, mAspectRatio,
-                    false /* useCurrentMinEdgeSize */);
+                    false /* useCurrentMinEdgeSize */, false /* useCurrentSize */);
         }
         displayInfo.copyFrom(mDisplayInfo);
     }
@@ -278,7 +278,9 @@ public class PipBoundsHandler {
             destinationBounds = new Rect(bounds);
         }
         if (isValidPictureInPictureAspectRatio(aspectRatio)) {
-            transformBoundsToAspectRatio(destinationBounds, aspectRatio, useCurrentMinEdgeSize);
+            boolean useCurrentSize = bounds == null && mReentrySize != null;
+            transformBoundsToAspectRatio(destinationBounds, aspectRatio, useCurrentMinEdgeSize,
+                    useCurrentSize);
         }
         mAspectRatio = aspectRatio;
         return destinationBounds;
@@ -384,7 +386,8 @@ public class PipBoundsHandler {
      * @param stackBounds
      */
     public void transformBoundsToAspectRatio(Rect stackBounds) {
-        transformBoundsToAspectRatio(stackBounds, mAspectRatio, true);
+        transformBoundsToAspectRatio(stackBounds, mAspectRatio, true /* useCurrentMinEdgeSize */,
+                true /* useCurrentSize */);
     }
 
     /**
@@ -392,18 +395,16 @@ public class PipBoundsHandler {
      * specified aspect ratio.
      */
     private void transformBoundsToAspectRatio(Rect stackBounds, float aspectRatio,
-            boolean useCurrentMinEdgeSize) {
+            boolean useCurrentMinEdgeSize, boolean useCurrentSize) {
         // 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;
+        final int minEdgeSize = useCurrentMinEdgeSize ? mCurrentMinSize : mDefaultMinSize;
         final Size size;
-        if (useCurrentMinEdgeSize) {
-            minEdgeSize = mCurrentMinSize;
+        if (useCurrentMinEdgeSize || useCurrentSize) {
             size = mSnapAlgorithm.getSizeForAspectRatio(
                     new Size(stackBounds.width(), stackBounds.height()), aspectRatio, minEdgeSize);
         } else {
-            minEdgeSize = mDefaultMinSize;
             size = mSnapAlgorithm.getSizeForAspectRatio(aspectRatio, minEdgeSize,
                     mDisplayInfo.logicalWidth, mDisplayInfo.logicalHeight);
         }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/pip/PipBoundsHandlerTest.java b/packages/SystemUI/tests/src/com/android/systemui/pip/PipBoundsHandlerTest.java
index f404f0489e01..70c2bba040a0 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/pip/PipBoundsHandlerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/pip/PipBoundsHandlerTest.java
@@ -270,6 +270,21 @@ public class PipBoundsHandlerTest extends SysuiTestCase {
     }
 
     @Test
+    public void onSaveReentryBounds_restoreLastSize() {
+        final Rect oldSize = mPipBoundsHandler.getDestinationBounds(mTestComponentName1,
+                DEFAULT_ASPECT_RATIO, EMPTY_CURRENT_BOUNDS, EMPTY_MINIMAL_SIZE);
+
+        oldSize.scale(1.25f);
+        mPipBoundsHandler.onSaveReentryBounds(mTestComponentName1, oldSize);
+
+        final Rect newSize = mPipBoundsHandler.getDestinationBounds(mTestComponentName1,
+                DEFAULT_ASPECT_RATIO, EMPTY_CURRENT_BOUNDS, EMPTY_MINIMAL_SIZE);
+
+        assertEquals(oldSize.width(), newSize.width());
+        assertEquals(oldSize.height(), newSize.height());
+    }
+
+    @Test
     public void onResetReentryBounds_useDefaultBounds() {
         final Rect defaultBounds = mPipBoundsHandler.getDestinationBounds(mTestComponentName1,
                 DEFAULT_ASPECT_RATIO, EMPTY_CURRENT_BOUNDS, EMPTY_MINIMAL_SIZE);
@@ -299,6 +314,22 @@ public class PipBoundsHandlerTest extends SysuiTestCase {
         assertBoundsInclusionWithMargin("restoreLastPosition", newBounds, actualBounds);
     }
 
+    @Test
+    public void onSaveReentryBounds_componentMismatch_restoreLastSize() {
+        final Rect oldSize = mPipBoundsHandler.getDestinationBounds(mTestComponentName1,
+                DEFAULT_ASPECT_RATIO, EMPTY_CURRENT_BOUNDS, EMPTY_MINIMAL_SIZE);
+
+        oldSize.scale(1.25f);
+        mPipBoundsHandler.onSaveReentryBounds(mTestComponentName1, oldSize);
+
+        mPipBoundsHandler.onResetReentryBounds(mTestComponentName2);
+        final Rect newSize = mPipBoundsHandler.getDestinationBounds(mTestComponentName1,
+                DEFAULT_ASPECT_RATIO, EMPTY_CURRENT_BOUNDS, EMPTY_MINIMAL_SIZE);
+
+        assertEquals(oldSize.width(), newSize.width());
+        assertEquals(oldSize.height(), newSize.height());
+    }
+
     private void assertBoundsInclusionWithMargin(String from, Rect expected, Rect actual) {
         final Rect expectedWithMargin = new Rect(expected);
         expectedWithMargin.inset(-ROUNDING_ERROR_MARGIN, -ROUNDING_ERROR_MARGIN);