4 files changed, 40 insertions, 4 deletions

diff --git a/libs/hwui/hwui/ImageDecoder.cpp b/libs/hwui/hwui/ImageDecoder.cpp
index a6c4e9db7280..4b2857f6c290 100644
--- a/libs/hwui/hwui/ImageDecoder.cpp
+++ b/libs/hwui/hwui/ImageDecoder.cpp
@@ -31,7 +31,7 @@ ImageDecoder::ImageDecoder(std::unique_ptr<SkAndroidCodec> codec, sk_sp<SkPngChu
     , mDecodeSize(mTargetSize)
     , mOutColorType(mCodec->computeOutputColorType(kN32_SkColorType))
     , mUnpremultipliedRequired(false)
-    , mOutColorSpace(mCodec->getInfo().refColorSpace())
+    , mOutColorSpace(mCodec->computeOutputColorSpace(mOutColorType, nullptr))
     , mSampleSize(1)
 {
 }
@@ -111,7 +111,6 @@ bool ImageDecoder::setOutColorType(SkColorType colorType) {
             if (!gray()) {
                 return false;
             }
-            mOutColorSpace = nullptr;
             break;
         case kN32_SkColorType:
             break;
@@ -137,9 +136,15 @@ void ImageDecoder::setOutColorSpace(sk_sp<SkColorSpace> colorSpace) {
     mOutColorSpace = std::move(colorSpace);
 }
 
+sk_sp<SkColorSpace> ImageDecoder::getOutputColorSpace() const {
+    // kGray_8 is used for ALPHA_8, which ignores the color space.
+    return mOutColorType == kGray_8_SkColorType ? nullptr : mOutColorSpace;
+}
+
+
 SkImageInfo ImageDecoder::getOutputInfo() const {
     SkISize size = mCropRect ? mCropRect->size() : mTargetSize;
-    return SkImageInfo::Make(size, mOutColorType, getOutAlphaType(), mOutColorSpace);
+    return SkImageInfo::Make(size, mOutColorType, getOutAlphaType(), getOutputColorSpace());
 }
 
 bool ImageDecoder::opaque() const {
@@ -154,7 +159,7 @@ SkCodec::Result ImageDecoder::decode(void* pixels, size_t rowBytes) {
     void* decodePixels = pixels;
     size_t decodeRowBytes = rowBytes;
     auto decodeInfo = SkImageInfo::Make(mDecodeSize, mOutColorType, getOutAlphaType(),
-                                        mOutColorSpace);
+                                        getOutputColorSpace());
     // Used if we need a temporary before scaling or subsetting.
     // FIXME: Use scanline decoding on only a couple lines to save memory. b/70709380.
     SkBitmap tmp;
diff --git a/libs/hwui/hwui/ImageDecoder.h b/libs/hwui/hwui/ImageDecoder.h
index 96f97e5421f3..0c99f84cbb72 100644
--- a/libs/hwui/hwui/ImageDecoder.h
+++ b/libs/hwui/hwui/ImageDecoder.h
@@ -66,6 +66,7 @@ private:
     ImageDecoder& operator=(const ImageDecoder&) = delete;
 
     SkAlphaType getOutAlphaType() const;
+    sk_sp<SkColorSpace> getOutputColorSpace() const;
 };
 
 } // namespace android
diff --git a/native/graphics/jni/imagedecoder.cpp b/native/graphics/jni/imagedecoder.cpp
index bd5914a55042..1c45ea6aaecc 100644
--- a/native/graphics/jni/imagedecoder.cpp
+++ b/native/graphics/jni/imagedecoder.cpp
@@ -18,12 +18,14 @@
 
 #include <android/asset_manager.h>
 #include <android/bitmap.h>
+#include <android/data_space.h>
 #include <android/imagedecoder.h>
 #include <android/graphics/MimeType.h>
 #include <android/rect.h>
 #include <hwui/ImageDecoder.h>
 #include <log/log.h>
 #include <SkAndroidCodec.h>
+#include <utils/Color.h>
 
 #include <fcntl.h>
 #include <optional>
@@ -165,6 +167,18 @@ int AImageDecoder_setAndroidBitmapFormat(AImageDecoder* decoder, int32_t format)
             ? ANDROID_IMAGE_DECODER_SUCCESS : ANDROID_IMAGE_DECODER_INVALID_CONVERSION;
 }
 
