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

src/opts/SkColorXform_opts.h

Index: src/opts/SkColorXform_opts.h
diff --git a/src/opts/SkColorXform_opts.h b/src/opts/SkColorXform_opts.h
index b4eb9a2552b556563e1fd34f87a69d91a9a93750..b3da55c1fdd758d2b3b7d74ea41b050b5918eb18 100644
--- a/src/opts/SkColorXform_opts.h
+++ b/src/opts/SkColorXform_opts.h
@@ -16,14 +16,19 @@
 
 namespace SK_OPTS_NS {
 
-static Sk4i linear_to_2dot2(const Sk4f& x) {
+// Strange that we need a wrapper on SkNx_cast to use as a function ptr.
+static Sk4i Sk4f_trunc(const Sk4f& x) {
+    return SkNx_cast<int>(x);
+}
+
+static Sk4f linear_to_2dot2(const Sk4f& x) {
     // x^(29/64) is a very good approximation of the true value, x^(1/2.2).
     auto x2  = x.rsqrt(),                            // x^(-1/2)
          x32 = x2.rsqrt().rsqrt().rsqrt().rsqrt(),   // x^(-1/32)
          x64 = x32.rsqrt();                          // x^(+1/64)
 
     // 29 = 32 - 2 - 1
-    return Sk4f_round(sk_clamp_0_255(255.0f * x2.invert() * x32 * x64.invert()));
+    return 255.0f * x2.invert() * x32 * x64.invert();
 }
 
 enum DstGamma {
@@ -36,13 +41,23 @@ enum DstGamma {
     kLinear_DstGamma,
 };
 
-template <DstGamma kDstGamma>
+template <DstGamma kDstGamma, bool kSwapRB>
 static void color_xform_RGB1(void* dst, const uint32_t* src, int len,
                              const float* const srcTables[3], const float matrix[16],
                              const uint8_t* const dstTables[3]) {
-    Sk4f rXgXbX = Sk4f::Load(matrix + 0),
-         rYgYbY = Sk4f::Load(matrix + 4),
-         rZgZbZ = Sk4f::Load(matrix + 8);
+    int kRShift = 0;
+    int kGShift = 8;
+    int kBShift = 16;
+    int kAShift = 24;
+    if (kSwapRB) {
+        kBShift = 0;
+        kRShift = 16;
+    }
+
+    Sk4f rXgXbX = Sk4f::Load(matrix +  0),
+         rYgYbY = Sk4f::Load(matrix +  4),
+         rZgZbZ = Sk4f::Load(matrix +  8),
+         rTgTbT = Sk4f::Load(matrix + 12);
 
     if (len >= 4) {
         Sk4f reds, greens, blues;
@@ -65,25 +80,32 @@ static void color_xform_RGB1(void* dst, const uint32_t* src, int len,
 
         Sk4f dstReds, dstGreens, dstBlues;
         auto transform_4 = [&reds, &greens, &blues, &dstReds, &dstGreens, &dstBlues, &rXgXbX,
-                            &rYgYbY, &rZgZbZ] {
-            dstReds   = rXgXbX[0]*reds + rYgYbY[0]*greens + rZgZbZ[0]*blues;
-            dstGreens = rXgXbX[1]*reds + rYgYbY[1]*greens + rZgZbZ[1]*blues;
-            dstBlues  = rXgXbX[2]*reds + rYgYbY[2]*greens + rZgZbZ[2]*blues;
+                            &rYgYbY, &rZgZbZ, &rTgTbT] {
+            dstReds   = rXgXbX[0]*reds + rYgYbY[0]*greens + rZgZbZ[0]*blues + rTgTbT[0];
+            dstGreens = rXgXbX[1]*reds + rYgYbY[1]*greens + rZgZbZ[1]*blues + rTgTbT[1];
+            dstBlues  = rXgXbX[2]*reds + rYgYbY[2]*greens + rZgZbZ[2]*blues + rTgTbT[2];
         };
 
-        auto store_4 = [&dstReds, &dstGreens, &dstBlues, &dst, &dstTables] {
+        auto store_4 = [&dstReds, &dstGreens, &dstBlues, &dst, &dstTables, kRShift, kGShift,
+                        kBShift, kAShift] {
             if (kSRGB_DstGamma == kDstGamma || k2Dot2_DstGamma == kDstGamma) {
-                Sk4i (*linear_to_curve)(const Sk4f&) =
-                        (kSRGB_DstGamma == kDstGamma) ? sk_linear_to_srgb : linear_to_2dot2;
-
-                auto reds   = linear_to_curve(dstReds);
-                auto greens = linear_to_curve(dstGreens);
-                auto blues  = linear_to_curve(dstBlues);
-
-                auto rgba = (reds       << SK_R32_SHIFT)
-                          | (greens     << SK_G32_SHIFT)
-                          | (blues      << SK_B32_SHIFT)
-                          | (Sk4i{0xFF} << SK_A32_SHIFT);
+                Sk4f (*linear_to_curve)(const Sk4f&) = (kSRGB_DstGamma == kDstGamma) ?
+                        sk_linear_to_srgb_needs_trunc : linear_to_2dot2;
+                Sk4i (*float_to_int)(const Sk4f&) = (kSRGB_DstGamma == kDstGamma) ?
+                        Sk4f_trunc : Sk4f_round;
+
+                dstReds   = linear_to_curve(dstReds);
+                dstGreens = linear_to_curve(dstGreens);
+                dstBlues  = linear_to_curve(dstBlues);
+
+                dstReds   = sk_clamp_0_255(dstReds);
+                dstGreens = sk_clamp_0_255(dstGreens);
+                dstBlues  = sk_clamp_0_255(dstBlues);
+
+                auto rgba = (float_to_int(dstReds)   << kRShift)
+                          | (float_to_int(dstGreens) << kGShift)
+                          | (float_to_int(dstBlues)  << kBShift)
+                          | (Sk4i{0xFF}              << kAShift);
                 rgba.store((uint32_t*) dst);
 
                 dst = SkTAddOffset<void>(dst, 4 * sizeof(uint32_t));
@@ -97,22 +119,22 @@ static void color_xform_RGB1(void* dst, const uint32_t* src, int len,
                 Sk4i indicesBlues  = Sk4f_round(scaledBlues);
 
                 uint32_t* dst32 = (uint32_t*) dst;
-                dst32[0] = dstTables[0][indicesReds  [0]] << SK_R32_SHIFT
-                         | dstTables[1][indicesGreens[0]] << SK_G32_SHIFT
-                         | dstTables[2][indicesBlues [0]] << SK_B32_SHIFT
-                         | 0xFF                           << SK_A32_SHIFT;
-                dst32[1] = dstTables[0][indicesReds  [1]] << SK_R32_SHIFT
-                         | dstTables[1][indicesGreens[1]] << SK_G32_SHIFT
-                         | dstTables[2][indicesBlues [1]] << SK_B32_SHIFT
-                         | 0xFF                           << SK_A32_SHIFT;
-                dst32[2] = dstTables[0][indicesReds  [2]] << SK_R32_SHIFT
-                         | dstTables[1][indicesGreens[2]] << SK_G32_SHIFT
-                         | dstTables[2][indicesBlues [2]] << SK_B32_SHIFT
-                         | 0xFF                           << SK_A32_SHIFT;
-                dst32[3] = dstTables[0][indicesReds  [3]] << SK_R32_SHIFT
-                         | dstTables[1][indicesGreens[3]] << SK_G32_SHIFT
-                         | dstTables[2][indicesBlues [3]] << SK_B32_SHIFT
-                         | 0xFF                           << SK_A32_SHIFT;
+                dst32[0] = dstTables[0][indicesReds  [0]] << kRShift
+                         | dstTables[1][indicesGreens[0]] << kGShift
+                         | dstTables[2][indicesBlues [0]] << kBShift
+                         | 0xFF                           << kAShift;
+                dst32[1] = dstTables[0][indicesReds  [1]] << kRShift
+                         | dstTables[1][indicesGreens[1]] << kGShift
+                         | dstTables[2][indicesBlues [1]] << kBShift
+                         | 0xFF                           << kAShift;
+                dst32[2] = dstTables[0][indicesReds  [2]] << kRShift
+                         | dstTables[1][indicesGreens[2]] << kGShift
+                         | dstTables[2][indicesBlues [2]] << kBShift
+                         | 0xFF                           << kAShift;
+                dst32[3] = dstTables[0][indicesReds  [3]] << kRShift
+                         | dstTables[1][indicesGreens[3]] << kGShift
+                         | dstTables[2][indicesBlues [3]] << kBShift
+                         | 0xFF                           << kAShift;
 
                 dst = SkTAddOffset<void>(dst, 4 * sizeof(uint32_t));
             } else {
@@ -142,28 +164,34 @@ static void color_xform_RGB1(void* dst, const uint32_t* src, int len,
              g = Sk4f{srcTables[1][(*src >>  8) & 0xFF]},
              b = Sk4f{srcTables[2][(*src >> 16) & 0xFF]};
 
-        auto dstPixel = rXgXbX*r + rYgYbY*g + rZgZbZ*b;
+        auto dstPixel = rXgXbX*r + rYgYbY*g + rZgZbZ*b + rTgTbT;
 
         if (kSRGB_DstGamma == kDstGamma || k2Dot2_DstGamma == kDstGamma) {
-            Sk4i (*linear_to_curve)(const Sk4f&) =
-                    (kSRGB_DstGamma == kDstGamma) ? sk_linear_to_srgb : linear_to_2dot2;
+            Sk4f (*linear_to_curve)(const Sk4f&) = (kSRGB_DstGamma == kDstGamma) ?
+                    sk_linear_to_srgb_needs_trunc : linear_to_2dot2;
+            Sk4i (*float_to_int)(const Sk4f&) = (kSRGB_DstGamma == kDstGamma) ?
+                    Sk4f_trunc : Sk4f_round;
 
-            auto pixel = linear_to_curve(dstPixel);
+            dstPixel = sk_clamp_0_255(linear_to_curve(dstPixel));
 
             uint32_t rgba;
-            SkNx_cast<uint8_t>(pixel).store(&rgba);
+            SkNx_cast<uint8_t>(float_to_int(dstPixel)).store(&rgba);
             rgba |= 0xFF000000;
-            *((uint32_t*) dst) = SkSwizzle_RGBA_to_PMColor(rgba);
+            if (kSwapRB) {
+                *((uint32_t*) dst) = SkSwizzle_RB(rgba);
+            } else {
+                *((uint32_t*) dst) = rgba;
+            }
             dst = SkTAddOffset<void>(dst, sizeof(uint32_t));
         } else if (kTable_DstGamma == kDstGamma) {
             Sk4f scaledPixel = Sk4f::Min(Sk4f::Max(1023.0f * dstPixel, 0.0f), 1023.0f);
 
             Sk4i indices = Sk4f_round(scaledPixel);
 
-            *((uint32_t*) dst) = dstTables[0][indices[0]] << SK_R32_SHIFT
-                               | dstTables[1][indices[1]] << SK_G32_SHIFT
-                               | dstTables[2][indices[2]] << SK_B32_SHIFT
-                               | 0xFF                     << SK_A32_SHIFT;
+            *((uint32_t*) dst) = dstTables[0][indices[0]] << kRShift
+                               | dstTables[1][indices[1]] << kGShift
+                               | dstTables[2][indices[2]] << kBShift
+                               | 0xFF                     << kAShift;
 
             dst = SkTAddOffset<void>(dst, sizeof(uint32_t));
         } else {
@@ -181,23 +209,42 @@ static void color_xform_RGB1(void* dst, const uint32_t* src, int len,
 
 static void color_xform_RGB1_to_2dot2(uint32_t* dst, const uint32_t* src, int len,
                                       const float* const srcTables[3], const float matrix[16]) {
-    color_xform_RGB1<k2Dot2_DstGamma>(dst, src, len, srcTables, matrix, nullptr);
+    color_xform_RGB1<k2Dot2_DstGamma, false>(dst, src, len, srcTables, matrix, nullptr);
 }
 
 static void color_xform_RGB1_to_srgb(uint32_t* dst, const uint32_t* src, int len,
                                      const float* const srcTables[3], const float matrix[16]) {
-    color_xform_RGB1<kSRGB_DstGamma>(dst, src, len, srcTables, matrix, nullptr);
+    color_xform_RGB1<kSRGB_DstGamma, false>(dst, src, len, srcTables, matrix, nullptr);
 }
 
 static void color_xform_RGB1_to_table(uint32_t* dst, const uint32_t* src, int len,
                                       const float* const srcTables[3], const float matrix[16],
                                       const uint8_t* const dstTables[3]) {
-    color_xform_RGB1<kTable_DstGamma>(dst, src, len, srcTables, matrix, dstTables);
+    color_xform_RGB1<kTable_DstGamma, false>(dst, src, len, srcTables, matrix, dstTables);
 }
 
 static void color_xform_RGB1_to_linear(uint64_t* dst, const uint32_t* src, int len,
                                        const float* const srcTables[3], const float matrix[16]) {
-    color_xform_RGB1<kLinear_DstGamma>(dst, src, len, srcTables, matrix, nullptr);
+    color_xform_RGB1<kLinear_DstGamma, false>(dst, src, len, srcTables, matrix, nullptr);
+}
+
+static void color_xform_RGB1_to_2dot2_swaprb(uint32_t* dst, const uint32_t* src, int len,
+                                             const float* const srcTables[3],
+                                             const float matrix[16]) {
+    color_xform_RGB1<k2Dot2_DstGamma, true>(dst, src, len, srcTables, matrix, nullptr);
+}
+
+static void color_xform_RGB1_to_srgb_swaprb(uint32_t* dst, const uint32_t* src, int len,
+                                            const float* const srcTables[3],
+                                            const float matrix[16]) {
+    color_xform_RGB1<kSRGB_DstGamma, true>(dst, src, len, srcTables, matrix, nullptr);
+}
+
+static void color_xform_RGB1_to_table_swaprb(uint32_t* dst, const uint32_t* src, int len,
+                                             const float* const srcTables[3],
+                                             const float matrix[16],
+                                             const uint8_t* const dstTables[3]) {
+    color_xform_RGB1<kTable_DstGamma, true>(dst, src, len, srcTables, matrix, dstTables);
 }
 
 }  // namespace SK_OPTS_NS