Add gainmap support for screencap

There still would need to be support added to the PNG spec to properly
support gainmaps without affecting downstream clients. So, this patch:

1. Allows for screencap to (temporarily) export jpegs
2. Adds plumbing in HWUI's apex layer to encode gainmaps
3. Wires up the attachGainmap flag to allow screenshots to output a
   gainmap

Bug: 329470026
Flag: com.android.graphics.surfaceflinger.flags.true_hdr_screenshots
Test: adb screencap -j sdcard/test.jpeg
Change-Id: I210a3e24ad2cfd6e0c0a954f42b9171d9e82e991

5 files changed, 186 insertions, 46 deletions

diff --git a/cmds/screencap/Android.bp b/cmds/screencap/Android.bp
index c009c1f5b08b..16026eca2980 100644
--- a/cmds/screencap/Android.bp
+++ b/cmds/screencap/Android.bp
@@ -17,6 +17,7 @@ cc_binary {
         "libutils",
         "libbinder",
         "libjnigraphics",
+        "libhwui",
         "libui",
         "libgui",
     ],
diff --git a/cmds/screencap/screencap.cpp b/cmds/screencap/screencap.cpp
index 01b20f4a5267..12de82a46263 100644
--- a/cmds/screencap/screencap.cpp
+++ b/cmds/screencap/screencap.cpp
@@ -15,36 +15,28 @@
  */
 
 #include <android/bitmap.h>
+#include <android/graphics/bitmap.h>
 #include <android/gui/DisplayCaptureArgs.h>
 #include <binder/ProcessState.h>
 #include <errno.h>
-#include <unistd.h>
-#include <stdio.h>
 #include <fcntl.h>
-#include <stdlib.h>
-#include <string.h>
-#include <getopt.h>
-
-#include <linux/fb.h>
-#include <sys/ioctl.h>
-#include <sys/mman.h>
-#include <sys/wait.h>
-
-#include <android/bitmap.h>
-
-#include <binder/ProcessState.h>
-
 #include <ftl/concat.h>
 #include <ftl/optional.h>
+#include <getopt.h>
 #include <gui/ISurfaceComposer.h>
 #include <gui/SurfaceComposerClient.h>
 #include <gui/SyncScreenCaptureListener.h>
-
+#include <linux/fb.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <sys/ioctl.h>
+#include <sys/mman.h>
+#include <sys/wait.h>
+#include <system/graphics.h>
 #include <ui/GraphicTypes.h>
 #include <ui/PixelFormat.h>
 
-#include <system/graphics.h>
-
 using namespace android;
 
 #define COLORSPACE_UNKNOWN    0
@@ -85,11 +77,12 @@ enum {
 };
 }
 
-static const struct option LONG_OPTIONS[] = {
-        {"png", no_argument, nullptr, 'p'},
-        {"help", no_argument, nullptr, 'h'},
-        {"hint-for-seamless", no_argument, nullptr, LongOpts::HintForSeamless},
-        {0, 0, 0, 0}};
+static const struct option LONG_OPTIONS[] = {{"png", no_argument, nullptr, 'p'},
+                                             {"jpeg", no_argument, nullptr, 'j'},
+                                             {"help", no_argument, nullptr, 'h'},
+                                             {"hint-for-seamless", no_argument, nullptr,
+                                              LongOpts::HintForSeamless},
+                                             {0, 0, 0, 0}};
 
 static int32_t flinger2bitmapFormat(PixelFormat f)
 {
@@ -170,10 +163,11 @@ status_t capture(const DisplayId displayId,
     return 0;
 }
 
-status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& captureResults) {
+status_t saveImage(const char* fn, std::optional<AndroidBitmapCompressFormat> format,
+                   const ScreenCaptureResults& captureResults) {
     void* base = nullptr;
     ui::Dataspace dataspace = captureResults.capturedDataspace;
-    sp<GraphicBuffer> buffer = captureResults.buffer;
+    const sp<GraphicBuffer>& buffer = captureResults.buffer;
 
     status_t result = buffer->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &base);
 
@@ -188,22 +182,48 @@ status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& capture
         return 1;
     }
 
