Determine the dataspace of the screenshot buffer based on display color mode.

Previously when we capture layers into graphic buffer we always assume SRGB
dataspace, however, if we have an app that is in wide color gamut mode, we want
to show the difference. This patch adds the ability to determine the dataspace
screenshot graphic buffer based on the color mode of the display.

BUG: 116112787
Test: Build, flash and boot. Verify with WCG Photos.
Change-Id: Ie2df32cad056576c256b9299a67855ed73714f50
diff --git a/core/java/android/app/ActivityManager.java b/core/java/android/app/ActivityManager.java
index 395c867..a30dab9 100644
--- a/core/java/android/app/ActivityManager.java
+++ b/core/java/android/app/ActivityManager.java
@@ -1787,17 +1787,17 @@
         private final float mScale;
         private final int mSystemUiVisibility;
         private final boolean mIsTranslucent;
-
-        // TODO(b/116112787) TaskSnapshot must also book keep the color space from hardware bitmap
-        // when created.
-        private final ColorSpace mColorSpace = ColorSpace.get(ColorSpace.Named.SRGB);
+        // Must be one of the named color spaces, otherwise, always use SRGB color space.
+        private final ColorSpace mColorSpace;
 
         public TaskSnapshot(@NonNull ComponentName topActivityComponent, GraphicBuffer snapshot,
-                int orientation, Rect contentInsets, boolean reducedResolution, float scale,
-                boolean isRealSnapshot, int windowingMode, int systemUiVisibility,
-                boolean isTranslucent) {
+                @NonNull ColorSpace colorSpace, int orientation, Rect contentInsets,
+                boolean reducedResolution, float scale, boolean isRealSnapshot, int windowingMode,
+                int systemUiVisibility, boolean isTranslucent) {
             mTopActivityComponent = topActivityComponent;
             mSnapshot = snapshot;
+            mColorSpace = colorSpace.getId() < 0
+                    ? ColorSpace.get(ColorSpace.Named.SRGB) : colorSpace;
             mOrientation = orientation;
             mContentInsets = new Rect(contentInsets);
             mReducedResolution = reducedResolution;
@@ -1811,6 +1811,10 @@
         private TaskSnapshot(Parcel source) {
             mTopActivityComponent = ComponentName.readFromParcel(source);
             mSnapshot = source.readParcelable(null /* classLoader */);
+            int colorSpaceId = source.readInt();
+            mColorSpace = colorSpaceId >= 0
+                    ? ColorSpace.get(ColorSpace.Named.values()[colorSpaceId])
+                    : ColorSpace.get(ColorSpace.Named.SRGB);
             mOrientation = source.readInt();
             mContentInsets = source.readParcelable(null /* classLoader */);
             mReducedResolution = source.readBoolean();
@@ -1917,6 +1921,7 @@
         public void writeToParcel(Parcel dest, int flags) {
             ComponentName.writeToParcel(mTopActivityComponent, dest);
             dest.writeParcelable(mSnapshot, 0);
+            dest.writeInt(mColorSpace.getId());
             dest.writeInt(mOrientation);
             dest.writeParcelable(mContentInsets, 0);
             dest.writeBoolean(mReducedResolution);
@@ -1934,6 +1939,7 @@
             return "TaskSnapshot{"
                     + " mTopActivityComponent=" + mTopActivityComponent.flattenToShortString()
                     + " mSnapshot=" + mSnapshot + " (" + width + "x" + height + ")"
+                    + " mColorSpace=" + mColorSpace.toString()
                     + " mOrientation=" + mOrientation
                     + " mContentInsets=" + mContentInsets.toShortString()
                     + " mReducedResolution=" + mReducedResolution + " mScale=" + mScale
diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index 79363ed..d34d19b 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -91,8 +91,8 @@
     private static native ScreenshotGraphicBuffer nativeScreenshot(IBinder displayToken,
             Rect sourceCrop, int width, int height, boolean useIdentityTransform, int rotation,
             boolean captureSecureLayers);
-    private static native ScreenshotGraphicBuffer nativeCaptureLayers(IBinder layerHandleToken,
-            Rect sourceCrop, float frameScale, IBinder[] excludeLayers);
+    private static native ScreenshotGraphicBuffer nativeCaptureLayers(IBinder displayToken,
+            IBinder layerHandleToken, Rect sourceCrop, float frameScale, IBinder[] excludeLayers);
 
     private static native long nativeCreateTransaction();
     private static native long nativeGetNativeTransactionFinalizer();
@@ -2001,7 +2001,8 @@
      */
     public static ScreenshotGraphicBuffer captureLayers(IBinder layerHandleToken, Rect sourceCrop,
             float frameScale) {
-        return nativeCaptureLayers(layerHandleToken, sourceCrop, frameScale, null);
+        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+        return nativeCaptureLayers(displayToken, layerHandleToken, sourceCrop, frameScale, null);
     }
 
     /**
@@ -2010,7 +2011,8 @@
      */
     public static ScreenshotGraphicBuffer captureLayersExcluding(IBinder layerHandleToken,
             Rect sourceCrop, float frameScale, IBinder[] exclude) {
-        return nativeCaptureLayers(layerHandleToken, sourceCrop, frameScale, exclude);
+        final IBinder displayToken = SurfaceControl.getInternalDisplayToken();
+        return nativeCaptureLayers(displayToken, layerHandleToken, sourceCrop, frameScale, exclude);
     }
 
     /**
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 3135c62..c0b31e4 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -271,8 +271,9 @@
             capturedSecureLayers);
 }
 
-static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerHandleToken,
-        jobject sourceCropObj, jfloat frameScale, jobjectArray excludeArray) {
+static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject displayTokenObj,
+        jobject layerHandleToken, jobject sourceCropObj, jfloat frameScale,
+        jobjectArray excludeArray) {
 
     sp<IBinder> layerHandle = ibinderForJavaObject(env, layerHandleToken);
     if (layerHandle == NULL) {
@@ -301,7 +302,12 @@
     }
 
     sp<GraphicBuffer> buffer;
-    const ui::Dataspace dataspace = ui::Dataspace::V0_SRGB;
+    ui::Dataspace dataspace = ui::Dataspace::V0_SRGB;
+    sp<IBinder> displayToken = ibinderForJavaObject(env, displayTokenObj);
+    if (displayToken != nullptr) {
+        const ui::ColorMode colorMode = SurfaceComposerClient::getActiveColorMode(displayToken);
+        dataspace = pickDataspaceFromColorMode(colorMode);
+    }
     status_t res = ScreenshotClient::captureChildLayers(layerHandle, dataspace,
                                                         ui::PixelFormat::RGBA_8888, sourceCrop,
                                                         excludeHandles, frameScale, &buffer);
@@ -1373,7 +1379,8 @@
             "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeScreenshot },
     {"nativeCaptureLayers",
-            "(Landroid/os/IBinder;Landroid/graphics/Rect;F[Landroid/os/IBinder;)"
+            "(Landroid/os/IBinder;Landroid/os/IBinder;Landroid/graphics/Rect;"
+            "F[Landroid/os/IBinder;)"
             "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeCaptureLayers },
     {"nativeSetInputWindowInfo", "(JJLandroid/view/InputWindowHandle;)V",
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index 6fe8b43..f559848 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -292,9 +292,9 @@
         }
         final boolean isWindowTranslucent = mainWindow.getAttrs().format != PixelFormat.OPAQUE;
         return new TaskSnapshot(appWindowToken.mActivityComponent, buffer,
-                appWindowToken.getConfiguration().orientation, getInsets(mainWindow),
-                isLowRamDevice /* reduced */, scaleFraction /* scale */, true /* isRealSnapshot */,
-                task.getWindowingMode(), getSystemUiVisibility(task),
+                screenshotBuffer.getColorSpace(), appWindowToken.getConfiguration().orientation,
+                getInsets(mainWindow), isLowRamDevice /* reduced */, scaleFraction /* scale */,
+                true /* isRealSnapshot */, task.getWindowingMode(), getSystemUiVisibility(task),
                 !appWindowToken.fillsParent() || isWindowTranslucent);
     }
 
@@ -385,10 +385,10 @@
         // Note, the app theme snapshot is never translucent because we enforce a non-translucent
         // color above
         return new TaskSnapshot(topChild.mActivityComponent, hwBitmap.createGraphicBufferHandle(),
-                topChild.getConfiguration().orientation, mainWindow.getStableInsets(),
-                ActivityManager.isLowRamDeviceStatic() /* reduced */, 1.0f /* scale */,
-                false /* isRealSnapshot */, task.getWindowingMode(), getSystemUiVisibility(task),
-                false);
+                hwBitmap.getColorSpace(), topChild.getConfiguration().orientation,
+                mainWindow.getStableInsets(), ActivityManager.isLowRamDeviceStatic() /* reduced */,
+                1.0f /* scale */, 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 d30843b..f7b8945 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotLoader.java
@@ -89,7 +89,8 @@
             }
             ComponentName topActivityComponent = ComponentName.unflattenFromString(
                     proto.topActivityComponent);
