image/png: pack image data for small bitdepth paletted images

Bit packs image data when writing images with fewer than 16
colors in its palette. Reading of bit packed image data was
already implemented.

Fixes #19879

Change-Id: I0a06f9599a163931e20d3503fc3722e5101f0070
Reviewed-on: https://go-review.googlesource.com/134235
Reviewed-by: Nigel Tao <nigeltao@golang.org>

2 files changed, 170 insertions, 11 deletions

diff --git a/src/image/png/writer.go b/src/image/png/writer.go
index de8c28e919..c03335120e 100644
--- a/src/image/png/writer.go
+++ b/src/image/png/writer.go
@@ -137,6 +137,15 @@ func (e *encoder) writeIHDR() {
 	case cbP8:
 		e.tmp[8] = 8
 		e.tmp[9] = ctPaletted
+	case cbP4:
+		e.tmp[8] = 4
+		e.tmp[9] = ctPaletted
+	case cbP2:
+		e.tmp[8] = 2
+		e.tmp[9] = ctPaletted
+	case cbP1:
+		e.tmp[8] = 1
+		e.tmp[9] = ctPaletted
 	case cbTCA8:
 		e.tmp[8] = 8
 		e.tmp[9] = ctTrueColorAlpha
@@ -305,31 +314,38 @@ func (e *encoder) writeImage(w io.Writer, m image.Image, cb int, level int) erro
 	}
 	defer e.zw.Close()
 
-	bpp := 0 // Bytes per pixel.
+	bitsPerPixel := 0
 
 	switch cb {
 	case cbG8:
-		bpp = 1
+		bitsPerPixel = 8
 	case cbTC8:
-		bpp = 3
+		bitsPerPixel = 24
 	case cbP8:
-		bpp = 1
+		bitsPerPixel = 8
+	case cbP4:
+		bitsPerPixel = 4
+	case cbP2:
+		bitsPerPixel = 2
+	case cbP1:
+		bitsPerPixel = 1
 	case cbTCA8:
-		bpp = 4
+		bitsPerPixel = 32
 	case cbTC16:
-		bpp = 6
+		bitsPerPixel = 48
 	case cbTCA16:
-		bpp = 8
+		bitsPerPixel = 64
 	case cbG16:
-		bpp = 2
+		bitsPerPixel = 16
 	}
+
 	// cr[*] and pr are the bytes for the current and previous row.
 	// cr[0] is unfiltered (or equivalently, filtered with the ftNone filter).
 	// cr[ft], for non-zero filter types ft, are buffers for transforming cr[0] under the
 	// other PNG filter types. These buffers are allocated once and re-used for each row.
 	// The +1 is for the per-row filter type, which is at cr[*][0].
 	b := m.Bounds()
-	sz := 1 + bpp*b.Dx()
+	sz := 1 + (bitsPerPixel*b.Dx()+7)/8
 	for i := range e.cr {
 		if cap(e.cr[i]) < sz {
 			e.cr[i] = make([]uint8, sz)
@@ -405,6 +421,30 @@ func (e *encoder) writeImage(w io.Writer, m image.Image, cb int, level int) erro
 					i += 1
 				}
 			}
