Encode JPEG/R from YuvImage

Test: YuvImageTest
Bug: b/252835416
Change-Id: I010b4498487bf58a0eb1dec3f619fec60b0191aa
diff --git a/libs/hwui/Android.bp b/libs/hwui/Android.bp
index aeead5e..9112b1b 100644
--- a/libs/hwui/Android.bp
+++ b/libs/hwui/Android.bp
@@ -376,7 +376,10 @@
         "jni/text/TextShaper.cpp",
     ],
 
-    header_libs: ["android_graphics_jni_headers"],
+    header_libs: [
+        "android_graphics_jni_headers",
+        "libnativewindow_headers",
+    ],
 
     include_dirs: [
         "external/skia/include/private",
@@ -392,10 +395,14 @@
         "libbase",
         "libcutils",
         "libharfbuzz_ng",
+        "libimage_io",
+        "libjpeg",
+        "libjpegdecoder",
+        "libjpegencoder",
+        "libjpegrecoverymap",
         "liblog",
         "libminikin",
         "libz",
-        "libjpeg",
     ],
 
     static_libs: [
diff --git a/libs/hwui/jni/YuvToJpegEncoder.cpp b/libs/hwui/jni/YuvToJpegEncoder.cpp
index 1c5f126..80bca1f 100644
--- a/libs/hwui/jni/YuvToJpegEncoder.cpp
+++ b/libs/hwui/jni/YuvToJpegEncoder.cpp
@@ -1,3 +1,6 @@
+#undef LOG_TAG
+#define LOG_TAG "YuvToJpegEncoder"
+
 #include "CreateJavaOutputStreamAdaptor.h"
 #include "SkJPEGWriteUtility.h"
 #include "SkStream.h"
@@ -235,6 +238,99 @@
 }
 ///////////////////////////////////////////////////////////////////////////////
 
