Encode kIndex8 to PNG more efficiently

src/images/SkPNGImageEncoder.cpp

Index: src/images/SkPNGImageEncoder.cpp
diff --git a/src/images/SkPNGImageEncoder.cpp b/src/images/SkPNGImageEncoder.cpp
index 9b8808a2e0e765fbccfaabd40f1c3a9a12673718..765df3f4659a77024322f7c359504d6af786d7d3 100644
--- a/src/images/SkPNGImageEncoder.cpp
+++ b/src/images/SkPNGImageEncoder.cpp
@@ -103,70 +103,65 @@ static int computeBitDepth(int colorCount) {
 #endif
 }
 
-/*  Pack palette[] with the corresponding colors, and if hasAlpha is true, also
-    pack trans[] and return the number of trans[] entries written. If hasAlpha
-    is false, the return value will always be 0.
-
-    Note: this routine takes care of unpremultiplying the RGB values when we
-    have alpha in the colortable, since png doesn't support premul colors
+/*  Pack palette[] with the corresponding colors, and if the image has alpha, also
+    pack trans[] and return the number of alphas[] entries written. If the image is
+    opaque, the return value will always be 0.
 */
-static inline int pack_palette(SkColorTable* ctable,
-                               png_color* SK_RESTRICT palette,
-                               png_byte* SK_RESTRICT trans, SkAlphaType alphaType) {
-    const SkPMColor* SK_RESTRICT colors = ctable ? ctable->readColors() : nullptr;
-    const int ctCount = ctable->count();
-    int i, num_trans = 0;
-
+static inline int pack_palette(SkColorTable* ctable, png_color* SK_RESTRICT palette,
+                               png_byte* SK_RESTRICT alphas, SkAlphaType alphaType) {
+    const SkPMColor* SK_RESTRICT colors = ctable->readColors();
+    const int count = ctable->count();
+    int numWithAlpha = 0;
     if (kOpaque_SkAlphaType != alphaType) {
-        /*  first see if we have some number of fully opaque at the end of the
-            ctable. PNG allows num_trans < num_palette, but all of the trans
-            entries must come first in the palette. If I was smarter, I'd
-            reorder the indices and ctable so that all non-opaque colors came
-            first in the palette. But, since that would slow down the encode,
-            I'm leaving the indices and ctable order as is, and just looking
-            at the tail of the ctable for opaqueness.
-        */
-        num_trans = ctCount;
-        for (i = ctCount - 1; i >= 0; --i) {
-            if (SkGetPackedA32(colors[i]) != 0xFF) {
-                break;
-            }
-            num_trans -= 1;
-        }
-
-        if (kPremul_SkAlphaType == alphaType) {
-            const SkUnPreMultiply::Scale* SK_RESTRICT table = SkUnPreMultiply::GetScaleTable();
-            for (i = 0; i < num_trans; i++) {
-                const SkPMColor c = *colors++;
-                const unsigned a = SkGetPackedA32(c);
-                const SkUnPreMultiply::Scale s = table[a];
-                trans[i] = a;
-                palette[i].red = SkUnPreMultiply::ApplyScale(s, SkGetPackedR32(c));
-                palette[i].green = SkUnPreMultiply::ApplyScale(s,SkGetPackedG32(c));
-                palette[i].blue = SkUnPreMultiply::ApplyScale(s, SkGetPackedB32(c));
+        auto getUnpremulColor = [alphaType](uint8_t color, uint8_t alpha) {
+            if (kPremul_SkAlphaType == alphaType) {
+                const SkUnPreMultiply::Scale* table = SkUnPreMultiply::GetScaleTable();
+                const SkUnPreMultiply::Scale scale = table[alpha];
+                return (uint8_t) SkUnPreMultiply::ApplyScale(scale, color);
+            } else {
+                return color;
             }
-        } else {
-            for (i = 0; i < num_trans; i++) {
-                const SkPMColor c = *colors++;
-                trans[i] = SkGetPackedA32(c);
-                palette[i].red = SkGetPackedR32(c);
-                palette[i].green = SkGetPackedG32(c);
-                palette[i].blue = SkGetPackedB32(c);
+        };
+
+        // PNG requires that all non-opaque colors come first in the palette.  Write these first.
+        for (int i = 0; i < count; i++) {
+            uint8_t alpha = SkGetPackedA32(colors[i]);
+            if (0xFF != alpha) {
+                alphas[numWithAlpha] = alpha;
+                palette[numWithAlpha].red   = getUnpremulColor(SkGetPackedR32(colors[i]), alpha);
+                palette[numWithAlpha].green = getUnpremulColor(SkGetPackedG32(colors[i]), alpha);
+                palette[numWithAlpha].blue  = getUnpremulColor(SkGetPackedB32(colors[i]), alpha);
+                numWithAlpha++;
             }
         }
 
-        // now fall out of this if-block to use common code for the trailing
-        // opaque entries
     }
 
-    // these (remaining) entries are opaque
-    for (i = num_trans; i < ctCount; i++) {
-        SkPMColor c = *colors++;
-        palette[i].red = SkGetPackedR32(c);
-        palette[i].green = SkGetPackedG32(c);
-        palette[i].blue = SkGetPackedB32(c);
+    if (0 == numWithAlpha) {
+        // All of the entries are opaque.
+        for (int i = 0; i < count; i++) {
+            SkPMColor c = *colors++;
+            palette[i].red   = SkGetPackedR32(c);
+            palette[i].green = SkGetPackedG32(c);
+            palette[i].blue  = SkGetPackedB32(c);
+        }
+    } else {
+        // We have already written the non-opaque colors.  Now just write the opaque colors.
+        int currIndex = numWithAlpha;
+        int i = 0;
+        while (i < count && currIndex != count) {

[scroggo, 2016/09/12 17:12:48]

I don't think you need to check currIndex on each iteration through the loop. If
you do, I think you've made a mistake somewhere above here. I was thinking
something like

else if (numWithAlpha < count) {

Maybe you want to put some asserts in here?

[msarett, 2016/09/12 17:17:13]

Yeah I wrote that at first.  I think the advantage of this approach is that we
will exit the loop early once we hit the last opaque alpha.

I'd guess which is "better" depends on the input.  I don't feel strongly.

[scroggo, 2016/09/12 17:19:44]

Ah, so this would let us exit early if for example the input had all the
non-opaque colors at the end.

Either way is fine by me.

+            uint8_t alpha = SkGetPackedA32(colors[i]);
+            if (0xFF == alpha) {
+                palette[currIndex].red   = SkGetPackedR32(colors[i]);
+                palette[currIndex].green = SkGetPackedG32(colors[i]);
+                palette[currIndex].blue  = SkGetPackedB32(colors[i]);
+                currIndex++;
+                i++;
+            }
+        }
     }
-    return num_trans;
+
+    return numWithAlpha;
 }
 
 class SkPNGImageEncoder : public SkImageEncoder {
@@ -327,6 +322,7 @@ bool SkPNGImageEncoder::doEncode(SkWStream* stream, const SkBitmap& bitmap,
     png_byte trans[256];
     if (kIndex_8_SkColorType == ct) {
         SkColorTable* ct = bitmap.getColorTable();
+        SkASSERT(ct);
         int numTrans = pack_palette(ct, paletteColors, trans, alphaType);
         png_set_PLTE(png_ptr, info_ptr, paletteColors, ct->count());
         if (numTrans > 0) {