image/color: have Palette.Index honor alpha for closest match, not just

red, green and blue.

Fixes #9902

Change-Id: Ibffd0aa2f98996170e39a919296f69e9d5c71545
Reviewed-on: https://go-review.googlesource.com/8907Reviewed-by: Rob Pike <r@golang.org>

src/image/color/color.go

@@ -271,32 +271,39 @@ func (p Palette) Convert(c Color) Color {
 }
 
 // Index returns the index of the palette color closest to c in Euclidean
-// R,G,B space.
+// R,G,B,A space.
 func (p Palette) Index(c Color) int {
 	// A batch version of this computation is in image/draw/draw.go.
 
-	cr, cg, cb, _ := c.RGBA()
-	ret, bestSSD := 0, uint32(1<<32-1)
+	cr, cg, cb, ca := c.RGBA()
+	ret, bestSum := 0, uint32(1<<32-1)
 	for i, v := range p {
-		vr, vg, vb, _ := v.RGBA()
-		// We shift by 1 bit to avoid potential uint32 overflow in
-		// sum-squared-difference.
-		delta := (int32(cr) - int32(vr)) >> 1
-		ssd := uint32(delta * delta)
-		delta = (int32(cg) - int32(vg)) >> 1
-		ssd += uint32(delta * delta)
-		delta = (int32(cb) - int32(vb)) >> 1
-		ssd += uint32(delta * delta)
-		if ssd < bestSSD {
-			if ssd == 0 {
+		vr, vg, vb, va := v.RGBA()
+		sum := sqDiff(cr, vr) + sqDiff(cg, vg) + sqDiff(cb, vb) + sqDiff(ca, va)
+		if sum < bestSum {
+			if sum == 0 {
 				return i
 			}
-			ret, bestSSD = i, ssd
+			ret, bestSum = i, sum
 		}
 	}
 	return ret
 }
 
+// sqDiff returns the squared-difference of x and y, shifted by 2 so that
+// adding four of those won't overflow a uint32.
+//
+// x and y are both assumed to be in the range [0, 0xffff].
+func sqDiff(x, y uint32) uint32 {
+	var d uint32
+	if x > y {
+		d = x - y
+	} else {
+		d = y - x
+	}
+	return (d * d) >> 2
+}
+
 // Standard colors.
 var (
 	Black       = Gray16{0}

src/image/color/ycbcr_test.go

@@ -90,3 +90,29 @@ func TestCMYKToRGBConsistency(t *testing.T) {
 		}
 	}
 }
+
+func TestPalette(t *testing.T) {
+	p := Palette{
+		RGBA{0xff, 0xff, 0xff, 0xff},
+		RGBA{0x80, 0x00, 0x00, 0xff},
+		RGBA{0x7f, 0x00, 0x00, 0x7f},
+		RGBA{0x00, 0x00, 0x00, 0x7f},
+		RGBA{0x00, 0x00, 0x00, 0x00},
+		RGBA{0x40, 0x40, 0x40, 0x40},
+	}
+	// Check that, for a Palette with no repeated colors, the closest color to
+	// each element is itself.
+	for i, c := range p {
+		j := p.Index(c)
+		if i != j {
+			t.Errorf("Index(%v): got %d (color = %v), want %d", c, j, p[j], i)
+		}
+	}
+	// Check that finding the closest color considers alpha, not just red,
+	// green and blue.
+	got := p.Convert(RGBA{0x80, 0x00, 0x00, 0x80})
+	want := RGBA{0x7f, 0x00, 0x00, 0x7f}
+	if got != want {
+		t.Errorf("got %v, want %v", got, want)
+	}
+}

src/image/draw/draw.go

@@ -554,6 +554,20 @@ func clamp(i int32) int32 {
 	return i
 }
 
+// sqDiff returns the squared-difference of x and y, shifted by 2 so that
+// adding four of those won't overflow a uint32.
+//
+// x and y are both assumed to be in the range [0, 0xffff].
+func sqDiff(x, y int32) uint32 {
+	var d uint32
+	if x > y {
+		d = uint32(x - y)
+	} else {
+		d = uint32(y - x)
+	}
+	return (d * d) >> 2
+}
+
 func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point, floydSteinberg bool) {
 	// TODO(nigeltao): handle the case where the dst and src overlap.
 	// Does it even make sense to try and do Floyd-Steinberg whilst
@@ -563,14 +577,15 @@ func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point,
 	// dst.At. The dst.Set equivalent is a batch version of the algorithm
 	// used by color.Palette's Index method in image/color/color.go, plus
 	// optional Floyd-Steinberg error diffusion.
-	palette, pix, stride := [][3]int32(nil), []byte(nil), 0
+	palette, pix, stride := [][4]int32(nil), []byte(nil), 0
 	if p, ok := dst.(*image.Paletted); ok {
-		palette = make([][3]int32, len(p.Palette))
+		palette = make([][4]int32, len(p.Palette))
 		for i, col := range p.Palette {
-			r, g, b, _ := col.RGBA()
+			r, g, b, a := col.RGBA()
 			palette[i][0] = int32(r)
 			palette[i][1] = int32(g)
 			palette[i][2] = int32(b)
+			palette[i][3] = int32(a)
 		}
 		pix, stride = p.Pix[p.PixOffset(r.Min.X, r.Min.Y):], p.Stride
 	}
@@ -578,10 +593,10 @@ func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point,
 	// quantErrorCurr and quantErrorNext are the Floyd-Steinberg quantization
 	// errors that have been propagated to the pixels in the current and next
 	// rows. The +2 simplifies calculation near the edges.
-	var quantErrorCurr, quantErrorNext [][3]int32
+	var quantErrorCurr, quantErrorNext [][4]int32
 	if floydSteinberg {
-		quantErrorCurr = make([][3]int32, r.Dx()+2)
-		quantErrorNext = make([][3]int32, r.Dx()+2)
+		quantErrorCurr = make([][4]int32, r.Dx()+2)
+		quantErrorNext = make([][4]int32, r.Dx()+2)
 	}
 
 	// Loop over each source pixel.
@@ -590,30 +605,25 @@ func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point,
 		for x := 0; x != r.Dx(); x++ {
 			// er, eg and eb are the pixel's R,G,B values plus the
 			// optional Floyd-Steinberg error.
-			sr, sg, sb, _ := src.At(sp.X+x, sp.Y+y).RGBA()
-			er, eg, eb := int32(sr), int32(sg), int32(sb)
+			sr, sg, sb, sa := src.At(sp.X+x, sp.Y+y).RGBA()
+			er, eg, eb, ea := int32(sr), int32(sg), int32(sb), int32(sa)
 			if floydSteinberg {
 				er = clamp(er + quantErrorCurr[x+1][0]/16)
 				eg = clamp(eg + quantErrorCurr[x+1][1]/16)
 				eb = clamp(eb + quantErrorCurr[x+1][2]/16)
+				ea = clamp(ea + quantErrorCurr[x+1][3]/16)
 			}
 
 			if palette != nil {
-				// Find the closest palette color in Euclidean R,G,B space: the
-				// one that minimizes sum-squared-difference. We shift by 1 bit
-				// to avoid potential uint32 overflow in sum-squared-difference.
+				// Find the closest palette color in Euclidean R,G,B,A space:
+				// the one that minimizes sum-squared-difference.
 				// TODO(nigeltao): consider smarter algorithms.
-				bestIndex, bestSSD := 0, uint32(1<<32-1)
+				bestIndex, bestSum := 0, uint32(1<<32-1)
 				for index, p := range palette {
-					delta := (er - p[0]) >> 1
-					ssd := uint32(delta * delta)
-					delta = (eg - p[1]) >> 1
-					ssd += uint32(delta * delta)
-					delta = (eb - p[2]) >> 1
-					ssd += uint32(delta * delta)
-					if ssd < bestSSD {
-						bestIndex, bestSSD = index, ssd
-						if ssd == 0 {
+					sum := sqDiff(er, p[0]) + sqDiff(eg, p[1]) + sqDiff(eb, p[2]) + sqDiff(ea, p[3])
+					if sum < bestSum {
+						bestIndex, bestSum = index, sum
+						if sum == 0 {
 							break
 						}
 					}
@@ -626,11 +636,13 @@ func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point,
 				er -= int32(palette[bestIndex][0])
 				eg -= int32(palette[bestIndex][1])
 				eb -= int32(palette[bestIndex][2])
+				ea -= int32(palette[bestIndex][3])
 
 			} else {
 				out.R = uint16(er)
 				out.G = uint16(eg)
 				out.B = uint16(eb)
+				out.A = uint16(ea)
 				// The third argument is &out instead of out (and out is
 				// declared outside of the inner loop) to avoid the implicit
 				// conversion to color.Color here allocating memory in the
@@ -640,32 +652,37 @@ func drawPaletted(dst Image, r image.Rectangle, src image.Image, sp image.Point,
 				if !floydSteinberg {
 					continue
 				}
-				sr, sg, sb, _ = dst.At(r.Min.X+x, r.Min.Y+y).RGBA()
+				sr, sg, sb, sa = dst.At(r.Min.X+x, r.Min.Y+y).RGBA()
 				er -= int32(sr)
 				eg -= int32(sg)
 				eb -= int32(sb)
+				ea -= int32(sa)
 			}
 
 			// Propagate the Floyd-Steinberg quantization error.
 			quantErrorNext[x+0][0] += er * 3
 			quantErrorNext[x+0][1] += eg * 3
 			quantErrorNext[x+0][2] += eb * 3
+			quantErrorNext[x+0][3] += ea * 3
 			quantErrorNext[x+1][0] += er * 5
 			quantErrorNext[x+1][1] += eg * 5
 			quantErrorNext[x+1][2] += eb * 5
+			quantErrorNext[x+1][3] += ea * 5
 			quantErrorNext[x+2][0] += er * 1
 			quantErrorNext[x+2][1] += eg * 1
 			quantErrorNext[x+2][2] += eb * 1
+			quantErrorNext[x+2][3] += ea * 1
 			quantErrorCurr[x+2][0] += er * 7
 			quantErrorCurr[x+2][1] += eg * 7
 			quantErrorCurr[x+2][2] += eb * 7
+			quantErrorCurr[x+2][3] += ea * 7
 		}
 
 		// Recycle the quantization error buffers.
 		if floydSteinberg {
 			quantErrorCurr, quantErrorNext = quantErrorNext, quantErrorCurr
 			for i := range quantErrorNext {
-				quantErrorNext[i] = [3]int32{}
+				quantErrorNext[i] = [4]int32{}
 			}
 		}
 	}