Store original task width and height instead of scale

Store the original task snapshot size instead of the scale from which
the bitmap was saved. This simplifies the logic around restoring and
saving from the proto, as both the reduced scale and full scale
snapshots make use and share the same state.

Also remove scale from TaskSnapshot, and remove and reducedScale from
TaskSnapshot.Builder.

Test: TaskSnapshotCacheTest
Test: TaskSnapshotControllerTest
Test: TaskSnapshotPersisterLoaderTest
Test: TaskSnapshotSurfaceTest
Bug: 148491788
Bug: 148617404
Bug: 142063079
Change-Id: I1dccaba87c3d8b95bf4156f41f9fd5d40019f675
diff --git a/core/java/android/app/ActivityManager.java b/core/java/android/app/ActivityManager.java
index 82fdb90..25d3043 100644
--- a/core/java/android/app/ActivityManager.java
+++ b/core/java/android/app/ActivityManager.java
@@ -2012,6 +2012,8 @@
         /** See {@link android.view.Surface.Rotation} */
         @Surface.Rotation
         private int mRotation;
+        /** The size of the snapshot before scaling */
+        private final Point mTaskSize;
         private final Rect mContentInsets;
         // Whether this snapshot is a down-sampled version of the full resolution, used mainly for
         // low-ram devices
@@ -2020,7 +2022,6 @@
         // the task having a secure window or having previews disabled
         private final boolean mIsRealSnapshot;
         private final int mWindowingMode;
-        private final float mScale;
         private final int mSystemUiVisibility;
         private final boolean mIsTranslucent;
         // Must be one of the named color spaces, otherwise, always use SRGB color space.
@@ -2028,9 +2029,9 @@
 
         public TaskSnapshot(long id,
                 @NonNull ComponentName topActivityComponent, GraphicBuffer snapshot,
-                @NonNull ColorSpace colorSpace, int orientation, int rotation, Rect contentInsets,
-                boolean isLowResolution, float scale, boolean isRealSnapshot, int windowingMode,
-                int systemUiVisibility, boolean isTranslucent) {
+                @NonNull ColorSpace colorSpace, int orientation, int rotation, Point taskSize,
+                Rect contentInsets, boolean isLowResolution, boolean isRealSnapshot,
+                int windowingMode, int systemUiVisibility, boolean isTranslucent) {
             mId = id;
             mTopActivityComponent = topActivityComponent;
             mSnapshot = snapshot;
@@ -2038,9 +2039,9 @@
                     ? ColorSpace.get(ColorSpace.Named.SRGB) : colorSpace;
             mOrientation = orientation;
             mRotation = rotation;
+            mTaskSize = new Point(taskSize);
             mContentInsets = new Rect(contentInsets);
             mIsLowResolution = isLowResolution;
-            mScale = scale;
             mIsRealSnapshot = isRealSnapshot;
             mWindowingMode = windowingMode;
             mSystemUiVisibility = systemUiVisibility;
@@ -2057,9 +2058,9 @@
                     : ColorSpace.get(ColorSpace.Named.SRGB);
             mOrientation = source.readInt();
             mRotation = source.readInt();
+            mTaskSize = source.readParcelable(null /* classLoader */);
             mContentInsets = source.readParcelable(null /* classLoader */);
             mIsLowResolution = source.readBoolean();
-            mScale = source.readFloat();
             mIsRealSnapshot = source.readBoolean();
             mWindowingMode = source.readInt();
             mSystemUiVisibility = source.readInt();
@@ -2111,6 +2112,14 @@
         }
 
         /**
+         * @return The size of the task at the point this snapshot was taken.
+         */
+        @UnsupportedAppUsage
+        public Point getTaskSize() {
+            return mTaskSize;
+        }
+
+        /**
          * @return The system/content insets on the snapshot. These can be clipped off in order to
          *         remove any areas behind system bars in the snapshot.
          */
@@ -2159,14 +2168,6 @@
             return mSystemUiVisibility;
         }
 