+    void* gainmapBase = nullptr;
+    sp<GraphicBuffer> gainmap = captureResults.optionalGainMap;
+
+    if (gainmap) {
+        result = gainmap->lock(GraphicBuffer::USAGE_SW_READ_OFTEN, &gainmapBase);
+        if (gainmapBase == nullptr || result != NO_ERROR) {
+            fprintf(stderr, "Failed to capture gainmap with error code (%d)\n", result);
+            gainmapBase = nullptr;
+            // Fall-through: just don't attempt to write the gainmap
+        }
+    }
+
     int fd = -1;
     if (fn == nullptr) {
         fd = dup(STDOUT_FILENO);
         if (fd == -1) {
             fprintf(stderr, "Error writing to stdout. (%s)\n", strerror(errno));
+            if (gainmapBase) {
+                gainmap->unlock();
+            }
+
+            if (base) {
+                buffer->unlock();
+            }
             return 1;
         }
     } else {
         fd = open(fn, O_WRONLY | O_CREAT | O_TRUNC, 0664);
         if (fd == -1) {
             fprintf(stderr, "Error opening file: %s (%s)\n", fn, strerror(errno));
+            if (gainmapBase) {
+                gainmap->unlock();
+            }
+
+            if (base) {
+                buffer->unlock();
+            }
             return 1;
         }
     }
 
