jpegrecoverymap: Implement RGBA decoding

Bug: b/252835416
Test: manual
Change-Id: Ifa167dc5b75f2d171263969d150a028c8b64f1fb
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
index 39c79c9..5993546 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
@@ -36,11 +36,11 @@
     JpegDecoder();
     ~JpegDecoder();
     /*
-     * Decompresses JPEG image to raw image (YUV420planer or grey-scale) format. After calling
-     * this method, call getDecompressedImage() to get the image.
+     * Decompresses JPEG image to raw image (YUV420planer, grey-scale or RGBA) format. After
+     * calling this method, call getDecompressedImage() to get the image.
      * Returns false if decompressing the image fails.
      */
-    bool decompressImage(const void* image, int length);
+    bool decompressImage(const void* image, int length, bool decodeToRGBA = false);
     /*
      * Returns the decompressed raw image buffer pointer. This method must be called only after
      * calling decompressImage().
@@ -98,10 +98,11 @@
     bool extractEXIF(const void* image, int length);
 
 private:
-    bool decode(const void* image, int length);
+    bool decode(const void* image, int length, bool decodeToRGBA);
     // Returns false if errors occur.
     bool decompress(jpeg_decompress_struct* cinfo, const uint8_t* dest, bool isSingleChannel);
     bool decompressYUV(jpeg_decompress_struct* cinfo, const uint8_t* dest);
+    bool decompressRGBA(jpeg_decompress_struct* cinfo, const uint8_t* dest);
     bool decompressSingleChannel(jpeg_decompress_struct* cinfo, const uint8_t* dest);
     // Process 16 lines of Y and 16 lines of U/V each time.
     // We must pass at least 16 scanlines according to libjpeg documentation.
diff --git a/libs/jpegrecoverymap/jpegdecoder.cpp b/libs/jpegrecoverymap/jpegdecoder.cpp
index c2a8f45..07b527c 100644
--- a/libs/jpegrecoverymap/jpegdecoder.cpp
+++ b/libs/jpegrecoverymap/jpegdecoder.cpp
@@ -93,7 +93,7 @@
 JpegDecoder::~JpegDecoder() {
 }
 
-bool JpegDecoder::decompressImage(const void* image, int length) {
+bool JpegDecoder::decompressImage(const void* image, int length, bool decodeToRGBA) {
     if (image == nullptr || length <= 0) {
         ALOGE("Image size can not be handled: %d", length);
         return false;
@@ -101,7 +101,7 @@
 
     mResultBuffer.clear();
     mXMPBuffer.clear();
-    if (!decode(image, length)) {
+    if (!decode(image, length, decodeToRGBA)) {
         return false;
     }
 
@@ -140,7 +140,7 @@
     return mHeight;
 }
 
-bool JpegDecoder::decode(const void* image, int length) {
+bool JpegDecoder::decode(const void* image, int length, bool decodeToRGBA) {
     jpeg_decompress_struct cinfo;
     jpegr_source_mgr mgr(static_cast<const uint8_t*>(image), length);
     jpegrerror_mgr myerr;
@@ -210,15 +210,26 @@
     mWidth = cinfo.image_width;
     mHeight = cinfo.image_height;
 
-    if (cinfo.jpeg_color_space == JCS_YCbCr) {
-        mResultBuffer.resize(cinfo.image_width * cinfo.image_height * 3 / 2, 0);
-    } else if (cinfo.jpeg_color_space == JCS_GRAYSCALE) {
-        mResultBuffer.resize(cinfo.image_width * cinfo.image_height, 0);
+    if (decodeToRGBA) {
+        if (cinfo.jpeg_color_space == JCS_GRAYSCALE) {
+            // We don't intend to support decoding grayscale to RGBA
+            return false;
+        }
+        // 4 bytes per pixel
+        mResultBuffer.resize(cinfo.image_width * cinfo.image_height * 4);
+        cinfo.out_color_space = JCS_EXT_RGBA;
+    } else {
+        if (cinfo.jpeg_color_space == JCS_YCbCr) {
+            // 1 byte per pixel for Y, 0.5 byte per pixel for U+V
+            mResultBuffer.resize(cinfo.image_width * cinfo.image_height * 3 / 2, 0);
+        } else if (cinfo.jpeg_color_space == JCS_GRAYSCALE) {
+            mResultBuffer.resize(cinfo.image_width * cinfo.image_height, 0);
+        }
+        cinfo.out_color_space = cinfo.jpeg_color_space;
+        cinfo.raw_data_out = TRUE;
     }
 
-    cinfo.raw_data_out = TRUE;
     cinfo.dct_method = JDCT_IFAST;
-    cinfo.out_color_space = cinfo.jpeg_color_space;
 
     jpeg_start_decompress(&cinfo);
 
@@ -292,7 +303,10 @@
     if (isSingleChannel) {
         return decompressSingleChannel(cinfo, dest);
     }
-    return decompressYUV(cinfo, dest);
+    if (cinfo->out_color_space == JCS_EXT_RGBA)
+        return decompressRGBA(cinfo, dest);
+    else
+        return decompressYUV(cinfo, dest);
 }
 
 bool JpegDecoder::getCompressedImageParameters(const void* image, int length,
@@ -331,6 +345,20 @@
     return true;
 }
 
+bool JpegDecoder::decompressRGBA(jpeg_decompress_struct* cinfo, const uint8_t* dest) {
+    JSAMPLE* decodeDst = (JSAMPLE*) dest;
+    uint32_t lines = 0;
+    // TODO: use batches for more effectiveness
+    while (lines < cinfo->image_height) {
+        uint32_t ret = jpeg_read_scanlines(cinfo, &decodeDst, 1);
+        if (ret == 0) {
+            break;
+        }
+        decodeDst += cinfo->image_width * 4;
+        lines++;
+    }
+    return lines == cinfo->image_height;
+}
 
 bool JpegDecoder::decompressYUV(jpeg_decompress_struct* cinfo, const uint8_t* dest) {
 
diff --git a/libs/jpegrecoverymap/recoverymap.cpp b/libs/jpegrecoverymap/recoverymap.cpp
index 53fa8ce..27a756f 100644
--- a/libs/jpegrecoverymap/recoverymap.cpp
+++ b/libs/jpegrecoverymap/recoverymap.cpp
@@ -505,23 +505,37 @@
   if (compressed_jpegr_image == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
-
   // TODO: fill EXIF data
   (void) exif;
 
+  if (request_sdr) {
+    JpegDecoder jpeg_decoder;
+    if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length,
+                                      true)) {
+        return ERROR_JPEGR_DECODE_ERROR;
+    }
+    jpegr_uncompressed_struct uncompressed_rgba_image;
+    uncompressed_rgba_image.data = jpeg_decoder.getDecompressedImagePtr();
+    uncompressed_rgba_image.width = jpeg_decoder.getDecompressedImageWidth();
+    uncompressed_rgba_image.height = jpeg_decoder.getDecompressedImageHeight();
+    memcpy(dest->data, uncompressed_rgba_image.data,
+           uncompressed_rgba_image.width * uncompressed_rgba_image.height * 4);
+    dest->width = uncompressed_rgba_image.width;
+    dest->height = uncompressed_rgba_image.height;
+    return NO_ERROR;
+  }
+
   jpegr_compressed_struct compressed_map;
   jpegr_metadata metadata;
   JPEGR_CHECK(extractRecoveryMap(compressed_jpegr_image, &compressed_map));
 
-
   JpegDecoder jpeg_decoder;
   if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length)) {
     return ERROR_JPEGR_DECODE_ERROR;
   }
 
   JpegDecoder recovery_map_decoder;
-  if (!recovery_map_decoder.decompressImage(compressed_map.data,
-                                    compressed_map.length)) {
+  if (!recovery_map_decoder.decompressImage(compressed_map.data, compressed_map.length)) {
     return ERROR_JPEGR_DECODE_ERROR;
   }
 
@@ -530,26 +544,17 @@
   map.width = recovery_map_decoder.getDecompressedImageWidth();
   map.height = recovery_map_decoder.getDecompressedImageHeight();
 
-
   jpegr_uncompressed_struct uncompressed_yuv_420_image;
   uncompressed_yuv_420_image.data = jpeg_decoder.getDecompressedImagePtr();
   uncompressed_yuv_420_image.width = jpeg_decoder.getDecompressedImageWidth();
   uncompressed_yuv_420_image.height = jpeg_decoder.getDecompressedImageHeight();
 
   if (!getMetadataFromXMP(static_cast<uint8_t*>(jpeg_decoder.getXMPPtr()),
-                                       jpeg_decoder.getXMPSize(), &metadata)) {
+                          jpeg_decoder.getXMPSize(), &metadata)) {
     return ERROR_JPEGR_DECODE_ERROR;
   }
 
-  if (request_sdr) {
-    memcpy(dest->data, uncompressed_yuv_420_image.data,
-            uncompressed_yuv_420_image.width*uncompressed_yuv_420_image.height *3 / 2);
-    dest->width = uncompressed_yuv_420_image.width;
-    dest->height = uncompressed_yuv_420_image.height;
-  } else {
-    JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, &metadata, dest));
-  }
-
+  JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, &metadata, dest));
   return NO_ERROR;
 }