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

6 files changed, 138 insertions, 41 deletions

diff --git a/core/java/android/view/SurfaceControl.java b/core/java/android/view/SurfaceControl.java
index cd075bf65e4a..6ae5167b57a1 100644
--- a/core/java/android/view/SurfaceControl.java
+++ b/core/java/android/view/SurfaceControl.java
@@ -88,10 +88,10 @@ public final class SurfaceControl implements Parcelable {
     private static native void nativeDestroy(long nativeObject);
     private static native void nativeDisconnect(long nativeObject);
 
-    private static native GraphicBuffer nativeScreenshot(IBinder displayToken,
+    private static native ScreenshotGraphicBuffer nativeScreenshot(IBinder displayToken,
             Rect sourceCrop, int width, int height, boolean useIdentityTransform, int rotation,
             boolean captureSecureLayers);
-    private static native GraphicBuffer nativeCaptureLayers(IBinder layerHandleToken,
+    private static native ScreenshotGraphicBuffer nativeCaptureLayers(IBinder layerHandleToken,
             Rect sourceCrop, float frameScale);
 
     private static native long nativeCreateTransaction();
@@ -431,6 +431,46 @@ public final class SurfaceControl implements Parcelable {
     public static final int METADATA_TASK_ID = 3;
 
     /**
+     * A wrapper around GraphicBuffer that contains extra information about how to
+     * interpret the screenshot GraphicBuffer.
+     * @hide
+     */
+    public static class ScreenshotGraphicBuffer {
+        private final GraphicBuffer mGraphicBuffer;
+        private final ColorSpace mColorSpace;
+
+        public ScreenshotGraphicBuffer(GraphicBuffer graphicBuffer, ColorSpace colorSpace) {
+            mGraphicBuffer = graphicBuffer;
+            mColorSpace = colorSpace;
+        }
+
+       /**
+        * Create ScreenshotGraphicBuffer from existing native GraphicBuffer object.
+        * @param width The width in pixels of the buffer
+        * @param height The height in pixels of the buffer
+        * @param format The format of each pixel as specified in {@link PixelFormat}
+        * @param usage Hint indicating how the buffer will be used
+        * @param unwrappedNativeObject The native object of GraphicBuffer
+        * @param namedColorSpace Integer value of a named color space {@link ColorSpace.Named}
+        */
+        private static ScreenshotGraphicBuffer createFromNative(int width, int height, int format,
+                int usage, long unwrappedNativeObject, int namedColorSpace) {
+            GraphicBuffer graphicBuffer = GraphicBuffer.createFromExisting(width, height, format,
+                    usage, unwrappedNativeObject);
+            ColorSpace colorSpace = ColorSpace.get(ColorSpace.Named.values()[namedColorSpace]);
+            return new ScreenshotGraphicBuffer(graphicBuffer, colorSpace);
+        }
+
+        public ColorSpace getColorSpace() {
+            return mColorSpace;
+        }
+
+        public GraphicBuffer getGraphicBuffer() {
+            return mGraphicBuffer;
+        }
+    }
+
+    /**
      * Builder class for {@link SurfaceControl} objects.
      */
     public static class Builder {
@@ -1815,10 +1855,10 @@ public final class SurfaceControl implements Parcelable {
             throw new IllegalArgumentException("consumer must not be null");
         }
 
-        final GraphicBuffer buffer = screenshotToBuffer(display, sourceCrop, width, height,
-                useIdentityTransform, rotation);
+        final ScreenshotGraphicBuffer buffer = screenshotToBuffer(display, sourceCrop, width,
+                height, useIdentityTransform, rotation);
         try {
-            consumer.attachAndQueueBuffer(buffer);
+            consumer.attachAndQueueBuffer(buffer.getGraphicBuffer());
         } catch (RuntimeException e) {
             Log.w(TAG, "Failed to take screenshot - " + e.getMessage());
         }
@@ -1861,17 +1901,16 @@ public final class SurfaceControl implements Parcelable {
         }
 
         SurfaceControl.rotateCropForSF(sourceCrop, rotation);
-        final GraphicBuffer buffer = screenshotToBuffer(displayToken, sourceCrop, width, height,
-                useIdentityTransform, rotation);
+        final ScreenshotGraphicBuffer buffer = screenshotToBuffer(displayToken, sourceCrop, width,
+                height, useIdentityTransform, rotation);
 
         if (buffer == null) {
             Log.w(TAG, "Failed to take screenshot");
             return null;
         }
-        // TODO(b/116112787) Now that hardware bitmap creation can take color space, we
-        // should continue to fix screenshot.
-        return Bitmap.wrapHardwareBuffer(HardwareBuffer.createFromGraphicBuffer(buffer),
-                ColorSpace.get(ColorSpace.Named.SRGB));
+        return Bitmap.wrapHardwareBuffer(
+                HardwareBuffer.createFromGraphicBuffer(buffer.getGraphicBuffer()),
+                buffer.getColorSpace());
     }
 
     /**
@@ -1897,8 +1936,8 @@ public final class SurfaceControl implements Parcelable {
      * @return Returns a GraphicBuffer that contains the captured content.
      * @hide
      */
-    public static GraphicBuffer screenshotToBuffer(IBinder display, Rect sourceCrop, int width,
-            int height, boolean useIdentityTransform, int rotation) {
+    public static ScreenshotGraphicBuffer screenshotToBuffer(IBinder display, Rect sourceCrop,
+            int width, int height, boolean useIdentityTransform, int rotation) {
         if (display == null) {
             throw new IllegalArgumentException("displayToken must not be null");
         }
@@ -1917,7 +1956,7 @@ public final class SurfaceControl implements Parcelable {
      *
      * @hide
      */
-    public static GraphicBuffer screenshotToBufferWithSecureLayersUnsafe(IBinder display,
+    public static ScreenshotGraphicBuffer screenshotToBufferWithSecureLayersUnsafe(IBinder display,
             Rect sourceCrop, int width, int height, boolean useIdentityTransform,
             int rotation) {
         if (display == null) {
@@ -1951,7 +1990,7 @@ public final class SurfaceControl implements Parcelable {
      * @return Returns a GraphicBuffer that contains the layer capture.
      * @hide
      */
-    public static GraphicBuffer captureLayers(IBinder layerHandleToken, Rect sourceCrop,
+    public static ScreenshotGraphicBuffer captureLayers(IBinder layerHandleToken, Rect sourceCrop,
             float frameScale) {
         return nativeCaptureLayers(layerHandleToken, sourceCrop, frameScale);
     }
diff --git a/core/jni/android_view_SurfaceControl.cpp b/core/jni/android_view_SurfaceControl.cpp
index 4a6c72bda283..94f96babaf90 100644
--- a/core/jni/android_view_SurfaceControl.cpp
+++ b/core/jni/android_view_SurfaceControl.cpp
@@ -126,6 +126,41 @@ static struct {
     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 @@ static jobject nativeScreenshot(JNIEnv* env, jclass clazz,
     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 @@ static jobject nativeScreenshot(JNIEnv* env, jclass clazz,
         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 @@ static jobject nativeCaptureLayers(JNIEnv* env, jclass clazz, jobject layerHandl
     }
 
     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 @@ static const JNINativeMethod sSurfaceControlMethods[] = {
             (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 @@ int register_android_view_SurfaceControl(JNIEnv* env)
     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);
diff --git a/services/core/java/com/android/server/display/DisplayManagerService.java b/services/core/java/com/android/server/display/DisplayManagerService.java
index 32f34b8378ce..3010324488b8 100644
--- a/services/core/java/com/android/server/display/DisplayManagerService.java
+++ b/services/core/java/com/android/server/display/DisplayManagerService.java
@@ -37,7 +37,6 @@ import android.content.pm.ParceledListSlice;
 import android.content.res.Resources;
 import android.content.res.TypedArray;
 import android.graphics.ColorSpace;
-import android.graphics.GraphicBuffer;
 import android.graphics.Point;
 import android.graphics.Rect;
 import android.hardware.SensorManager;
@@ -1275,11 +1274,12 @@ public final class DisplayManagerService extends SystemService {
         if (token == null) {
             return false;
         }
-        final GraphicBuffer gb = SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(
-                token, new Rect(), 0 /* width */, 0 /* height */, false /* useIdentityTransform */,
-                0 /* rotation */);
+        final SurfaceControl.ScreenshotGraphicBuffer gb =
+                SurfaceControl.screenshotToBufferWithSecureLayersUnsafe(
+                        token, new Rect(), 0 /* width */, 0 /* height */,
+                        false /* useIdentityTransform */, 0 /* rotation */);
         try {
-            outSurface.attachAndQueueBuffer(gb);
+            outSurface.attachAndQueueBuffer(gb.getGraphicBuffer());
         } catch (RuntimeException e) {
             Slog.w(TAG, "Failed to take screenshot - " + e.getMessage());
         }
diff --git a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
index e0d85e876476..4379b7cdace6 100644
--- a/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
+++ b/services/core/java/com/android/server/wm/TaskScreenshotAnimatable.java
@@ -41,7 +41,7 @@ class TaskScreenshotAnimatable implements SurfaceAnimator.Animatable {
         return new TaskScreenshotAnimatable(task, getBufferFromTask(task));
     }
 
-    private static GraphicBuffer getBufferFromTask(Task task) {
+    private static SurfaceControl.ScreenshotGraphicBuffer getBufferFromTask(Task task) {
         if (task == null) {
             return null;
         }
@@ -51,7 +51,10 @@ class TaskScreenshotAnimatable implements SurfaceAnimator.Animatable {
                 task.getSurfaceControl().getHandle(), tmpRect, 1f);
     }
 
-    private TaskScreenshotAnimatable(Task task, GraphicBuffer buffer) {
+    private TaskScreenshotAnimatable(Task task,
+            SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer) {
+        GraphicBuffer buffer = screenshotBuffer == null
+                ? null : screenshotBuffer.getGraphicBuffer();
         mTask = task;
         mWidth = (buffer != null) ? buffer.getWidth() : 1;
         mHeight = (buffer != null) ? buffer.getHeight() : 1;
diff --git a/services/core/java/com/android/server/wm/TaskSnapshotController.java b/services/core/java/com/android/server/wm/TaskSnapshotController.java
index beb3d82c05fb..6fe8b43db139 100644
--- a/services/core/java/com/android/server/wm/TaskSnapshotController.java
+++ b/services/core/java/com/android/server/wm/TaskSnapshotController.java
@@ -279,8 +279,11 @@ class TaskSnapshotController {
             Slog.w(TAG_WM, "Failed to take screenshot. No main window for " + task);
             return null;
         }
-        final GraphicBuffer buffer = SurfaceControl.captureLayers(
-                task.getSurfaceControl().getHandle(), mTmpRect, scaleFraction);
+        final SurfaceControl.ScreenshotGraphicBuffer screenshotBuffer =
+                SurfaceControl.captureLayers(
+                        task.getSurfaceControl().getHandle(), mTmpRect, scaleFraction);
+        final GraphicBuffer buffer = screenshotBuffer != null ? screenshotBuffer.getGraphicBuffer()
+                : null;
         if (buffer == null || buffer.getWidth() <= 1 || buffer.getHeight() <= 1) {
             if (DEBUG_SCREENSHOT) {
                 Slog.w(TAG_WM, "Failed to take screenshot for " + task);
diff --git a/services/core/java/com/android/server/wm/WallpaperController.java b/services/core/java/com/android/server/wm/WallpaperController.java
index dddc6b755db0..166a33d7a560 100644
--- a/services/core/java/com/android/server/wm/WallpaperController.java
+++ b/services/core/java/com/android/server/wm/WallpaperController.java
@@ -34,8 +34,6 @@ import static com.android.server.wm.WindowManagerDebugConfig.TAG_WM;
 import static com.android.server.wm.WindowManagerService.H.WALLPAPER_DRAW_PENDING_TIMEOUT;
 
 import android.graphics.Bitmap;
-import android.graphics.ColorSpace;
-import android.graphics.GraphicBuffer;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
 import android.os.Bundle;
@@ -738,17 +736,16 @@ class WallpaperController {
         final Rect bounds = wallpaperWindowState.getBounds();
         bounds.offsetTo(0, 0);
 
-        GraphicBuffer wallpaperBuffer = SurfaceControl.captureLayers(
+        SurfaceControl.ScreenshotGraphicBuffer wallpaperBuffer = SurfaceControl.captureLayers(
                 wallpaperWindowState.getSurfaceControl().getHandle(), bounds, 1 /* frameScale */);
 
         if (wallpaperBuffer == null) {
             Slog.w(TAG_WM, "Failed to screenshot wallpaper");
             return null;
         }
-        // TODO(b/116112787) Now that hardware bitmap creation can take color space, we
-        // should continue to fix screenshot.
-        return Bitmap.wrapHardwareBuffer(HardwareBuffer.createFromGraphicBuffer(wallpaperBuffer),
-                                         ColorSpace.get(ColorSpace.Named.SRGB));
+        return Bitmap.wrapHardwareBuffer(
+                HardwareBuffer.createFromGraphicBuffer(wallpaperBuffer.getGraphicBuffer()),
+                wallpaperBuffer.getColorSpace());
     }
 
     private WindowState getTopVisibleWallpaper() {