Revert "Revert "Allow screenshot to carry more information.""

This reverts commit bf9d677afd8a820f96d75bd174264a94d7e12b82.

Below is the original commit message:
Previously screenshot returns a GraphicBuffer without any other information
about how to interpret this buffer, especially color space. We could not take
any WCG screenshot, and the color space information is also lost by the time we
create bitmap from the GraphicBuffer. This patch adds a ScreenshotGraphicBuffer
as a wrapper around GraphicBuffer with more information, for example color
space. And make screenshot native method return ScreenshotGraphicBuffer.

BUG: 116112787
Test: Build, flash and boot. Verify with WCG Photos.
Change-Id: I4db6aeaf086a183a8be4f731d9f1a10205d3b8ae
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 4a6c72b..94f96ba 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -126,6 +126,41 @@
     jfieldID white;
 } gDisplayPrimariesClassInfo;
 
+static struct {
+    jclass clazz;
+    jmethodID builder;
+} gScreenshotGraphicBufferClassInfo;
+
+class JNamedColorSpace {
+public:
+    // ColorSpace.Named.SRGB.ordinal() = 0;
+    static constexpr jint SRGB = 0;
+
+    // ColorSpace.Named.DISPLAY_P3.ordinal() = 6;
+    static constexpr jint DISPLAY_P3 = 6;
+};
+
+constexpr jint fromDataspaceToNamedColorSpaceValue(const ui::Dataspace dataspace) {
+    switch (dataspace) {
+        case ui::Dataspace::DISPLAY_P3:
+            return JNamedColorSpace::DISPLAY_P3;
+        default:
+            return JNamedColorSpace::SRGB;
+    }
+}
+
+constexpr ui::Dataspace pickDataspaceFromColorMode(const ui::ColorMode colorMode) {
+    switch (colorMode) {
+        case ui::ColorMode::DISPLAY_P3:
+        case ui::ColorMode::BT2100_PQ:
+        case ui::ColorMode::BT2100_HLG:
+        case ui::ColorMode::DISPLAY_BT2020:
+            return ui::Dataspace::DISPLAY_P3;
+        default:
+            return ui::Dataspace::V0_SRGB;
+    }
+}
+
 // ----------------------------------------------------------------------------
 
 static jlong nativeCreateTransaction(JNIEnv* env, jclass clazz) {
@@ -210,9 +245,12 @@
     if (displayToken == NULL) {
         return NULL;
     }
+    const ui::ColorMode colorMode = SurfaceComposerClient::getActiveColorMode(displayToken);
+    const ui::Dataspace dataspace = pickDataspaceFromColorMode(colorMode);
+
     Rect sourceCrop = rectFromObj(env, sourceCropObj);
     sp<GraphicBuffer> buffer;
-    status_t res = ScreenshotClient::capture(displayToken, ui::Dataspace::V0_SRGB,
+    status_t res = ScreenshotClient::capture(displayToken, dataspace,
             ui::PixelFormat::RGBA_8888,
             sourceCrop, width, height,
             useIdentityTransform, rotation, captureSecureLayers, &buffer);
@@ -220,13 +258,15 @@
         return NULL;
     }
 
-    return env->CallStaticObjectMethod(gGraphicBufferClassInfo.clazz,
-            gGraphicBufferClassInfo.builder,
+    const jint namedColorSpace = fromDataspaceToNamedColorSpaceValue(dataspace);
+    return env->CallStaticObjectMethod(gScreenshotGraphicBufferClassInfo.clazz,
+            gScreenshotGraphicBufferClassInfo.builder,
             buffer->getWidth(),
             buffer->getHeight(),
             buffer->getPixelFormat(),
             (jint)buffer->getUsage(),
-            (jlong)buffer.get());
+            (jlong)buffer.get(),
+            namedColorSpace);
 }
 
 static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerHandleToken,
@@ -243,20 +283,23 @@
     }
 
     sp<GraphicBuffer> buffer;
-    status_t res = ScreenshotClient::captureChildLayers(layerHandle, ui::Dataspace::V0_SRGB,
+    const ui::Dataspace dataspace = ui::Dataspace::V0_SRGB;
+    status_t res = ScreenshotClient::captureChildLayers(layerHandle, dataspace,
                                                         ui::PixelFormat::RGBA_8888, sourceCrop,
                                                         frameScale, &buffer);
     if (res != NO_ERROR) {
         return NULL;
     }
 
-    return env->CallStaticObjectMethod(gGraphicBufferClassInfo.clazz,
-                                       gGraphicBufferClassInfo.builder,
+    const jint namedColorSpace = fromDataspaceToNamedColorSpaceValue(dataspace);
+    return env->CallStaticObjectMethod(gScreenshotGraphicBufferClassInfo.clazz,
+                                       gScreenshotGraphicBufferClassInfo.builder,
                                        buffer->getWidth(),
                                        buffer->getHeight(),
                                        buffer->getPixelFormat(),
                                        (jint)buffer->getUsage(),
-                                       (jlong)buffer.get());
+                                       (jlong)buffer.get(),
+                                       namedColorSpace);
 }
 
 static void nativeApplyTransaction(JNIEnv* env, jclass clazz, jlong transactionObj, jboolean sync) {
@@ -1306,9 +1349,13 @@
             (void*)nativeSetOverrideScalingMode },
     {"nativeGetHandle", "(J)Landroid/os/IBinder;",
             (void*)nativeGetHandle },
-    {"nativeScreenshot", "(Landroid/os/IBinder;Landroid/graphics/Rect;IIZIZ)Landroid/graphics/GraphicBuffer;",
+    {"nativeScreenshot",
+            "(Landroid/os/IBinder;Landroid/graphics/Rect;IIZIZ)"
+            "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeScreenshot },
-    {"nativeCaptureLayers", "(Landroid/os/IBinder;Landroid/graphics/Rect;F)Landroid/graphics/GraphicBuffer;",
+    {"nativeCaptureLayers",
+            "(Landroid/os/IBinder;Landroid/graphics/Rect;F)"
+            "Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;",
             (void*)nativeCaptureLayers },
     {"nativeSetInputWindowInfo", "(JJLandroid/view/InputWindowHandle;)V",
             (void*)nativeSetInputWindowInfo },
@@ -1386,6 +1433,14 @@
     gGraphicBufferClassInfo.builder = GetStaticMethodIDOrDie(env, graphicsBufferClazz,
             "createFromExisting", "(IIIIJ)Landroid/graphics/GraphicBuffer;");
 
+    jclass screenshotGraphicsBufferClazz = FindClassOrDie(env,
+            "android/view/SurfaceControl$ScreenshotGraphicBuffer");
+    gScreenshotGraphicBufferClassInfo.clazz =
+            MakeGlobalRefOrDie(env, screenshotGraphicsBufferClazz);
+    gScreenshotGraphicBufferClassInfo.builder = GetStaticMethodIDOrDie(env,
+            screenshotGraphicsBufferClazz,
+            "createFromNative", "(IIIIJI)Landroid/view/SurfaceControl$ScreenshotGraphicBuffer;");
+
     jclass displayedContentSampleClazz = FindClassOrDie(env,
             "android/hardware/display/DisplayedContentSample");
     gDisplayedContentSampleClassInfo.clazz = MakeGlobalRefOrDie(env, displayedContentSampleClazz);