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;

[mtklein, 2016/07/27 13:51:06]

// Our color transformation code requires RGBA order inputs, but it'll swizzle
to BGRA for us.

[msarett, 2016/07/27 15:13:50]

Done.

             } 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.

[mtklein, 2016/07/27 13:51:06]

Yeah, the new standard line is

// F16 makes no sense without a color space.
return false;

[msarett, 2016/07/27 15:13:50]

Deleting this whole block.  It's unnecessary because SkColorSpaceXform::New()
will return nullptr and we will catch this elsewhere.

+                // 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
-                // 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()) ||

[mtklein, 2016/07/27 13:51:06]

?

Is this, we know we need to end up in F16 by the time we're done, and the only
way to get there is by calling the colorspace transform code, even if the
transform itself will be an identity transform?

(... pipeline ...)

[msarett, 2016/07/27 15:13:50]

The transform will never be an identity transform because we need to convert to
linear gamma.

The gamut transform may be identity.  As an optimization, we should catch this
and skip the step.

+           (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;

[mtklein, 2016/07/27 13:51:05]

Some comments about what aliases what, when, and why it's safe would be nice
here.

[msarett, 2016/07/27 15:13:50]

SGTM

+    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) {

[mtklein, 2016/07/27 13:51:06]

// i.e. == 0 ?

[msarett, 2016/07/27 15:13:50]

Done.

+            return y;
+        }
+
+        if (fSwizzler) {
+            fSwizzler->swizzle(swizzleDst, decodeDst);

[mtklein, 2016/07/27 13:51:05]

... oh hey, this would be a good pipeline stage too ...

[msarett, 2016/07/27 15:13:50]

The swizzler performs sampling and subsetting.  It's potentially possible but
would be a bit tricky.

+        }
+
+        if (fColorXform) {
+            switch (dstInfo.colorType()) {
+                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())));

[mtklein, 2016/07/27 13:51:06]

Aren't these moves redundant?

[msarett, 2016/07/27 15:13:50]

Because we ref them as parameters?  Sure, seems like they ought to be. 
Removing.

+        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) {

[mtklein, 2016/07/27 13:51:06]

Hmm.  I feel like I've read this before.  Think this chunk of logic wants to be
its own method called here and in onGetPixels()?

[msarett, 2016/07/27 15:13:50]

Done.

+        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;
         }
     }