Add color space xform support to SkJpegCodec (includes F16!)

src/codec/SkJpegCodec.cpp

Index: src/codec/SkJpegCodec.cpp
diff --git a/src/codec/SkJpegCodec.cpp b/src/codec/SkJpegCodec.cpp
index f4116e324f82ca4e584e9080362735469abde463..64fa04e0ff52cb6f65a04fd06aa05baba010820b 100644
--- a/src/codec/SkJpegCodec.cpp
+++ b/src/codec/SkJpegCodec.cpp
@@ -263,7 +263,8 @@ SkJpegCodec::SkJpegCodec(int width, int height, const SkEncodedInfo& info, SkStr
     : INHERITED(width, height, info, stream, std::move(colorSpace), origin)
     , fDecoderMgr(decoderMgr)
     , fReadyState(decoderMgr->dinfo()->global_state)
-    , fSrcRow(nullptr)
+    , fSwizzleSrcRow(nullptr)
+    , fColorXformSrcRow(nullptr)
     , fSwizzlerSubset(SkIRect::MakeEmpty())
     , fICCData(std::move(iccData))
 {}
@@ -342,8 +343,10 @@ bool SkJpegCodec::onRewind() {
     fDecoderMgr.reset(decoderMgr);
 
     fSwizzler.reset(nullptr);
-    fSrcRow = nullptr;
+    fSwizzleSrcRow = nullptr;
+    fColorXformSrcRow = nullptr;
     fStorage.reset();
+    fColorXform.reset(nullptr);
 
     return true;
 }
@@ -353,22 +356,23 @@ bool SkJpegCodec::onRewind() {
  * image has been implemented
  * Sets the output color space
  */
-bool SkJpegCodec::setOutputColorSpace(const SkImageInfo& dst) {
-    if (kUnknown_SkAlphaType == dst.alphaType()) {
+bool SkJpegCodec::setOutputColorSpace(const SkImageInfo& dstInfo, bool needsColorXform) {
+    if (kUnknown_SkAlphaType == dstInfo.alphaType()) {
         return false;
     }
 
-    if (kOpaque_SkAlphaType != dst.alphaType()) {
+    if (kOpaque_SkAlphaType != dstInfo.alphaType()) {
         SkCodecPrintf("Warning: an opaque image should be decoded as opaque "
                       "- it is being decoded as non-opaque, which will draw slower\n");
     }
 
-    // Check if we will decode to CMYK because a conversion to RGBA is not supported
-    J_COLOR_SPACE colorSpace = fDecoderMgr->dinfo()->jpeg_color_space;
-    bool isCMYK = JCS_CMYK == colorSpace || JCS_YCCK == colorSpace;
+    // Check if we will decode to CMYK.  libjpeg-turbo does not convert CMYK to RGBA, so
+    // we must do it ourselves.
+    J_COLOR_SPACE encodedColorType = fDecoderMgr->dinfo()->jpeg_color_space;
+    bool isCMYK = (JCS_CMYK == encodedColorType || JCS_YCCK == encodedColorType);
 
     // Check for valid color types and set the output color space
-    switch (dst.colorType()) {
+    switch (dstInfo.colorType()) {
         case kRGBA_8888_SkColorType:
             if (isCMYK) {
                 fDecoderMgr->dinfo()->out_color_space = JCS_CMYK;
@@ -379,11 +383,17 @@ bool SkJpegCodec::setOutputColorSpace(const SkImageInfo& dst) {
         case kBGRA_8888_SkColorType:
             if (isCMYK) {
                 fDecoderMgr->dinfo()->out_color_space = JCS_CMYK;
+            } else if (needsColorXform) {
+                fDecoderMgr->dinfo()->out_color_space = JCS_EXT_RGBA;
             } else {
                 fDecoderMgr->dinfo()->out_color_space = JCS_EXT_BGRA;
             }
             return true;
         case kRGB_565_SkColorType:
+            if (needsColorXform) {
+                return false;
+            }
+
             if (isCMYK) {
                 fDecoderMgr->dinfo()->out_color_space = JCS_CMYK;
             } else {
@@ -396,12 +406,26 @@ bool SkJpegCodec::setOutputColorSpace(const SkImageInfo& dst) {
             }
             return true;
         case kGray_8_SkColorType:
-            if (isCMYK) {
+            if (needsColorXform || JCS_GRAYSCALE != encodedColorType) {
+                return false;
+            }
+
+            fDecoderMgr->dinfo()->out_color_space = JCS_GRAYSCALE;
+            return true;
+        case kRGBA_F16_SkColorType:
+            SkASSERT(needsColorXform);
+            if (!dstInfo.colorSpace()) {
+                // Disable F16 in legacy mode.
+                // TODO (msarett):
+                // Alternatively, should this mean "xform to half floats, but don't xform
+                // the gamut"?
                 return false;
+            }
+
+            if (isCMYK) {
+                fDecoderMgr->dinfo()->out_color_space = JCS_CMYK;
             } else {
-                // We will enable decodes to gray even if the image is color because this is

[msarett, 2016/07/22 15:39:46]

This was a dumb feature that nobody uses and currently isn't being tested.

-                // much faster than decoding to color and then converting
-                fDecoderMgr->dinfo()->out_color_space = JCS_GRAYSCALE;
+                fDecoderMgr->dinfo()->out_color_space = JCS_EXT_RGBA;
             }
             return true;
         default:
@@ -450,6 +474,72 @@ bool SkJpegCodec::onDimensionsSupported(const SkISize& size) {
     return true;
 }
 
+static bool needs_color_xform(const SkImageInfo& dstInfo, const SkImageInfo& srcInfo) {
+    // FIXME (msarett):
+    // Do a better check for color space equality.
+    return (kRGBA_F16_SkColorType == dstInfo.colorType()) ||
+           (dstInfo.colorSpace() && (dstInfo.colorSpace() != srcInfo.colorSpace()));
+}
+
+int SkJpegCodec::readRows(const SkImageInfo& dstInfo, void* dst, size_t rowBytes, int count) {
+    int width = dstInfo.width();
+
+    JSAMPLE* decodeDst = (JSAMPLE*) dst;

[msarett, 2016/07/22 15:39:45]

Gladly taking suggestions to make this function more readable and/or efficient.

If we swizzle, it will never be "in place" (since it will probably scale and/or
subset).

We can color xform "in place", unless we must convert to F16 (in which case we
have a fColorXformSrcRow).

Probably the most common/important case will be:
(1) No swizzler
(2) Yes color xform (may be "in place" or not)

+    uint32_t* swizzleDst = (uint32_t*) dst;
+    size_t decodeDstRowBytes = rowBytes;
+    size_t swizzleDstRowBytes = rowBytes;
+    if (fSwizzleSrcRow && fColorXformSrcRow) {
+        decodeDst = (JSAMPLE*) fSwizzleSrcRow;
+        swizzleDst = fColorXformSrcRow;
+        decodeDstRowBytes = 0;
+        swizzleDstRowBytes = 0;
+    } else if (fColorXformSrcRow) {
+        decodeDst = (JSAMPLE*) fColorXformSrcRow;
+        swizzleDst = fColorXformSrcRow;
+        decodeDstRowBytes = 0;
+        swizzleDstRowBytes = 0;
+    } else if (fSwizzleSrcRow) {
+        decodeDst = (JSAMPLE*) fSwizzleSrcRow;
+        decodeDstRowBytes = 0;
+    }
+
+    for (int y = 0; y < count; y++) {
+        uint32_t lines = jpeg_read_scanlines(fDecoderMgr->dinfo(), &decodeDst, 1);
+        sk_msan_mark_initialized(decodeDst, decodeDst + rowBytes, "skbug.com/4550");
+        if (lines != 1) {
+            return y;
+        }
+
+        if (fSwizzler) {
+            fSwizzler->swizzle(swizzleDst, decodeDst);
+        }
+
+        if (fColorXform) {
+            switch (dstInfo.colorType()) {

[msarett, 2016/07/22 15:39:45]

It'd be nice if we could call a function ptr here, but we need all the state on
the color xform object...

+                case kRGBA_8888_SkColorType:
+                    fColorXform->applyToRGBA((uint32_t*) dst, swizzleDst, width);
+                    break;
+                case kBGRA_8888_SkColorType:
+                    fColorXform->applyToBGRA((uint32_t*) dst, swizzleDst, width);
+                    break;
+                case kRGBA_F16_SkColorType:
+                    fColorXform->applyToF16((uint64_t*) dst, swizzleDst, width);
+                    break;
+                default:
+                    SkASSERT(false);
+                    break;
+            }
+
+            dst = SkTAddOffset<void>(dst, rowBytes);
+        }
+
+        decodeDst = SkTAddOffset<JSAMPLE>(decodeDst, decodeDstRowBytes);
+        swizzleDst = SkTAddOffset<uint32_t>(swizzleDst, swizzleDstRowBytes);
+    }
+
+    return count;
+}
+
 /*
  * Performs the jpeg decode
  */
@@ -471,11 +561,19 @@ SkCodec::Result SkJpegCodec::onGetPixels(const SkImageInfo& dstInfo,
     }
 
     // Check if we can decode to the requested destination and set the output color space
-    if (!this->setOutputColorSpace(dstInfo)) {
-        return fDecoderMgr->returnFailure("conversion_possible", kInvalidConversion);
+    bool needsColorXform = needs_color_xform(dstInfo, this->getInfo());
+    if (!this->setOutputColorSpace(dstInfo, needsColorXform)) {
+        return fDecoderMgr->returnFailure("setOutputColorSpace", kInvalidConversion);
+    }
+
+    if (needsColorXform) {
+        fColorXform = SkColorSpaceXform::New(std::move(sk_ref_sp(this->getInfo().colorSpace())),
+                                             std::move(sk_ref_sp(dstInfo.colorSpace())));
+        if (!fColorXform && kRGBA_F16_SkColorType == dstInfo.colorType()) {
+            return kInvalidInput;
+        }
     }
 
-    // Now, given valid output dimensions, we can start the decompress
     if (!jpeg_start_decompress(dinfo)) {
         return fDecoderMgr->returnFailure("startDecompress", kInvalidInput);
     }
@@ -489,39 +587,36 @@ SkCodec::Result SkJpegCodec::onGetPixels(const SkImageInfo& dstInfo,
         this->initializeSwizzler(dstInfo, options);
     }
 
-    // Perform the decode a single row at a time
-    uint32_t dstHeight = dstInfo.height();
+    this->allocateStorage(dstInfo);
 
-    JSAMPLE* dstRow;
-    if (fSwizzler) {
-        // write data to storage row, then sample using swizzler
-        dstRow = fSrcRow;
-    } else {
-        // write data directly to dst
-        dstRow = (JSAMPLE*) dst;
+    int rows = this->readRows(dstInfo, dst, dstRowBytes, dstInfo.height());
+    if (rows < dstInfo.height()) {
+        *rowsDecoded = rows;
+        return fDecoderMgr->returnFailure("Incomplete image data", kIncompleteInput);
     }
 
-    for (uint32_t y = 0; y < dstHeight; y++) {
-        // Read rows of the image
-        uint32_t lines = jpeg_read_scanlines(dinfo, &dstRow, 1);
-        sk_msan_mark_initialized(dstRow, dstRow + dstRowBytes, "skbug.com/4550");
+    return kSuccess;
+}
 
-        // If we cannot read enough rows, assume the input is incomplete
-        if (lines != 1) {
-            *rowsDecoded = y;
-            return fDecoderMgr->returnFailure("Incomplete image data", kIncompleteInput);
-        }
+void SkJpegCodec::allocateStorage(const SkImageInfo& dstInfo) {
+    size_t swizzleBytes = 0;
+    if (fSwizzler) {
+        swizzleBytes = get_row_bytes(fDecoderMgr->dinfo());
+    }
 
-        if (fSwizzler) {
-            // use swizzler to sample row
-            fSwizzler->swizzle(dst, dstRow);
-            dst = SkTAddOffset<JSAMPLE>(dst, dstRowBytes);
-        } else {
-            dstRow = SkTAddOffset<JSAMPLE>(dstRow, dstRowBytes);
-        }
+    size_t xformBytes = 0;
+    if (kRGBA_F16_SkColorType == dstInfo.colorType()) {
+        SkASSERT(fColorXform);
+        xformBytes = dstInfo.width() * sizeof(SkColorSpaceXform::RGBA32);
     }
 
-    return kSuccess;
+    size_t totalBytes = swizzleBytes + xformBytes;
+    if (totalBytes > 0) {
+        fStorage.reset(totalBytes);
+        fSwizzleSrcRow = (swizzleBytes > 0) ? fStorage.get() : nullptr;
+        fColorXformSrcRow = (xformBytes > 0) ?
+                SkTAddOffset<uint32_t>(fStorage.get(), swizzleBytes) : nullptr;
+    }
 }
 
 void SkJpegCodec::initializeSwizzler(const SkImageInfo& dstInfo, const Options& options) {
@@ -538,9 +633,9 @@ void SkJpegCodec::initializeSwizzler(const SkImageInfo& dstInfo, const Options&
             break;
         case JCS_CMYK:
             preSwizzled = false;
-            swizzlerInfo = SkEncodedInfo::Make(
-                    SkEncodedInfo::kInvertedCMYK_Color, swizzlerInfo.alpha(),
-                    swizzlerInfo.bitsPerComponent());
+            swizzlerInfo = SkEncodedInfo::Make(SkEncodedInfo::kInvertedCMYK_Color,
+                                               swizzlerInfo.alpha(),
+                                               swizzlerInfo.bitsPerComponent());
             break;
         default:
             break;
@@ -558,17 +653,16 @@ void SkJpegCodec::initializeSwizzler(const SkImageInfo& dstInfo, const Options&
     fSwizzler.reset(SkSwizzler::CreateSwizzler(swizzlerInfo, nullptr, dstInfo, swizzlerOptions,
                                                nullptr, preSwizzled));
     SkASSERT(fSwizzler);
-    fStorage.reset(get_row_bytes(fDecoderMgr->dinfo()));
-    fSrcRow = fStorage.get();
 }
 
 SkSampler* SkJpegCodec::getSampler(bool createIfNecessary) {
     if (!createIfNecessary || fSwizzler) {
-        SkASSERT(!fSwizzler || (fSrcRow && fStorage.get() == fSrcRow));
+        SkASSERT(!fSwizzler || (fSwizzleSrcRow && fStorage.get() == fSwizzleSrcRow));
         return fSwizzler;
     }
 
     this->initializeSwizzler(this->dstInfo(), this->options());
+    this->allocateStorage(this->dstInfo());
     return fSwizzler;
 }
 
@@ -581,11 +675,19 @@ SkCodec::Result SkJpegCodec::onStartScanlineDecode(const SkImageInfo& dstInfo,
     }
 
     // Check if we can decode to the requested destination and set the output color space
-    if (!this->setOutputColorSpace(dstInfo)) {
+    bool needsColorXform = needs_color_xform(dstInfo, this->getInfo());
+    if (!this->setOutputColorSpace(dstInfo, needsColorXform)) {
         return kInvalidConversion;
     }
 
-    // Now, given valid output dimensions, we can start the decompress
+    if (needsColorXform) {
+        fColorXform = SkColorSpaceXform::New(std::move(sk_ref_sp(this->getInfo().colorSpace())),
+                                             std::move(sk_ref_sp(dstInfo.colorSpace())));
+        if (!fColorXform && kRGBA_F16_SkColorType == dstInfo.colorType()) {
+            return kInvalidInput;
+        }
+    }
+
     if (!jpeg_start_decompress(fDecoderMgr->dinfo())) {
         SkCodecPrintf("start decompress failed\n");
         return kInvalidInput;
@@ -646,6 +748,8 @@ SkCodec::Result SkJpegCodec::onStartScanlineDecode(const SkImageInfo& dstInfo,
     }
 #endif
 
+    this->allocateStorage(dstInfo);
+
     return kSuccess;
 }
 
@@ -654,36 +758,14 @@ int SkJpegCodec::onGetScanlines(void* dst, int count, size_t dstRowBytes) {
     if (setjmp(fDecoderMgr->getJmpBuf())) {
         return fDecoderMgr->returnFailure("setjmp", kInvalidInput);
     }
-    // Read rows one at a time
-    JSAMPLE* dstRow;
-    size_t srcRowBytes = get_row_bytes(fDecoderMgr->dinfo());
-    if (fSwizzler) {
-        // write data to storage row, then sample using swizzler
-        dstRow = fSrcRow;
-    } else {
-        // write data directly to dst
-        SkASSERT(count == 1 || dstRowBytes >= srcRowBytes);
-        dstRow = (JSAMPLE*) dst;
-    }
 
-    for (int y = 0; y < count; y++) {
-        // Read row of the image
-        uint32_t rowsDecoded = jpeg_read_scanlines(fDecoderMgr->dinfo(), &dstRow, 1);
-        sk_msan_mark_initialized(dstRow, dstRow + srcRowBytes, "skbug.com/4550");
-        if (rowsDecoded != 1) {
-            fDecoderMgr->dinfo()->output_scanline = this->dstInfo().height();
-            return y;
-        }
-
-        if (fSwizzler) {
-            // use swizzler to sample row
-            fSwizzler->swizzle(dst, dstRow);
-            dst = SkTAddOffset<JSAMPLE>(dst, dstRowBytes);
-        } else {
-            dstRow = SkTAddOffset<JSAMPLE>(dstRow, dstRowBytes);
-        }
+    int rows = this->readRows(this->dstInfo(), dst, dstRowBytes, count);
+    if (rows < count) {
+        // This allows us to skip calling jpeg_finish_decompress().
+        fDecoderMgr->dinfo()->output_scanline = this->dstInfo().height();
     }
-    return count;
+
+    return rows;
 }
 
 bool SkJpegCodec::onSkipScanlines(int count) {
@@ -695,13 +777,13 @@ bool SkJpegCodec::onSkipScanlines(int count) {
 #ifdef TURBO_HAS_SKIP
     return (uint32_t) count == jpeg_skip_scanlines(fDecoderMgr->dinfo(), count);
 #else
-    if (!fSrcRow) {
+    if (!fSwizzleSrcRow) {
         fStorage.reset(get_row_bytes(fDecoderMgr->dinfo()));
-        fSrcRow = fStorage.get();
+        fSwizzleSrcRow = fStorage.get();
     }
 
     for (int y = 0; y < count; y++) {
-        if (1 != jpeg_read_scanlines(fDecoderMgr->dinfo(), &fSrcRow, 1)) {
+        if (1 != jpeg_read_scanlines(fDecoderMgr->dinfo(), &fSwizzleSrcRow, 1)) {
             return false;
         }
     }