+
+		case cbP4, cbP2, cbP1:
+			pi := m.(image.PalettedImage)
+
+			var a uint8
+			var c int
+			for x := b.Min.X; x < b.Max.X; x++ {
+				a = a<<uint(bitsPerPixel) | pi.ColorIndexAt(x, y)
+				c++
+				if c == 8/bitsPerPixel {
+					cr[0][i] = a
+					i += 1
+					a = 0
+					c = 0
+				}
+			}
+			if c != 0 {
+				for c != 8/bitsPerPixel {
+					a = a << uint(bitsPerPixel)
+					c++
+				}
+				cr[0][i] = a
+			}
+
 		case cbTCA8:
 			if nrgba != nil {
 				offset := (y - b.Min.Y) * nrgba.Stride
@@ -460,7 +500,10 @@ func (e *encoder) writeImage(w io.Writer, m image.Image, cb int, level int) erro
 		// "filters are rarely useful on palette images" and will result
 		// in larger files (see http://www.libpng.org/pub/png/book/chapter09.html).
 		f := ftNone
-		if level != zlib.NoCompression && cb != cbP8 {
+		if level != zlib.NoCompression && cb != cbP8 && cb != cbP4 && cb != cbP2 && cb != cbP1 {
+			// Since we skip paletted images we don't have to worry about
+			// bitsPerPixel not being a multiple of 8
+			bpp := bitsPerPixel / 8
 			f = filter(&cr, pr, bpp)
 		}
 
@@ -551,7 +594,15 @@ func (enc *Encoder) Encode(w io.Writer, m image.Image) error {
 		pal, _ = m.ColorModel().(color.Palette)
 	}
 	if pal != nil {
-		e.cb = cbP8
+		if len(pal) <= 2 {
+			e.cb = cbP1
+		} else if len(pal) <= 4 {
+			e.cb = cbP2
+		} else if len(pal) <= 16 {
+			e.cb = cbP4
+		} else {
+			e.cb = cbP8
+		}
 	} else {
 		switch m.ColorModel() {
 		case color.GrayModel:
diff --git a/src/image/png/writer_test.go b/src/image/png/writer_test.go
index 6c5e942310..5d131ff823 100644
--- a/src/image/png/writer_test.go
+++ b/src/image/png/writer_test.go
@@ -6,9 +6,12 @@ package png
 
 import (
 	"bytes"
+	"compress/zlib"
+	"encoding/binary"
 	"fmt"
 	"image"
 	"image/color"
+	"io"
 	"io/ioutil"
 	"testing"
 )
@@ -77,6 +80,111 @@ func TestWriter(t *testing.T) {
 	}
 }
 
+func TestWriterPaletted(t *testing.T) {
+	const width, height = 32, 16
+
+	testCases := []struct {
+		plen     int
+		bitdepth uint8
+		datalen  int
+	}{
+
+		{
+			plen:     256,
+			bitdepth: 8,
+			datalen:  (1 + width) * height,
+		},
+
+		{
+			plen:     128,
+			bitdepth: 8,
+			datalen:  (1 + width) * height,
+		},
+
+		{
+			plen:     16,
+			bitdepth: 4,
+			datalen:  (1 + width/2) * height,
+		},
+
+		{
+			plen:     4,
+			bitdepth: 2,
+			datalen:  (1 + width/4) * height,
+		},
+
+		{
+			plen:     2,
+			bitdepth: 1,
+			datalen:  (1 + width/8) * height,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("plen-%d", tc.plen), func(t *testing.T) {
+			// Create a paletted image with the correct palette length
+			palette := make(color.Palette, tc.plen)
+			for i := range palette {
+				palette[i] = color.NRGBA{
+					R: uint8(i),
+					G: uint8(i),
+					B: uint8(i),
+					A: 255,
+				}
+			}
+			m0 := image.NewPaletted(image.Rect(0, 0, width, height), palette)
+
+			i := 0
+			for y := 0; y < height; y++ {
+				for x := 0; x < width; x++ {
+					m0.SetColorIndex(x, y, uint8(i%tc.plen))
+					i++
+				}
+			}
+
+			// Encode the image
+			var b bytes.Buffer
+			if err := Encode(&b, m0); err != nil {
+				t.Error(err)
+				return
+			}
+			const chunkFieldsLength = 12 // 4 bytes for length, name and crc
+			data := b.Bytes()
+			i = len(pngHeader)
+
+			for i < len(data)-chunkFieldsLength {
+				length := binary.BigEndian.Uint32(data[i : i+4])
+				name := string(data[i+4 : i+8])
+
+				switch name {
+				case "IHDR":
+					bitdepth := data[i+8+8]
+					if bitdepth != tc.bitdepth {
+						t.Errorf("got bitdepth %d, want %d", bitdepth, tc.bitdepth)
+					}
+				case "IDAT":
+					// Uncompress the image data
+					r, err := zlib.NewReader(bytes.NewReader(data[i+8 : i+8+int(length)]))
+					if err != nil {
+						t.Error(err)
+						return
+					}
+					n, err := io.Copy(ioutil.Discard, r)
+					if err != nil {
+						t.Errorf("got error while reading image data: %v", err)
+					}
+					if n != int64(tc.datalen) {
+						t.Errorf("got uncompressed data length %d, want %d", n, tc.datalen)
+					}
+				}
+
+				i += chunkFieldsLength + int(length)
+			}
+		})
+
+	}
+}
+
 func TestWriterLevels(t *testing.T) {
 	m := image.NewNRGBA(image.Rect(0, 0, 100, 100))