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",