-            return new TaskSnapshot(topActivityComponent, buffer, proto.orientation,
+            return new TaskSnapshot(topActivityComponent, buffer, bitmap.getColorSpace(),
+                    proto.orientation,
                     new Rect(proto.insetLeft, proto.insetTop, proto.insetRight, proto.insetBottom),
                     reducedResolution, reducedResolution ? mPersister.getReducedScale() : 1f,
                     proto.isRealSnapshot, proto.windowingMode, proto.systemUiVisibility,
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 d29e3fa..c3b0a67 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotPersisterTestBase.java
@@ -27,6 +27,7 @@
 import android.content.ComponentName;
 import android.graphics.Canvas;
 import android.graphics.Color;
+import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
@@ -127,8 +128,9 @@
             Canvas c = buffer.lockCanvas();
             c.drawColor(Color.RED);
             buffer.unlockCanvasAndPost(c);
-            return new TaskSnapshot(new ComponentName("", ""), buffer, ORIENTATION_PORTRAIT,
-                    TEST_INSETS, mScale < 1f /* reducedResolution */, mScale, mIsRealSnapshot,
+            return new TaskSnapshot(new ComponentName("", ""), buffer,
+                    ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT, TEST_INSETS,
+                    mScale < 1f /* reducedResolution */, mScale, 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 ca815ec..fc45456 100644
--- a/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
+++ b/services/tests/wmtests/src/com/android/server/wm/TaskSnapshotSurfaceTest.java
@@ -34,11 +34,11 @@
 import android.content.ComponentName;
 import android.graphics.Canvas;
 import android.graphics.Color;
+import android.graphics.ColorSpace;
 import android.graphics.GraphicBuffer;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.platform.test.annotations.Presubmit;
-import android.view.Surface;
 import android.view.SurfaceControl;
 
 import androidx.test.filters.SmallTest;
@@ -64,8 +64,9 @@
         final GraphicBuffer buffer = GraphicBuffer.create(width, height, PixelFormat.RGBA_8888,
                 GraphicBuffer.USAGE_SW_READ_RARELY | GraphicBuffer.USAGE_SW_WRITE_NEVER);
         final TaskSnapshot snapshot = new TaskSnapshot(new ComponentName("", ""), buffer,
-                ORIENTATION_PORTRAIT, contentInsets, false, 1.0f, true /* isRealSnapshot */,
-                WINDOWING_MODE_FULLSCREEN, 0 /* systemUiVisibility */, false /* isTranslucent */);
+                ColorSpace.get(ColorSpace.Named.SRGB), ORIENTATION_PORTRAIT, contentInsets, false,
+                1.0f, true /* isRealSnapshot */, WINDOWING_MODE_FULLSCREEN,
+                0 /* systemUiVisibility */, false /* isTranslucent */);
         mSurface = new TaskSnapshotSurface(mWm, new Window(), new SurfaceControl(), snapshot, "Test",
                 Color.WHITE, Color.RED, Color.BLUE, sysuiVis, windowFlags, 0, taskBounds,
                 ORIENTATION_PORTRAIT);