Encode JPEG/R from YuvImage

Test: YuvImageTest
Bug: b/252835416
Change-Id: I010b4498487bf58a0eb1dec3f619fec60b0191aa
diff --git a/core/api/current.txt b/core/api/current.txt
index 801da28..90f9dba 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -16338,7 +16338,10 @@
 
   public class YuvImage {
     ctor public YuvImage(byte[], int, int, int, int[]);
+    ctor public YuvImage(@NonNull byte[], int, int, int, @Nullable int[], @NonNull android.graphics.ColorSpace);
     method public boolean compressToJpeg(android.graphics.Rect, int, java.io.OutputStream);
+    method public boolean compressToJpegR(@NonNull android.graphics.YuvImage, int, @NonNull java.io.OutputStream);
+    method @NonNull public android.graphics.ColorSpace getColorSpace();
     method public int getHeight();
     method public int[] getStrides();
     method public int getWidth();
diff --git a/graphics/java/android/graphics/YuvImage.java b/graphics/java/android/graphics/YuvImage.java
index af3f276..6b5238b 100644
--- a/graphics/java/android/graphics/YuvImage.java
+++ b/graphics/java/android/graphics/YuvImage.java
@@ -16,6 +16,8 @@
 
 package android.graphics;
 
+import android.annotation.NonNull;
+import android.annotation.Nullable;
 import java.io.OutputStream;
 
 /**
@@ -63,7 +65,70 @@
     private int mHeight;
 
     /**
-     * Construct an YuvImage.
+     *  The color space of the image, defaults to SRGB
+     */
+    @NonNull private ColorSpace mColorSpace;
+
+    /**
+     * Array listing all supported ImageFormat that are supported by this class
+     */
+    private final static String[] sSupportedFormats =
+            {"NV21", "YUY2", "YCBCR_P010", "YUV_420_888"};
+
+    private static String printSupportedFormats() {
+        StringBuilder sb = new StringBuilder();
+        for (int i = 0; i < sSupportedFormats.length; ++i) {
+            sb.append(sSupportedFormats[i]);
+            if (i != sSupportedFormats.length - 1) {
+                sb.append(", ");
+            }
+        }
+        return sb.toString();
+    }
+
+    /**
+     * Array listing all supported HDR ColorSpaces that are supported by JPEG/R encoding
+     */
+    private final static ColorSpace.Named[] sSupportedJpegRHdrColorSpaces = {
+        ColorSpace.Named.BT2020_HLG,
+        ColorSpace.Named.BT2020_PQ
+    };
+
+    /**
+     * Array listing all supported SDR ColorSpaces that are supported by JPEG/R encoding
+     */
+    private final static ColorSpace.Named[] sSupportedJpegRSdrColorSpaces = {
+        ColorSpace.Named.SRGB,
+        ColorSpace.Named.DISPLAY_P3
+    };
+
+    private static String printSupportedJpegRColorSpaces(boolean isHdr) {
+        ColorSpace.Named[] colorSpaces = isHdr ? sSupportedJpegRHdrColorSpaces :
+                sSupportedJpegRSdrColorSpaces;
+        StringBuilder sb = new StringBuilder();
+        for (int i = 0; i < colorSpaces.length; ++i) {
+            sb.append(ColorSpace.get(colorSpaces[i]).getName());
+            if (i != colorSpaces.length - 1) {
+                sb.append(", ");
+            }
+        }
+        return sb.toString();
+    }
+
+    private static boolean isSupportedJpegRColorSpace(boolean isHdr, int colorSpace) {
+        ColorSpace.Named[] colorSpaces = isHdr ? sSupportedJpegRHdrColorSpaces :
+              sSupportedJpegRSdrColorSpaces;
+        for (ColorSpace.Named cs : colorSpaces) {
+            if (cs.ordinal() == colorSpace) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+
+    /**
+     * Construct an YuvImage. Use SRGB for as default {@link ColorSpace}.
      *
      * @param yuv     The YUV data. In the case of more than one image plane, all the planes must be
      *                concatenated into a single byte array.
@@ -77,11 +142,33 @@
      *                null.
      */
     public YuvImage(byte[] yuv, int format, int width, int height, int[] strides) {
+        this(yuv, format, width, height, strides, ColorSpace.get(ColorSpace.Named.SRGB));
+    }
+
+    /**
+     * Construct an YuvImage.
+     *
+     * @param yuv        The YUV data. In the case of more than one image plane, all the planes
+     *                   must be concatenated into a single byte array.
+     * @param format     The YUV data format as defined in {@link ImageFormat}.
+     * @param width      The width of the YuvImage.
+     * @param height     The height of the YuvImage.
+     * @param strides    (Optional) Row bytes of each image plane. If yuv contains padding, the
+     *                   stride of each image must be provided. If strides is null, the method
+     *                   assumes no padding and derives the row bytes by format and width itself.
+     * @param colorSpace The YUV image color space as defined in {@link ColorSpace}.
+     *                   If the parameter is null, SRGB will be set as the default value.
+     * @throws IllegalArgumentException if format is not support; width or height <= 0; or yuv is
+     *                null.
+     */
+    public YuvImage(@NonNull byte[] yuv, int format, int width, int height,
+            @Nullable int[] strides, @NonNull ColorSpace colorSpace) {
         if (format != ImageFormat.NV21 &&
-                format != ImageFormat.YUY2) {
+                format != ImageFormat.YUY2 &&
+                format != ImageFormat.YCBCR_P010 &&
+                format != ImageFormat.YUV_420_888) {
             throw new IllegalArgumentException(
-                    "only support ImageFormat.NV21 " +
-                    "and ImageFormat.YUY2 for now");
+                    "only supports the following ImageFormat:" + printSupportedFormats());
         }
 
         if (width <= 0  || height <= 0) {
@@ -93,6 +180,10 @@
             throw new IllegalArgumentException("yuv cannot be null");
         }
 
+        if (colorSpace == null) {
+            throw new IllegalArgumentException("ColorSpace cannot be null");
+        }
+
         if (strides == null) {
             mStrides = calculateStrides(width, format);
         } else {
@@ -103,12 +194,13 @@
         mFormat = format;
         mWidth = width;
         mHeight = height;
+        mColorSpace = colorSpace;
     }
 
     /**
      * Compress a rectangle region in the YuvImage to a jpeg.
-     * Only ImageFormat.NV21 and ImageFormat.YUY2
-     * are supported for now.
+     * For image format, only ImageFormat.NV21 and ImageFormat.YUY2 are supported.
+     * For color space, only SRGB is supported.
      *
      * @param rectangle The rectangle region to be compressed. The medthod checks if rectangle is
      *                  inside the image. Also, the method modifies rectangle if the chroma pixels
@@ -117,10 +209,18 @@
      *                  small size, 100 meaning compress for max quality.
      * @param stream    OutputStream to write the compressed data.
      * @return          True if the compression is successful.
-     * @throws IllegalArgumentException if rectangle is invalid; quality is not within [0,
-     *                  100]; or stream is null.
+     * @throws IllegalArgumentException if rectangle is invalid; color space or image format
+     *                  is not supported; quality is not within [0, 100]; or stream is null.
      */
     public boolean compressToJpeg(Rect rectangle, int quality, OutputStream stream) {
+        if (mFormat != ImageFormat.NV21 && mFormat != ImageFormat.YUY2) {
+            throw new IllegalArgumentException(
+                    "Only ImageFormat.NV21 and ImageFormat.YUY2 are supported.");
+        }
+        if (mColorSpace.getId() != ColorSpace.Named.SRGB.ordinal()) {
+            throw new IllegalArgumentException("Only SRGB color space is supported.");
+        }
+
         Rect wholeImage = new Rect(0, 0, mWidth, mHeight);
         if (!wholeImage.contains(rectangle)) {
             throw new IllegalArgumentException(
@@ -143,6 +243,70 @@
                 new byte[WORKING_COMPRESS_STORAGE]);
     }
 
+    /**
+     * Compress the HDR image into JPEG/R format.
+     *
+     * Sample usage:
+     *     hdr_image.compressToJpegR(sdr_image, 90, stream);
+     *
+     * For the SDR image, only YUV_420_888 image format is supported, and the following
+     * color spaces are supported:
+     *     ColorSpace.Named.SRGB,
+     *     ColorSpace.Named.DISPLAY_P3
+     *
+     * For the HDR image, only YCBCR_P010 image format is supported, and the following
+     * color spaces are supported:
+     *     ColorSpace.Named.BT2020_HLG,
+     *     ColorSpace.Named.BT2020_PQ
+     *
+     * @param sdr       The SDR image, only ImageFormat.YUV_420_888 is supported.
+     * @param quality   Hint to the compressor, 0-100. 0 meaning compress for
+     *                  small size, 100 meaning compress for max quality.
+     * @param stream    OutputStream to write the compressed data.
+     * @return          True if the compression is successful.
+     * @throws IllegalArgumentException if input images are invalid; quality is not within [0,
+     *                  100]; or stream is null.
+     */
+    public boolean compressToJpegR(@NonNull YuvImage sdr, int quality,
+            @NonNull OutputStream stream) {
+        if (sdr == null) {
+            throw new IllegalArgumentException("SDR input cannot be null");
+        }
+
+        if (mData.length == 0 || sdr.getYuvData().length == 0) {
+            throw new IllegalArgumentException("Input images cannot be empty");
+        }
+
+        if (mFormat != ImageFormat.YCBCR_P010 || sdr.getYuvFormat() != ImageFormat.YUV_420_888) {
+            throw new IllegalArgumentException(
+                "only support ImageFormat.YCBCR_P010 and ImageFormat.YUV_420_888");
+        }
+
+        if (sdr.getWidth() != mWidth || sdr.getHeight() != mHeight) {
+            throw new IllegalArgumentException("HDR and SDR resolution mismatch");
+        }
+
+        if (quality < 0 || quality > 100) {
+            throw new IllegalArgumentException("quality must be 0..100");
+        }
+
+        if (stream == null) {
+            throw new IllegalArgumentException("stream cannot be null");
+        }
+
+        if (!isSupportedJpegRColorSpace(true, mColorSpace.getId()) ||
+                !isSupportedJpegRColorSpace(false, sdr.getColorSpace().getId())) {
+            throw new IllegalArgumentException("Not supported color space. "
+                + "SDR only supports: " + printSupportedJpegRColorSpaces(false)
+                + "HDR only supports: " + printSupportedJpegRColorSpaces(true));
+        }
+
+      return nativeCompressToJpegR(mData, mColorSpace.getDataSpace(),
+                                   sdr.getYuvData(), sdr.getColorSpace().getDataSpace(),
+                                   mWidth, mHeight, quality, stream,
+                                   new byte[WORKING_COMPRESS_STORAGE]);
+  }
+
 
    /**
      * @return the YUV data.
@@ -179,6 +343,12 @@
         return mHeight;
     }
 
+
+    /**
+     * @return the color space of the image.
+     */
+    public @NonNull ColorSpace getColorSpace() { return mColorSpace; }
+
     int[] calculateOffsets(int left, int top) {
         int[] offsets = null;
         if (mFormat == ImageFormat.NV21) {
@@ -198,17 +368,23 @@
 
     private int[] calculateStrides(int width, int format) {
         int[] strides = null;
-        if (format == ImageFormat.NV21) {
+        switch (format) {
+          case ImageFormat.NV21:
             strides = new int[] {width, width};
             return strides;
-        }
-
-        if (format == ImageFormat.YUY2) {
+          case ImageFormat.YCBCR_P010:
+            strides = new int[] {width * 2, width * 2};
+            return strides;
+          case ImageFormat.YUV_420_888:
+            strides = new int[] {width, (width + 1) / 2, (width + 1) / 2};
+            return strides;
+          case ImageFormat.YUY2:
             strides = new int[] {width * 2};
             return strides;
+          default:
+            throw new IllegalArgumentException(
+                "only supports the following ImageFormat:" + printSupportedFormats());
         }
-
-        return strides;
     }
 
    private void adjustRectangle(Rect rect) {
@@ -237,4 +413,8 @@
     private static native boolean nativeCompressToJpeg(byte[] oriYuv,
             int format, int width, int height, int[] offsets, int[] strides,
             int quality, OutputStream stream, byte[] tempStorage);
+
+    private static native boolean nativeCompressToJpegR(byte[] hdr, int hdrColorSpaceId,
+            byte[] sdr, int sdrColorSpaceId, int width, int height, int quality,
+            OutputStream stream, byte[] tempStorage);
 }
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_