-    if (png) {
+    if (format) {
         AndroidBitmapInfo info;
         info.format = flinger2bitmapFormat(buffer->getPixelFormat());
         info.flags = ANDROID_BITMAP_FLAGS_ALPHA_PREMUL;
@@ -211,16 +231,31 @@ status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& capture
         info.height = buffer->getHeight();
         info.stride = buffer->getStride() * bytesPerPixel(buffer->getPixelFormat());
 
-        int result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base,
-                                            ANDROID_BITMAP_COMPRESS_FORMAT_PNG, 100, &fd,
+        int result;
+
+        if (gainmapBase) {
+            result = ABitmap_compressWithGainmap(&info, static_cast<ADataSpace>(dataspace), base,
+                                                 gainmapBase, captureResults.hdrSdrRatio, *format,
+                                                 100, &fd,
+                                                 [](void* fdPtr, const void* data,
+                                                    size_t size) -> bool {
+                                                     int bytesWritten =
+                                                             write(*static_cast<int*>(fdPtr), data,
+                                                                   size);
+                                                     return bytesWritten == size;
+                                                 });
+        } else {
+            result = AndroidBitmap_compress(&info, static_cast<int32_t>(dataspace), base, *format,
+                                            100, &fd,
                                             [](void* fdPtr, const void* data, size_t size) -> bool {
                                                 int bytesWritten = write(*static_cast<int*>(fdPtr),
                                                                          data, size);
                                                 return bytesWritten == size;
                                             });
+        }
 
         if (result != ANDROID_BITMAP_RESULT_SUCCESS) {
-            fprintf(stderr, "Failed to compress PNG (error code: %d)\n", result);
+            fprintf(stderr, "Failed to compress (error code: %d)\n", result);
         }
 
         if (fn != NULL) {
@@ -245,6 +280,14 @@ status_t saveImage(const char* fn, bool png, const ScreenCaptureResults& capture
     }
     close(fd);
 
+    if (gainmapBase) {
+        gainmap->unlock();
+    }
+
+    if (base) {
+        buffer->unlock();
+    }
+
     return 0;
 }
 
@@ -262,13 +305,17 @@ int main(int argc, char** argv)
     gui::CaptureArgs captureArgs;
     const char* pname = argv[0];
     bool png = false;
+    bool jpeg = false;
     bool all = false;
     int c;
-    while ((c = getopt_long(argc, argv, "aphd:", LONG_OPTIONS, nullptr)) != -1) {
+    while ((c = getopt_long(argc, argv, "apjhd:", LONG_OPTIONS, nullptr)) != -1) {
         switch (c) {
             case 'p':
                 png = true;
                 break;
+            case 'j':
+                jpeg = true;
+                break;
             case 'd': {
                 errno = 0;
                 char* end = nullptr;
@@ -325,6 +372,14 @@ int main(int argc, char** argv)
             baseName = filename.substr(0, filename.size()-4);
             suffix = ".png";
             png = true;
+        } else if (filename.ends_with(".jpeg")) {
+            baseName = filename.substr(0, filename.size() - 5);
+            suffix = ".jpeg";
+            jpeg = true;
+        } else if (filename.ends_with(".jpg")) {
+            baseName = filename.substr(0, filename.size() - 4);
+            suffix = ".jpg";
+            jpeg = true;
         } else {
             baseName = filename;
         }
@@ -350,6 +405,20 @@ int main(int argc, char** argv)
         }
     }
 
+    if (png && jpeg) {
+        fprintf(stderr, "Ambiguous file type");
+        return 1;
+    }
+
+    std::optional<AndroidBitmapCompressFormat> format = std::nullopt;
+
+    if (png) {
+        format = ANDROID_BITMAP_COMPRESS_FORMAT_PNG;
+    } else if (jpeg) {
+        format = ANDROID_BITMAP_COMPRESS_FORMAT_JPEG;
+        captureArgs.attachGainmap = true;
+    }
+
     // setThreadPoolMaxThreadCount(0) actually tells the kernel it's
     // not allowed to spawn any additional threads, but we still spawn
     // a binder thread from userspace when we call startThreadPool().
@@ -385,7 +454,7 @@ int main(int argc, char** argv)
         if (!filename.empty()) {
             fn = filename.c_str();
         }
-        if (const status_t saveImageStatus = saveImage(fn, png, result) != 0) {
+        if (const status_t saveImageStatus = saveImage(fn, format, result) != 0) {
             fprintf(stderr, "Saving image failed.\n");
             return saveImageStatus;
         }
diff --git a/libs/hwui/apex/android_bitmap.cpp b/libs/hwui/apex/android_bitmap.cpp
index c80a9b4ae97f..000f1092eb17 100644
--- a/libs/hwui/apex/android_bitmap.cpp
+++ b/libs/hwui/apex/android_bitmap.cpp
@@ -14,21 +14,21 @@
  * limitations under the License.
  */
 
-#include <log/log.h>
-
-#include "android/graphics/bitmap.h"
-#include "TypeCast.h"
-#include "GraphicsJNI.h"
-
+#include <Gainmap.h>
 #include <GraphicsJNI.h>
-#include <hwui/Bitmap.h>
 #include <SkBitmap.h>
 #include <SkColorSpace.h>
 #include <SkImageInfo.h>
 #include <SkRefCnt.h>
 #include <SkStream.h>
+#include <hwui/Bitmap.h>
+#include <log/log.h>
 #include <utils/Color.h>
 
+#include "GraphicsJNI.h"
+#include "TypeCast.h"
+#include "android/graphics/bitmap.h"
+
 using namespace android;
 
 ABitmap* ABitmap_acquireBitmapFromJava(JNIEnv* env, jobject bitmapObj) {
@@ -215,6 +215,14 @@ private:
 int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const void* pixels,
                      AndroidBitmapCompressFormat inFormat, int32_t quality, void* userContext,
                      AndroidBitmap_CompressWriteFunc fn) {
+    return ABitmap_compressWithGainmap(info, dataSpace, pixels, nullptr, -1.f, inFormat, quality,
+                                       userContext, fn);
+}
+
+int ABitmap_compressWithGainmap(const AndroidBitmapInfo* info, ADataSpace dataSpace,
+                                const void* pixels, const void* gainmapPixels, float hdrSdrRatio,
+                                AndroidBitmapCompressFormat inFormat, int32_t quality,
+                                void* userContext, AndroidBitmap_CompressWriteFunc fn) {
     Bitmap::JavaCompressFormat format;
     switch (inFormat) {
         case ANDROID_BITMAP_COMPRESS_FORMAT_JPEG:
@@ -275,7 +283,7 @@ int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const
         // besides ADATASPACE_UNKNOWN as an error?
         cs = nullptr;
     } else {
-        cs = uirenderer::DataSpaceToColorSpace((android_dataspace) dataSpace);
+        cs = uirenderer::DataSpaceToColorSpace((android_dataspace)dataSpace);
         // DataSpaceToColorSpace treats UNKNOWN as SRGB, but compress forces the
         // client to specify SRGB if that is what they want.
         if (!cs || dataSpace == ADATASPACE_UNKNOWN) {
@@ -292,16 +300,70 @@ int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const
 
     auto imageInfo =
             SkImageInfo::Make(info->width, info->height, colorType, alphaType, std::move(cs));
-    SkBitmap bitmap;
-    // We are not going to modify the pixels, but installPixels expects them to
-    // not be const, since for all it knows we might want to draw to the SkBitmap.
-    if (!bitmap.installPixels(imageInfo, const_cast<void*>(pixels), info->stride)) {
-        return ANDROID_BITMAP_RESULT_BAD_PARAMETER;
+
+    // Validate the image info
+    {
+        SkBitmap tempBitmap;
+        if (!tempBitmap.installPixels(imageInfo, const_cast<void*>(pixels), info->stride)) {
+            return ANDROID_BITMAP_RESULT_BAD_PARAMETER;
+        }
+    }
+
+    SkPixelRef pixelRef =
+            SkPixelRef(info->width, info->height, const_cast<void*>(pixels), info->stride);
+
+    auto bitmap = Bitmap::createFrom(imageInfo, pixelRef);
+
+    if (gainmapPixels) {
+        auto gainmap = sp<uirenderer::Gainmap>::make();
+        gainmap->info.fGainmapRatioMin = {
+                1.f,
+                1.f,
+                1.f,
+                1.f,
+        };
+        gainmap->info.fGainmapRatioMax = {
+                hdrSdrRatio,
+                hdrSdrRatio,
+                hdrSdrRatio,
+                1.f,
+        };
+        gainmap->info.fGainmapGamma = {
+                1.f,
+                1.f,
+                1.f,
+                1.f,
+        };
+
+        static constexpr auto kDefaultEpsilon = 1.f / 64.f;
+        gainmap->info.fEpsilonSdr = {
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                1.f,
+        };
+        gainmap->info.fEpsilonHdr = {
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                kDefaultEpsilon,
+                1.f,
+        };
+        gainmap->info.fDisplayRatioSdr = 1.f;
+        gainmap->info.fDisplayRatioHdr = hdrSdrRatio;
+
+        SkPixelRef gainmapPixelRef = SkPixelRef(info->width, info->height,
+                                                const_cast<void*>(gainmapPixels), info->stride);
+        auto gainmapBitmap = Bitmap::createFrom(imageInfo, gainmapPixelRef);
+        gainmap->bitmap = std::move(gainmapBitmap);
+        bitmap->setGainmap(std::move(gainmap));
     }
 
     CompressWriter stream(userContext, fn);
-    return Bitmap::compress(bitmap, format, quality, &stream) ? ANDROID_BITMAP_RESULT_SUCCESS
-                                                              : ANDROID_BITMAP_RESULT_JNI_EXCEPTION;
+
+    return bitmap->compress(format, quality, &stream)
+
+                   ? ANDROID_BITMAP_RESULT_SUCCESS
+                   : ANDROID_BITMAP_RESULT_JNI_EXCEPTION;
 }
 
 AHardwareBuffer* ABitmap_getHardwareBuffer(ABitmap* bitmapHandle) {
diff --git a/libs/hwui/apex/include/android/graphics/bitmap.h b/libs/hwui/apex/include/android/graphics/bitmap.h
index 8c4b439d2a2b..6f65e9eb0495 100644
--- a/libs/hwui/apex/include/android/graphics/bitmap.h
+++ b/libs/hwui/apex/include/android/graphics/bitmap.h
@@ -65,6 +65,13 @@ ANDROID_API jobject ABitmapConfig_getConfigFromFormat(JNIEnv* env, AndroidBitmap
 ANDROID_API int ABitmap_compress(const AndroidBitmapInfo* info, ADataSpace dataSpace, const void* pixels,
                      AndroidBitmapCompressFormat format, int32_t quality, void* userContext,
                      AndroidBitmap_CompressWriteFunc);
+// If gainmapPixels is null, then no gainmap is encoded, and hdrSdrRatio is
+// unused
+ANDROID_API int ABitmap_compressWithGainmap(const AndroidBitmapInfo* info, ADataSpace dataSpace,
+                                            const void* pixels, const void* gainmapPixels,
+                                            float hdrSdrRatio, AndroidBitmapCompressFormat format,
+                                            int32_t quality, void* userContext,
+                                            AndroidBitmap_CompressWriteFunc);
 /**
  *  Retrieve the native object associated with a HARDWARE Bitmap.
  *
diff --git a/libs/hwui/libhwui.map.txt b/libs/hwui/libhwui.map.txt
index d03ceb471d6c..2414299321a9 100644
--- a/libs/hwui/libhwui.map.txt
+++ b/libs/hwui/libhwui.map.txt
@@ -13,6 +13,7 @@ LIBHWUI { # platform-only /* HWUI isn't current a module, so all of these are st
     ABitmapConfig_getFormatFromConfig;
     ABitmapConfig_getConfigFromFormat;
     ABitmap_compress;
+    ABitmap_compressWithGainmap;
     ABitmap_getHardwareBuffer;
     ACanvas_isSupportedPixelFormat;
     ACanvas_getNativeHandleFromJava;