+int AImageDecoder_setDataSpace(AImageDecoder* decoder, int32_t dataspace) {
+    sk_sp<SkColorSpace> cs = uirenderer::DataSpaceToColorSpace((android_dataspace)dataspace);
+    // 0 is ADATASPACE_UNKNOWN. We need an explicit request for an ADataSpace.
+    if (!decoder || !dataspace || !cs) {
+        return ANDROID_IMAGE_DECODER_BAD_PARAMETER;
+    }
+
+    ImageDecoder* imageDecoder = toDecoder(decoder);
+    imageDecoder->setOutColorSpace(std::move(cs));
+    return ANDROID_IMAGE_DECODER_SUCCESS;
+}
+
 const AImageDecoderHeaderInfo* AImageDecoder_getHeaderInfo(const AImageDecoder* decoder) {
     return reinterpret_cast<const AImageDecoderHeaderInfo*>(decoder);
 }
@@ -201,6 +215,20 @@ bool AImageDecoderHeaderInfo_isAnimated(const AImageDecoderHeaderInfo* info) {
     return toDecoder(info)->mCodec->codec()->getFrameCount() > 1;
 }
 
+int32_t AImageDecoderHeaderInfo_getDataSpace(const AImageDecoderHeaderInfo* info) {
+    if (!info) {
+        return ANDROID_IMAGE_DECODER_BAD_PARAMETER;
+    }
+
+    // Note: This recomputes the data space because it's possible the client has
+    // changed the output color space, so we cannot rely on it. Alternatively,
+    // we could store the ADataSpace in the ImageDecoder.
+    const ImageDecoder* imageDecoder = toDecoder(info);
+    SkColorType colorType = imageDecoder->mCodec->computeOutputColorType(kN32_SkColorType);
+    sk_sp<SkColorSpace> colorSpace = imageDecoder->mCodec->computeOutputColorSpace(colorType);
+    return uirenderer::ColorSpaceToADataSpace(colorSpace.get(), colorType);
+}
+
 // FIXME: Share with getFormat in android_bitmap.cpp?
 static AndroidBitmapFormat getFormat(SkColorType colorType) {
     switch (colorType) {
diff --git a/native/graphics/jni/libjnigraphics.map.txt b/native/graphics/jni/libjnigraphics.map.txt
index 31a1b9c18ee8..1b396b893f4e 100644
--- a/native/graphics/jni/libjnigraphics.map.txt
+++ b/native/graphics/jni/libjnigraphics.map.txt
@@ -6,6 +6,7 @@ LIBJNIGRAPHICS {
     AImageDecoder_delete; # introduced=30
     AImageDecoder_setAndroidBitmapFormat; # introduced=30
     AImageDecoder_setUnpremultipliedRequired; # introduced=30
+    AImageDecoder_setDataSpace; # introduced=30
     AImageDecoder_getHeaderInfo; # introduced=30
     AImageDecoder_getMinimumStride; # introduced=30
     AImageDecoder_decodeImage; # introduced=30
@@ -18,6 +19,7 @@ LIBJNIGRAPHICS {
     AImageDecoderHeaderInfo_getAlphaFlags; # introduced=30
     AImageDecoderHeaderInfo_isAnimated; # introduced=30
     AImageDecoderHeaderInfo_getAndroidBitmapFormat; # introduced=30
+    AImageDecoderHeaderInfo_getDataSpace; # introduced=30
     AndroidBitmap_getInfo;
     AndroidBitmap_getDataSpace;
     AndroidBitmap_lockPixels;