-        /**
-         * @return The scale this snapshot was taken in.
-         */
-        @UnsupportedAppUsage
-        public float getScale() {
-            return mScale;
-        }
-
         @Override
         public int describeContents() {
             return 0;
@@ -2180,9 +2181,9 @@
             dest.writeInt(mColorSpace.getId());
             dest.writeInt(mOrientation);
             dest.writeInt(mRotation);
+            dest.writeParcelable(mTaskSize, 0);
             dest.writeParcelable(mContentInsets, 0);
             dest.writeBoolean(mIsLowResolution);
-            dest.writeFloat(mScale);
             dest.writeBoolean(mIsRealSnapshot);
             dest.writeInt(mWindowingMode);
             dest.writeInt(mSystemUiVisibility);
@@ -2200,9 +2201,11 @@
                     + " mColorSpace=" + mColorSpace.toString()
                     + " mOrientation=" + mOrientation
                     + " mRotation=" + mRotation
+                    + " mTaskSize=" + mTaskSize.toString()
                     + " mContentInsets=" + mContentInsets.toShortString()
-                    + " mIsLowResolution=" + mIsLowResolution + " mScale=" + mScale
-                    + " mIsRealSnapshot=" + mIsRealSnapshot + " mWindowingMode=" + mWindowingMode
+                    + " mIsLowResolution=" + mIsLowResolution
+                    + " mIsRealSnapshot=" + mIsRealSnapshot
+                    + " mWindowingMode=" + mWindowingMode
                     + " mSystemUiVisibility=" + mSystemUiVisibility
                     + " mIsTranslucent=" + mIsTranslucent;
         }
@@ -2224,9 +2227,9 @@
             private ColorSpace mColorSpace;
             private int mOrientation;
             private int mRotation;
+            private Point mTaskSize;
             private Rect mContentInsets;
             private boolean mIsLowResolution;
-            private float mScaleFraction;
             private boolean mIsRealSnapshot;
             private int mWindowingMode;
             private int mSystemUiVisibility;
@@ -2263,28 +2266,34 @@
                 return this;
             }
 
+            /**
+             * Sets the original size of the task
+             */
+            public Builder setTaskSize(Point size) {
+                mTaskSize = size;
+                return this;
+            }
+
             public Builder setContentInsets(Rect contentInsets) {
                 mContentInsets = contentInsets;
                 return this;
             }
 
             /**
-             * Set to true if this is a low-resolution snapshot stored in *_reduced.jpg.
+             * Returns {@code true} if this is meant to be a low-resolution
+             */
+            public boolean isLowResolution() {
+                return mIsLowResolution;
+            }
+
+            /**
+             * Set to {@code true} if this is a low-resolution snapshot stored in *_reduced.jpg.
              */
             public Builder setIsLowResolution(boolean isLowResolution) {
                 mIsLowResolution = isLowResolution;
                 return this;
             }
 
-            public float getScaleFraction() {
-                return mScaleFraction;
-            }
-
-            public Builder setScaleFraction(float scaleFraction) {
-                mScaleFraction = scaleFraction;
-                return this;
-            }
-
             public Builder setIsRealSnapshot(boolean realSnapshot) {
                 mIsRealSnapshot = realSnapshot;
                 return this;
@@ -2322,9 +2331,9 @@
                         mColorSpace,
                         mOrientation,
                         mRotation,
+                        mTaskSize,
                         mContentInsets,
                         mIsLowResolution,
-                        mScaleFraction,
                         mIsRealSnapshot,
                         mWindowingMode,
                         mSystemUiVisibility,
diff --git a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/ThumbnailData.java b/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/ThumbnailData.java
index 4474a49..eca6ebf 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/ThumbnailData.java
+++ b/packages/SystemUI/shared/src/com/android/systemui/shared/recents/model/ThumbnailData.java
@@ -62,7 +62,9 @@
         orientation = snapshot.getOrientation();
         rotation = snapshot.getRotation();
         reducedResolution = snapshot.isLowResolution();
-        scale = snapshot.getScale();
+        // TODO(b/149579527): Pass task size instead of computing scale.
+        // Assume width and height were scaled the same; compute scale only for width
+        scale = (float) thumbnail.getWidth() / snapshot.getTaskSize().x;
         isRealSnapshot = snapshot.isRealSnapshot();
         isTranslucent = snapshot.isTranslucent();
         windowingMode = snapshot.getWindowingMode();
diff --git a/proto/src/task_snapshot.proto b/proto/src/task_snapshot.proto
index 789019c..2006fb3 100644
--- a/proto/src/task_snapshot.proto
+++ b/proto/src/task_snapshot.proto
@@ -32,7 +32,12 @@
      int32 system_ui_visibility = 8;
      bool is_translucent = 9;
      string top_activity_component = 10;
-     float scale = 11;
+     // deprecated because original width and height are stored now instead of the scale.
+     float legacy_scale = 11 [deprecated=true];
      int64 id = 12;
      int32 rotation = 13;
+     // The task width when the snapshot was taken
+     int32 task_width = 14;
+     // The task height when the snapshot was taken
+     int32 task_height = 15;
  }
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 5c73f92..2092c01 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -31,6 +31,7 @@
 import android.graphics.Bitmap;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