+using namespace android::recoverymap;
+
+jpegr_color_gamut P010Yuv420ToJpegREncoder::findColorGamut(JNIEnv* env, int aDataSpace) {
+    switch (aDataSpace & ADataSpace::STANDARD_MASK) {
+        case ADataSpace::STANDARD_BT709:
+            return jpegr_color_gamut::JPEGR_COLORGAMUT_BT709;
+        case ADataSpace::STANDARD_DCI_P3:
+            return jpegr_color_gamut::JPEGR_COLORGAMUT_P3;
+        case ADataSpace::STANDARD_BT2020:
+            return jpegr_color_gamut::JPEGR_COLORGAMUT_BT2100;
+        default:
+            jclass IllegalArgumentException = env->FindClass("java/lang/IllegalArgumentException");
+            env->ThrowNew(IllegalArgumentException,
+                    "The requested color gamut is not supported by JPEG/R.");
+    }
+
+    return jpegr_color_gamut::JPEGR_COLORGAMUT_UNSPECIFIED;
+}
+
+jpegr_transfer_function P010Yuv420ToJpegREncoder::findHdrTransferFunction(JNIEnv* env,
+        int aDataSpace) {
+    switch (aDataSpace & ADataSpace::TRANSFER_MASK) {
+        case ADataSpace::TRANSFER_ST2084:
+            return jpegr_transfer_function::JPEGR_TF_PQ;
+        case ADataSpace::TRANSFER_HLG:
+            return jpegr_transfer_function::JPEGR_TF_HLG;
+        default:
+            jclass IllegalArgumentException = env->FindClass("java/lang/IllegalArgumentException");
+            env->ThrowNew(IllegalArgumentException,
+                    "The requested HDR transfer function is not supported by JPEG/R.");
+    }
+
+    return jpegr_transfer_function::JPEGR_TF_UNSPECIFIED;
+}
+
+bool P010Yuv420ToJpegREncoder::encode(JNIEnv* env,
+        SkWStream* stream, void* hdr, int hdrColorSpace, void* sdr, int sdrColorSpace,
+        int width, int height, int jpegQuality) {
+    // Check SDR color space. Now we only support SRGB transfer function
+    if ((sdrColorSpace & ADataSpace::TRANSFER_MASK) !=  ADataSpace::TRANSFER_SRGB) {
+        jclass IllegalArgumentException = env->FindClass("java/lang/IllegalArgumentException");
+        env->ThrowNew(IllegalArgumentException,
+            "The requested SDR color space is not supported. Transfer function must be SRGB");
+        return false;
+    }
+
+    jpegr_color_gamut hdrColorGamut = findColorGamut(env, hdrColorSpace);
+    jpegr_color_gamut sdrColorGamut = findColorGamut(env, sdrColorSpace);
+    jpegr_transfer_function hdrTransferFunction = findHdrTransferFunction(env, hdrColorSpace);
+
+    if (hdrColorGamut == jpegr_color_gamut::JPEGR_COLORGAMUT_UNSPECIFIED
+            || sdrColorGamut == jpegr_color_gamut::JPEGR_COLORGAMUT_UNSPECIFIED
+            || hdrTransferFunction == jpegr_transfer_function::JPEGR_TF_UNSPECIFIED) {
+        return false;
+    }
+
+    RecoveryMap recoveryMap;
+
+    jpegr_uncompressed_struct p010;
+    p010.data = hdr;
+    p010.width = width;
+    p010.height = height;
+    p010.colorGamut = hdrColorGamut;
+
+    jpegr_uncompressed_struct yuv420;
+    yuv420.data = sdr;
+    yuv420.width = width;
+    yuv420.height = height;
+    yuv420.colorGamut = sdrColorGamut;
+
+    jpegr_compressed_struct jpegR;
+    jpegR.maxLength = width * height * sizeof(uint8_t);
+
+    std::unique_ptr<uint8_t[]> jpegr_data = std::make_unique<uint8_t[]>(jpegR.maxLength);
+    jpegR.data = jpegr_data.get();
+
+    if (int success = recoveryMap.encodeJPEGR(&p010, &yuv420,
+            hdrTransferFunction,
+            &jpegR, jpegQuality, nullptr); success != android::OK) {
+        ALOGW("Encode JPEG/R failed, error code: %d.", success);
+        return false;
+    }
+
+    if (!stream->write(jpegR.data, jpegR.length)) {
+        ALOGW("Writing JPEG/R to stream failed.");
+        return false;
+    }
+
+    return true;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+
 static jboolean YuvImage_compressToJpeg(JNIEnv* env, jobject, jbyteArray inYuv,
         jint format, jint width, jint height, jintArray offsets,
         jintArray strides, jint jpegQuality, jobject jstream,
@@ -258,11 +354,34 @@
     delete strm;
     return result;
 }
+
+static jboolean YuvImage_compressToJpegR(JNIEnv* env, jobject, jbyteArray inHdr,
+        jint hdrColorSpace, jbyteArray inSdr, jint sdrColorSpace,
+        jint width, jint height, jint quality, jobject jstream,
+        jbyteArray jstorage) {
+    jbyte* hdr = env->GetByteArrayElements(inHdr, NULL);
+    jbyte* sdr = env->GetByteArrayElements(inSdr, NULL);
+    SkWStream* strm = CreateJavaOutputStreamAdaptor(env, jstream, jstorage);
+    P010Yuv420ToJpegREncoder encoder;
+
+    jboolean result = JNI_FALSE;
+    if (encoder.encode(env, strm, hdr, hdrColorSpace, sdr, sdrColorSpace,
+                       width, height, quality)) {
+        result = JNI_TRUE;
+    }
+
+    env->ReleaseByteArrayElements(inHdr, hdr, 0);
+    env->ReleaseByteArrayElements(inSdr, sdr, 0);
+    delete strm;
+    return result;
+}
 ///////////////////////////////////////////////////////////////////////////////
 
 static const JNINativeMethod gYuvImageMethods[] = {
     {   "nativeCompressToJpeg",  "([BIII[I[IILjava/io/OutputStream;[B)Z",
-        (void*)YuvImage_compressToJpeg }
+        (void*)YuvImage_compressToJpeg },
+    {   "nativeCompressToJpegR",  "([BI[BIIIILjava/io/OutputStream;[B)Z",
+        (void*)YuvImage_compressToJpegR }
 };
 
 int register_android_graphics_YuvImage(JNIEnv* env)