+import android.graphics.Point;
 import android.graphics.RecordingCanvas;
 import android.graphics.Rect;
 import android.graphics.RenderNode;
@@ -89,14 +90,6 @@
     @VisibleForTesting
     static final int SNAPSHOT_MODE_NONE = 2;
 
-    /**
-     * Constant for <code>scaleFactor</code> when calling {@link #snapshotTask} which is
-     * interpreted as using the most appropriate scale ratio for the system.
-     * This may yield a smaller ratio on low memory devices.
-     */
-    @VisibleForTesting
-    static final float SNAPSHOT_SCALE_AUTO = -1f;
-
     private final WindowManagerService mService;
 
     private final TaskSnapshotCache mCache;
@@ -273,8 +266,6 @@
      * information from the task and populates the builder.
      *
      * @param task the task to capture
-     * @param scaleFraction the scale fraction between 0-1.0, or {@link #SNAPSHOT_SCALE_AUTO}
-     *                      to automatically select
      * @param pixelFormat the desired pixel format, or {@link PixelFormat#UNKNOWN} to
      *                    automatically select
      * @param builder the snapshot builder to populate
@@ -282,8 +273,7 @@
      * @return true if the state of the task is ok to proceed
      */
     @VisibleForTesting
-    boolean prepareTaskSnapshot(Task task, float scaleFraction, int pixelFormat,
-            TaskSnapshot.Builder builder) {
+    boolean prepareTaskSnapshot(Task task, int pixelFormat, TaskSnapshot.Builder builder) {
         if (!mService.mPolicy.isScreenOn()) {
             if (DEBUG_SCREENSHOT) {
                 Slog.i(TAG_WM, "Attempted to take screenshot while display was off.");
@@ -310,21 +300,12 @@
             return false;
         }
 
+        final boolean isLowRamDevice = ActivityManager.isLowRamDeviceStatic();
+
         builder.setIsRealSnapshot(true);
         builder.setId(System.currentTimeMillis());
         builder.setContentInsets(getInsets(mainWindow));
-
-        final boolean isLowRamDevice = ActivityManager.isLowRamDeviceStatic();
-
-        if (scaleFraction == SNAPSHOT_SCALE_AUTO) {
-            builder.setScaleFraction(isLowRamDevice
-                    ? mPersister.getLowResScale()
-                    : mHighResTaskSnapshotScale);
-            builder.setIsLowResolution(isLowRamDevice);
-        } else {
-            builder.setScaleFraction(scaleFraction);
-            builder.setIsLowResolution(scaleFraction < 1.0f);
-        }
+        builder.setIsLowResolution(isLowRamDevice);
 
         final boolean isWindowTranslucent = mainWindow.getAttrs().format != PixelFormat.OPAQUE;
         final boolean isShowWallpaper = (mainWindow.getAttrs().flags & FLAG_SHOW_WALLPAPER) != 0;
@@ -351,13 +332,25 @@
 
     @Nullable
     SurfaceControl.ScreenshotGraphicBuffer createTaskSnapshot(@NonNull Task task,
-            float scaleFraction) {
-        return createTaskSnapshot(task, scaleFraction, PixelFormat.RGBA_8888);
+            TaskSnapshot.Builder builder) {
+        Point taskSize = new Point();
+        float scale = builder.isLowResolution()
+                ? mPersister.getLowResScale() : mHighResTaskSnapshotScale;
+        final SurfaceControl.ScreenshotGraphicBuffer taskSnapshot = createTaskSnapshot(task, scale,
+                builder.getPixelFormat(), taskSize);
+        builder.setTaskSize(taskSize);
+        return taskSnapshot;
     }
 
     @Nullable
     SurfaceControl.ScreenshotGraphicBuffer createTaskSnapshot(@NonNull Task task,
-            float scaleFraction, int pixelFormat) {
+            float scaleFraction) {
+        return createTaskSnapshot(task, scaleFraction, PixelFormat.RGBA_8888, null);
+    }
+
+    @Nullable
+    SurfaceControl.ScreenshotGraphicBuffer createTaskSnapshot(@NonNull Task task,
+            float scaleFraction, int pixelFormat, Point outTaskSize) {
         if (task.getSurfaceControl() == null) {
             if (DEBUG_SCREENSHOT) {
                 Slog.w(TAG_WM, "Failed to take screenshot. No surface control for " + task);
@@ -369,6 +362,10 @@
         final SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer =
                 SurfaceControl.captureLayers(
                         task.getSurfaceControl(), mTmpRect, scaleFraction, pixelFormat);
+        if (outTaskSize != null) {
+            outTaskSize.x = mTmpRect.width();
+            outTaskSize.y = mTmpRect.height();
+        }
         final GraphicBuffer buffer = screenshotBuffer != null ? screenshotBuffer.getGraphicBuffer()
                 : null;
         if (buffer == null || buffer.getWidth() <= 1 || buffer.getHeight() <= 1) {
@@ -379,21 +376,20 @@
 
     @Nullable
     TaskSnapshot snapshotTask(Task task) {
-        return snapshotTask(task, SNAPSHOT_SCALE_AUTO, PixelFormat.UNKNOWN);
+        return snapshotTask(task, PixelFormat.UNKNOWN);
     }
 
     @Nullable
-    TaskSnapshot snapshotTask(Task task, float scaleFraction, int pixelFormat) {
+    TaskSnapshot snapshotTask(Task task, int pixelFormat) {
         TaskSnapshot.Builder builder = new TaskSnapshot.Builder();
 
-        if (!prepareTaskSnapshot(task, scaleFraction, pixelFormat, builder)) {
+        if (!prepareTaskSnapshot(task, pixelFormat, builder)) {
             // Failed some pre-req. Has been logged.
             return null;
         }
 
         final SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer =
-                createTaskSnapshot(task, builder.getScaleFraction(),
-                builder.getPixelFormat());
+                createTaskSnapshot(task, builder);
 
         if (screenshotBuffer == null) {
             // Failed to acquire image. Has been logged.
@@ -472,8 +468,10 @@
         final SystemBarBackgroundPainter decorPainter = new SystemBarBackgroundPainter(attrs.flags,
                 attrs.privateFlags, attrs.systemUiVisibility, task.getTaskDescription(),
                 mHighResTaskSnapshotScale, mainWindow.getRequestedInsetsState());
-        final int width = (int) (task.getBounds().width() * mHighResTaskSnapshotScale);
-        final int height = (int) (task.getBounds().height() * mHighResTaskSnapshotScale);
+        final int taskWidth = task.getBounds().width();
+        final int taskHeight = task.getBounds().height();
+        final int width = (int) (taskWidth * mHighResTaskSnapshotScale);
+        final int height = (int) (taskHeight * mHighResTaskSnapshotScale);
 
         final RenderNode node = RenderNode.create("TaskSnapshotController", null);
         node.setLeftTopRightBottom(0, 0, width, height);
@@ -494,9 +492,9 @@
                 System.currentTimeMillis() /* id */,
                 topChild.mActivityComponent, hwBitmap.createGraphicBufferHandle(),
                 hwBitmap.getColorSpace(), mainWindow.getConfiguration().orientation,
-                mainWindow.getWindowConfiguration().getRotation(),
+                mainWindow.getWindowConfiguration().getRotation(), new Point(taskWidth, taskHeight),
                 getInsets(mainWindow), ActivityManager.isLowRamDeviceStatic() /* isLowResolution */,
-                mHighResTaskSnapshotScale, false /* isRealSnapshot */, task.getWindowingMode(),
+                false /* isRealSnapshot */, task.getWindowingMode(),
                 getSystemUiVisibility(task), false);
     }
 
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
index 01f3427..ed8dd90 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
@@ -19,6 +19,7 @@
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WITH_CLASS_NAME;
 import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 
+import android.app.ActivityManager;
 import android.app.ActivityManager.TaskSnapshot;
 import android.content.ComponentName;
 import android.graphics.Bitmap;
@@ -26,6 +27,7 @@
 import android.graphics.BitmapFactory;
 import android.graphics.BitmapFactory.Options;
 import android.graphics.GraphicBuffer;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.util.Slog;
 
@@ -46,6 +48,9 @@
 
     private static final String TAG = TAG_WITH_CLASS_NAME ? "TaskSnapshotLoader" : TAG_WM;
 
+    private static final float LEGACY_REDUCED_SCALE =
+            ActivityManager.isLowRamDeviceStatic() ? 0.6f : 0.5f;
+
     private final TaskSnapshotPersister mPersister;
 
     TaskSnapshotLoader(TaskSnapshotPersister persister) {
@@ -99,13 +104,25 @@
 
             final ComponentName topActivityComponent = ComponentName.unflattenFromString(
                     proto.topActivityComponent);
+
             // For legacy snapshots, restore the scale based on the reduced resolution state
-            final float legacyScale = isLowResolution ? mPersister.getLowResScale() : 1f;
-            final float scale = Float.compare(proto.scale, 0f) != 0 ? proto.scale : legacyScale;
+            Point taskSize;
+            if (proto.taskWidth == 0) {
+                // For legacy snapshots, restore the scale based on the reduced resolution state
+                final float preQLegacyScale = isLowResolution ? LEGACY_REDUCED_SCALE : 1f;
+                final float scale = Float.compare(proto.legacyScale, 0f) != 0
+                        ? proto.legacyScale : preQLegacyScale;
+                int taskWidth = (int) ((float) hwBitmap.getWidth() / scale);
+                int taskHeight = (int) ((float) hwBitmap.getHeight() / scale);
+                taskSize = new Point(taskWidth, taskHeight);
+            } else {
+                taskSize = new Point(proto.taskWidth, proto.taskHeight);
+            }
+
             return new TaskSnapshot(proto.id, topActivityComponent, buffer, hwBitmap.getColorSpace(),
-                    proto.orientation, proto.rotation,
+                    proto.orientation, proto.rotation, taskSize,
                     new Rect(proto.insetLeft, proto.insetTop, proto.insetRight, proto.insetBottom),
-                    isLowResolution, scale, proto.isRealSnapshot, proto.windowingMode,
+                    isLowResolution, proto.isRealSnapshot, proto.windowingMode,
                     proto.systemUiVisibility, proto.isTranslucent);
         } catch (IOException e) {
             Slog.w(TAG, "Unable to load task snapshot data for taskId=" + taskId);
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
index 31212b8..a5693b1 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotPersister.java
@@ -343,6 +343,8 @@
             final TaskSnapshotProto proto = new TaskSnapshotProto();
             proto.orientation = mSnapshot.getOrientation();
             proto.rotation = mSnapshot.getRotation();
+            proto.taskWidth = mSnapshot.getTaskSize().x;
+            proto.taskHeight = mSnapshot.getTaskSize().y;
             proto.insetLeft = mSnapshot.getContentInsets().left;
             proto.insetTop = mSnapshot.getContentInsets().top;
             proto.insetRight = mSnapshot.getContentInsets().right;
@@ -352,7 +354,6 @@
             proto.systemUiVisibility = mSnapshot.getSystemUiVisibility();
             proto.isTranslucent = mSnapshot.isTranslucent();
             proto.topActivityComponent = mSnapshot.getTopActivityComponent().flattenToString();
-            proto.scale = mSnapshot.getScale();
             proto.id = mSnapshot.getId();
             final byte[] bytes = TaskSnapshotProto.toByteArray(proto);
             final File file = getProtoFile(mTaskId, mUserId);
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
index f4e4245..eb005e0 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotSurface.java
@@ -36,6 +36,7 @@
 import static android.view.WindowManager.LayoutParams.FLAG_WATCH_OUTSIDE_TOUCH;
 import static android.view.WindowManager.LayoutParams.PRIVATE_FLAG_FORCE_DRAW_BAR_BACKGROUNDS;
 import static android.view.WindowManager.LayoutParams.TYPE_APPLICATION_STARTING;
+
 import static com.android.internal.policy.DecorView.NAVIGATION_BAR_COLOR_VIEW_ATTRIBUTES;
 import static com.android.internal.policy.DecorView.STATUS_BAR_COLOR_VIEW_ATTRIBUTES;
 import static com.android.internal.policy.DecorView.getColorViewLeftInset;
@@ -53,9 +54,11 @@
 import android.graphics.Canvas;
 import android.graphics.Color;
 import android.graphics.GraphicBuffer;
+import android.graphics.Matrix;
 import android.graphics.Paint;
 import android.graphics.Point;
 import android.graphics.Rect;
+import android.graphics.RectF;
 import android.os.Handler;
 import android.os.Looper;
 import android.os.Message;
@@ -131,6 +134,8 @@
     private final Rect mContentInsets = new Rect();
     private final Rect mFrame = new Rect();
     private TaskSnapshot mSnapshot;
+    private final RectF mTmpSnapshotSize = new RectF();
+    private final RectF mTmpDstFrame = new RectF();
     private final CharSequence mTitle;
     private boolean mHasDrawn;
     private long mShownTime;
@@ -141,6 +146,8 @@
     @VisibleForTesting final SystemBarBackgroundPainter mSystemBarBackgroundPainter;
     private final int mOrientationOnCreation;
     private final SurfaceControl.Transaction mTransaction;
+    private final Matrix mSnapshotMatrix = new Matrix();
+    private final float[] mTmpFloat9 = new float[9];
 
     static TaskSnapshotSurface create(WindowManagerService service, ActivityRecord activity,
             TaskSnapshot snapshot) {
@@ -365,13 +372,17 @@
             frame = calculateSnapshotFrame(crop);
             mTransaction.setWindowCrop(mChildSurfaceControl, crop);
             mTransaction.setPosition(mChildSurfaceControl, frame.left, frame.top);
+            mTmpDstFrame.set(frame);
         } else {
             frame = null;
+            mTmpDstFrame.set(mFrame);
         }
 
         // Scale the mismatch dimensions to fill the task bounds
-        final float scale = 1 / mSnapshot.getScale();
-        mTransaction.setMatrix(mChildSurfaceControl, scale, 0, 0, scale);
+        mTmpSnapshotSize.set(0, 0, buffer.getWidth(), buffer.getHeight());
+        mSnapshotMatrix.setRectToRect(mTmpSnapshotSize, mTmpDstFrame, Matrix.ScaleToFit.FILL);
+        mTransaction.setMatrix(mChildSurfaceControl, mSnapshotMatrix, mTmpFloat9);
+
         mTransaction.apply();
         surface.attachAndQueueBufferWithColorSpace(buffer, mSnapshot.getColorSpace());
         surface.release();
@@ -395,13 +406,17 @@
         rect.set(0, 0, mSnapshot.getSnapshot().getWidth(), mSnapshot.getSnapshot().getHeight());
         final Rect insets = mSnapshot.getContentInsets();
 
+        final float scaleX = (float) mSnapshot.getSnapshot().getWidth() / mSnapshot.getTaskSize().x;
+        final float scaleY =
+                (float) mSnapshot.getSnapshot().getHeight() / mSnapshot.getTaskSize().y;
+
         // Let's remove all system decorations except the status bar, but only if the task is at the
         // very top of the screen.
         final boolean isTop = mTaskBounds.top == 0 && mFrame.top == 0;
-        rect.inset((int) (insets.left * mSnapshot.getScale()),
-                isTop ? 0 : (int) (insets.top * mSnapshot.getScale()),
-                (int) (insets.right * mSnapshot.getScale()),
-                (int) (insets.bottom * mSnapshot.getScale()));
+        rect.inset((int) (insets.left * scaleX),
+                isTop ? 0 : (int) (insets.top * scaleY),
+                (int) (insets.right * scaleX),
+                (int) (insets.bottom * scaleY));
         return rect;
     }
 
@@ -412,14 +427,20 @@
      */
     @VisibleForTesting
     Rect calculateSnapshotFrame(Rect crop) {
-        final Rect frame = new Rect(crop);
-        final float scale = mSnapshot.getScale();
+        final float scaleX = (float) mSnapshot.getSnapshot().getWidth() / mSnapshot.getTaskSize().x;
+        final float scaleY =
+                (float) mSnapshot.getSnapshot().getHeight() / mSnapshot.getTaskSize().y;
 
         // Rescale the frame from snapshot to window coordinate space
-        frame.scale(1 / scale);
+        final Rect frame = new Rect(
+                (int) (crop.left / scaleX + 0.5f),
+                (int) (crop.top / scaleY + 0.5f),
+                (int) (crop.right / scaleX + 0.5f),
+                (int) (crop.bottom / scaleY + 0.5f)
+        );
 
         // By default, offset it to to top/left corner
-        frame.offsetTo((int) (-crop.left / scale), (int) (-crop.top / scale));
+        frame.offsetTo((int) (-crop.left / scaleX), (int) (-crop.top / scaleY));
 
         // However, we also need to make space for the navigation bar on the left side.
         final int colorViewLeftInset = getColorViewLeftInset(mStableInsets.left,
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
index bd8aacb..7a2707b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotControllerTest.java
@@ -36,6 +36,7 @@
 import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 import android.util.ArraySet;
@@ -138,6 +139,7 @@
         final int orientation = Configuration.ORIENTATION_PORTRAIT;
         final float scaleFraction = 0.25f;
         final Rect contentInsets = new Rect(1, 2, 3, 4);
+        final Point taskSize = new Point(5, 6);
 
         try {
             ActivityManager.TaskSnapshot.Builder builder =
@@ -151,10 +153,10 @@
             builder.setOrientation(orientation);
             builder.setContentInsets(contentInsets);
             builder.setIsTranslucent(true);
-            builder.setScaleFraction(0.25f);
             builder.setSnapshot(buffer);
             builder.setIsRealSnapshot(true);
             builder.setPixelFormat(pixelFormat);
+            builder.setTaskSize(taskSize);
 
             // Not part of TaskSnapshot itself, used in screenshot process
             assertEquals(pixelFormat, builder.getPixelFormat());
@@ -166,12 +168,13 @@
             assertEquals(windowingMode, snapshot.getWindowingMode());
             assertEquals(sRGB, snapshot.getColorSpace());
             assertTrue(snapshot.isLowResolution());
+            assertTrue(builder.isLowResolution());
             assertEquals(orientation, snapshot.getOrientation());
             assertEquals(contentInsets, snapshot.getContentInsets());
             assertTrue(snapshot.isTranslucent());
-            assertEquals(scaleFraction, builder.getScaleFraction(), 0.001f);
             assertSame(buffer, snapshot.getSnapshot());
             assertTrue(snapshot.isRealSnapshot());
+            assertEquals(taskSize, snapshot.getTaskSize());
         } finally {
             if (buffer != null) {
                 buffer.destroy();
@@ -188,11 +191,9 @@
 
         final ActivityManager.TaskSnapshot.Builder builder =
                 new ActivityManager.TaskSnapshot.Builder();
-        final float scaleFraction = 0.8f;
         mWm.mTaskSnapshotController.prepareTaskSnapshot(mAppWindow.mActivityRecord.getTask(),
-                scaleFraction, PixelFormat.UNKNOWN, builder);
+                PixelFormat.UNKNOWN, builder);
 
-        assertEquals(scaleFraction, builder.getScaleFraction(), 0 /* delta */);
         // The pixel format should be selected automatically.
         assertNotEquals(PixelFormat.UNKNOWN, builder.getPixelFormat());
     }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java
index 0b16e5c..1d1391b 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterLoaderTest.java
@@ -150,7 +150,7 @@
     @Test
     public void testLowResolutionPersistAndLoadSnapshot() {
         TaskSnapshot a = new TaskSnapshotBuilder()
-                .setScale(0.5f)
+                .setScaleFraction(0.5f)
                 .setIsLowResolution(true)
                 .build();
         assertTrue(a.isLowResolution());
@@ -271,13 +271,11 @@
     @Test
     public void testScalePersistAndLoadSnapshot() {
         TaskSnapshot a = new TaskSnapshotBuilder()
-                .setScale(0.25f)
+                .setScaleFraction(0.25f)
                 .build();
         TaskSnapshot b = new TaskSnapshotBuilder()
-                .setScale(0.75f)
+                .setScaleFraction(0.75f)
                 .build();
-        assertEquals(0.25f, a.getScale(), 1E-5);
-        assertEquals(0.75f, b.getScale(), 1E-5);
         mPersister.persistSnapshot(1, mTestUserId, a);
         mPersister.persistSnapshot(2, mTestUserId, b);
         mPersister.waitForQueueEmpty();
@@ -287,8 +285,6 @@
                 false /* isLowResolution */);
         assertNotNull(snapshotA);
         assertNotNull(snapshotB);
-        assertEquals(0.25f, snapshotA.getScale(), 1E-5);
-        assertEquals(0.75f, snapshotB.getScale(), 1E-5);
     }
 
     @Test
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
index 4612dba..fa6663c 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
@@ -30,6 +30,7 @@
 import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.os.UserManager;
 import android.view.Surface;
@@ -87,8 +88,10 @@
      * Builds a TaskSnapshot.
      */
     static class TaskSnapshotBuilder {
+        private static final int SNAPSHOT_WIDTH = 100;
+        private static final int SNAPSHOT_HEIGHT = 100;
 
-        private float mScale = 1f;
+        private float mScaleFraction = 1f;
         private boolean mIsLowResolution = false;
         private boolean mIsRealSnapshot = true;
         private boolean mIsTranslucent = false;
@@ -96,8 +99,11 @@
         private int mSystemUiVisibility = 0;
         private int mRotation = Surface.ROTATION_0;
 
-        TaskSnapshotBuilder setScale(float scale) {
-            mScale = scale;
+        TaskSnapshotBuilder() {
+        }
+
+        TaskSnapshotBuilder setScaleFraction(float scale) {
+            mScaleFraction = scale;
             return this;
         }
 
@@ -132,15 +138,20 @@
         }
 
         TaskSnapshot build() {
-            final GraphicBuffer buffer = GraphicBuffer.create(100, 100, PixelFormat.RGBA_8888,
+            // To satisfy existing tests, ensure the graphics buffer is always 100x100, and
+            // compute the ize of the task according to mScaleFraction.
+            Point taskSize = new Point((int) (SNAPSHOT_WIDTH / mScaleFraction),
+                    (int) (SNAPSHOT_HEIGHT / mScaleFraction));
+            final GraphicBuffer buffer = GraphicBuffer.create(SNAPSHOT_WIDTH, SNAPSHOT_HEIGHT,
+                    PixelFormat.RGBA_8888,
                     USAGE_HW_TEXTURE | USAGE_SW_READ_RARELY | USAGE_SW_READ_RARELY);
             Canvas c = buffer.lockCanvas();
             c.drawColor(Color.RED);
             buffer.unlockCanvasAndPost(c);
             return new TaskSnapshot(MOCK_SNAPSHOT_ID, new ComponentName("", ""), buffer,
                     ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT,
-                    mRotation, TEST_INSETS,
-                    mIsLowResolution, mScale, mIsRealSnapshot,
+                    mRotation, taskSize, TEST_INSETS,
+                    mIsLowResolution, mIsRealSnapshot,
                     mWindowingMode, mSystemUiVisibility, mIsTranslucent);
         }
     }
diff --git a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
index bb0e5ae..2164de9 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
@@ -38,6 +38,7 @@
 import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
+import android.graphics.Point;
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
 import android.view.Surface;
@@ -67,12 +68,22 @@
             int windowFlags, Rect taskBounds) {
         final GraphicBuffer buffer = GraphicBuffer.create(width, height, PixelFormat.RGBA_8888,
                 GraphicBuffer.USAGE_SW_READ_RARELY | GraphicBuffer.USAGE_SW_WRITE_NEVER);
+
+        // Previously when constructing TaskSnapshots for this test, scale was 1.0f, so to mimic
+        // this behavior set the taskSize to be the same as the taskBounds width and height. The
+        // taskBounds passed here are assumed to be the same task bounds as when the snapshot was
+        // taken. We assume there is no aspect ratio mismatch between the screenshot and the
+        // taskBounds
+        assertEquals(width, taskBounds.width());
+        assertEquals(height, taskBounds.height());
+        Point taskSize = new Point(taskBounds.width(), taskBounds.height());
+
         final TaskSnapshot snapshot = new TaskSnapshot(
                 System.currentTimeMillis(),
                 new ComponentName("", ""), buffer,
                 ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT,
-                Surface.ROTATION_0, contentInsets, false,
-                1.0f, true /* isRealSnapshot */, WINDOWING_MODE_FULLSCREEN,
+                Surface.ROTATION_0, taskSize, contentInsets, false,
+                true /* isRealSnapshot */, WINDOWING_MODE_FULLSCREEN,
                 0 /* systemUiVisibility */, false /* isTranslucent */);
         mSurface = new TaskSnapshotSurface(mWm, new Window(), new SurfaceControl(), snapshot, "Test",
                 createTaskDescription(Color.WHITE, Color.RED, Color.BLUE), sysuiVis, windowFlags, 0,
@@ -152,7 +163,7 @@
 
     @Test
     public void testCalculateSnapshotCrop_taskNotOnTop() {
-        setupSurface(100, 100, new Rect(0, 10, 0, 10), 0, 0, new Rect(0, 50, 100, 100));
+        setupSurface(100, 100, new Rect(0, 10, 0, 10), 0, 0, new Rect(0, 50, 100, 150));
         assertEquals(new Rect(0, 10, 100, 90), mSurface.calculateSnapshotCrop());
     }