diff --git a/libs/hwui/jni/YuvToJpegEncoder.h b/libs/hwui/jni/YuvToJpegEncoder.h
index a69726b1..3d6d1f3 100644
--- a/libs/hwui/jni/YuvToJpegEncoder.h
+++ b/libs/hwui/jni/YuvToJpegEncoder.h
@@ -1,6 +1,9 @@
 #ifndef _ANDROID_GRAPHICS_YUV_TO_JPEG_ENCODER_H_
 #define _ANDROID_GRAPHICS_YUV_TO_JPEG_ENCODER_H_
 
+#include <android/data_space.h>
+#include <jpegrecoverymap/recoverymap.h>
+
 extern "C" {
     #include "jpeglib.h"
     #include "jerror.h"
@@ -24,7 +27,7 @@
      *
      *  @param stream The jpeg output stream.
      *  @param inYuv The input yuv data.
-     *  @param width Width of the the Yuv data in terms of pixels.
+     *  @param width Width of the Yuv data in terms of pixels.
      *  @param height Height of the Yuv data in terms of pixels.
      *  @param offsets The offsets in each image plane with respect to inYuv.
      *  @param jpegQuality Picture quality in [0, 100].
@@ -71,4 +74,46 @@
             uint8_t* vRows, int rowIndex, int width, int height);
 };
 
+class P010Yuv420ToJpegREncoder {
+public:
+    /** Encode YUV data to jpeg/r,  which is output to a stream.
+     *  This method will call RecoveryMap::EncodeJPEGR() method. If encoding failed,
+     *  Corresponding error code (defined in jpegrerrorcode.h) will be printed and this
+     *  method will be terminated and return false.
+     *
+     *  @param env JNI environment.
+     *  @param stream The jpeg output stream.
+     *  @param hdr The input yuv data (p010 format).
+     *  @param hdrColorSpaceId color space id for the input hdr.
+     *  @param sdr The input yuv data (yuv420p format).
+     *  @param sdrColorSpaceId color space id for the input sdr.
+     *  @param width Width of the Yuv data in terms of pixels.
+     *  @param height Height of the Yuv data in terms of pixels.
+     *  @param jpegQuality Picture quality in [0, 100].
+     *  @return true if successfully compressed the stream.
+     */
+    bool encode(JNIEnv* env,
+            SkWStream* stream, void* hdr, int hdrColorSpace, void* sdr, int sdrColorSpace,
+            int width, int height, int jpegQuality);
+
+    /** Map data space (defined in DataSpace.java and data_space.h) to the color gamut
+     *  used in JPEG/R
+     *
+     *  @param env JNI environment.
+     *  @param aDataSpace data space defined in data_space.h.
+     *  @return color gamut for JPEG/R.
+     */
+    static android::recoverymap::jpegr_color_gamut findColorGamut(JNIEnv* env, int aDataSpace);
+
+    /** Map data space (defined in DataSpace.java and data_space.h) to the transfer function
+     *  used in JPEG/R
+     *
+     *  @param env JNI environment.
+     *  @param aDataSpace data space defined in data_space.h.
+     *  @return color gamut for JPEG/R.
+     */
+    static android::recoverymap::jpegr_transfer_function findHdrTransferFunction(
+            JNIEnv* env, int aDataSpace);
+};
+
 #endif  // _ANDROID_GRAPHICS_YUV_TO_JPEG_ENCODER_H_