perf(screenshot): encode PNG straight from the capture buffer

related #3056

core/cmd/dms/commands_screenshot.go

@@ -381,7 +381,7 @@ func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, q
 		}
 	default:
 		mimeType = "image/png"
-		if err := screenshot.EncodePNGTagged(&data, screenshot.BufferToImageDeep(buf, pixelFormat), cicp); err != nil {
+		if err := screenshot.EncodeBufferPNG(&data, buf, pixelFormat, cicp); err != nil {
 			return err
 		}
 	}
@@ -394,7 +394,7 @@ func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, qua
 	case screenshot.FormatJPEG:
 		return screenshot.EncodeJPEG(os.Stdout, screenshot.BufferToImageWithFormat(buf, pixelFormat), quality)
 	default:
-		return screenshot.EncodePNGTagged(os.Stdout, screenshot.BufferToImageDeep(buf, pixelFormat), cicp)
+		return screenshot.EncodeBufferPNG(os.Stdout, buf, pixelFormat, cicp)
 	}
 }
 

core/internal/screenshot/cicp.go

@@ -3,8 +3,6 @@ package screenshot
 import (
 	"encoding/binary"
 	"hash/crc32"
-	"image"
-	"io"
 
 	"github.com/AvengeMedia/DankMaterialShell/core/internal/proto/wp_color_management"
 )
@@ -77,6 +75,13 @@ func cicpFromNamed(primaries, tf uint32) *CICP {
 	return &CICP{Primaries: p, TF: t}
 }
 
+func (c *CICP) chunks() [][]byte {
+	if c == nil {
+		return nil
+	}
+	return [][]byte{c.chunk()}
+}
+
 func (c *CICP) chunk() []byte {
 	payload := []byte{'c', 'I', 'C', 'P', c.Primaries, c.TF, 0, 1} // matrix=RGB, full range
 	chunk := binary.BigEndian.AppendUint32(make([]byte, 0, 16), 4)
@@ -84,13 +89,6 @@ func (c *CICP) chunk() []byte {
 	return binary.BigEndian.AppendUint32(chunk, crc32.ChecksumIEEE(payload))
 }
 
-func EncodePNGTagged(w io.Writer, img image.Image, cicp *CICP) error {
-	if cicp == nil {
-		return encodePNG(w, img)
-	}
-	return encodePNG(w, img, cicp.chunk())
-}
-
 // outputCICP queries the output's image description via wp_color_management_v1.
 // Returns nil unless the output uses a wide-gamut or HDR encoding.
 func (s *Screenshoter) outputCICP(output *WaylandOutput) *CICP {

core/internal/screenshot/encode.go

@@ -47,39 +47,6 @@ func tenBitToRGBA(buf *ShmBuffer, format uint32) *image.RGBA {
 	return img
 }
 
-// BufferToImageDeep preserves 10-bit sources as 16-bit-per-channel; others decode to 8-bit RGBA.
-func BufferToImageDeep(buf *ShmBuffer, format uint32) image.Image {
-	if !PixelFormat(format).Is10Bit() {
-		return BufferToImageWithFormat(buf, format)
-	}
-
-	img := image.NewNRGBA64(image.Rect(0, 0, buf.Width, buf.Height))
-	data := buf.Data()
-	swapRB := tenBitSwapRB(format)
-
-	for y := 0; y < buf.Height; y++ {
-		srcOff := y * buf.Stride
-		dstOff := y * img.Stride
-		for x := 0; x < buf.Width; x++ {
-			si := srcOff + x*4
-			di := dstOff + x*8
-			if si+4 > len(data) || di+8 > len(img.Pix) {
-				continue
-			}
-			v := binary.LittleEndian.Uint32(data[si:])
-			c0, c1, c2 := v&0x3FF, (v>>10)&0x3FF, (v>>20)&0x3FF
-			if swapRB {
-				c0, c2 = c2, c0
-			}
-			binary.BigEndian.PutUint16(img.Pix[di+0:], uint16(c0<<6|c0>>4))
-			binary.BigEndian.PutUint16(img.Pix[di+2:], uint16(c1<<6|c1>>4))
-			binary.BigEndian.PutUint16(img.Pix[di+4:], uint16(c2<<6|c2>>4))
-			binary.BigEndian.PutUint16(img.Pix[di+6:], 0xFFFF)
-		}
-	}
-	return img
-}
-
 func BufferToImageWithFormat(buf *ShmBuffer, format uint32) *image.RGBA {
 	if PixelFormat(format).Is10Bit() {
 		return tenBitToRGBA(buf, format)
@@ -172,10 +139,6 @@ func ImageToBuffer(img image.Image) (*ShmBuffer, error) {
 	return buf, nil
 }
 
-func EncodePNG(w io.Writer, img image.Image) error {
-	return encodePNG(w, img)
-}
-
 func EncodeJPEG(w io.Writer, img image.Image, quality int) error {
 	return jpeg.Encode(w, img, &jpeg.Options{Quality: quality})
 }
@@ -251,6 +214,6 @@ func WriteToFileWithFormat(buf *ShmBuffer, path string, format Format, quality i
 	case FormatPPM:
 		return EncodePPM(f, BufferToImageWithFormat(buf, pixelFormat))
 	default:
-		return EncodePNGTagged(f, BufferToImageDeep(buf, pixelFormat), cicp)
+		return EncodeBufferPNG(f, buf, pixelFormat, cicp)
 	}
 }

core/internal/screenshot/encode_test.go

@@ -2,7 +2,6 @@ package screenshot
 
 import (
 	"encoding/binary"
-	"image"
 	"testing"
 )
 
@@ -44,27 +43,3 @@ func TestBufferToImageWithFormat10Bit(t *testing.T) {
 		})
 	}
 }
-
-func TestBufferToImageDeep10Bit(t *testing.T) {
-	const r10, g10, b10 = uint32(1023), uint32(512), uint32(0)
-	buf := newTenBitBuffer(t, FormatXRGB2101010, r10<<20|g10<<10|b10)
-
-	img, ok := BufferToImageDeep(buf, uint32(FormatXRGB2101010)).(*image.NRGBA64)
-	if !ok {
-		t.Fatal("expected *image.NRGBA64 for 10-bit source")
-	}
-
-	want := [4]uint16{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF}
-	for i, w := range want {
-		if got := binary.BigEndian.Uint16(img.Pix[i*2:]); got != w {
-			t.Errorf("channel %d = %d, want %d", i, got, w)
-		}
-	}
-}
-
-func TestBufferToImageDeep8BitFallback(t *testing.T) {
-	buf := newTenBitBuffer(t, FormatXRGB8888, 0)
-	if _, ok := BufferToImageDeep(buf, uint32(FormatXRGB8888)).(*image.RGBA); !ok {
-		t.Error("expected *image.RGBA for 8-bit source")
-	}
-}

core/internal/screenshot/png.go

@@ -25,6 +25,7 @@ const (
 	adlerBase      = 65521
 )
 
+// pngSource reads rows in place from pix; pack converts a row to PNG channel order, nil when it already is.
 type pngSource struct {
 	pix           []byte
 	stride        int
@@ -32,71 +33,107 @@ type pngSource struct {
 	depth         byte
 	colorType     byte
 	srcBpp        int
-	dstBpp        int
+	pack          func(dst, src []byte)
 }
 
-func pngSourceOf(pix []byte, stride, width, height int, depth byte, opaque bool) pngSource {
-	s := pngSource{pix: pix, stride: stride, width: width, height: height, depth: depth}
-	s.srcBpp = 4 * int(depth) / 8
-	s.dstBpp, s.colorType = s.srcBpp, pngColorRGBA
-	if opaque {
-		s.dstBpp, s.colorType = 3*int(depth)/8, pngColorRGB
+func (s pngSource) dstBpp() int {
+	channels := 4
+	if s.colorType == pngColorRGB {
+		channels = 3
+	}
+	return channels * int(s.depth) / 8
+}
+
+func packRGBX(dst, src []byte) {
+	for x := range len(dst) / 3 {
+		d, p := dst[x*3:x*3+3], src[x*4:x*4+4]
+		d[0], d[1], d[2] = p[0], p[1], p[2]
+	}
+}
+
+func packBGRX(dst, src []byte) {
+	for x := range len(dst) / 3 {
+		d, p := dst[x*3:x*3+3], src[x*4:x*4+4]
+		d[0], d[1], d[2] = p[2], p[1], p[0]
+	}
+}
+
+func packRGBX16(dst, src []byte) {
+	for x := range len(dst) / 6 {
+		copy(dst[x*6:x*6+6], src[x*8:x*8+6])
+	}
+}
+
+func pack2101010(swapRB bool) func(dst, src []byte) {
+	return func(dst, src []byte) {
+		for x := range len(dst) / 6 {
+			v := binary.LittleEndian.Uint32(src[x*4:])
+			c0, c1, c2 := v&0x3FF, v>>10&0x3FF, v>>20&0x3FF
+			if swapRB {
+				c0, c2 = c2, c0
+			}
+			d := dst[x*6 : x*6+6]
+			binary.BigEndian.PutUint16(d[0:], uint16(c0<<6|c0>>4))
+			binary.BigEndian.PutUint16(d[2:], uint16(c1<<6|c1>>4))
+			binary.BigEndian.PutUint16(d[4:], uint16(c2<<6|c2>>4))
+		}
 	}
-	return s
 }
 
 func newPNGSource(img image.Image) pngSource {
 	b := img.Bounds()
+	s := pngSource{width: b.Dx(), height: b.Dy(), colorType: pngColorRGBA}
 	switch src := img.(type) {
 	case *image.RGBA:
-		if src.Opaque() {
-			return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 8, true)
+		if !src.Opaque() {
+			break
 		}
+		s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 8, 4
+		s.colorType, s.pack = pngColorRGB, packRGBX
+		return s
 	case *image.NRGBA:
-		return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 8, src.Opaque())
+		s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 8, 4
+		if src.Opaque() {
+			s.colorType, s.pack = pngColorRGB, packRGBX
+		}
+		return s
 	case *image.NRGBA64:
-		return pngSourceOf(src.Pix, src.Stride, b.Dx(), b.Dy(), 16, src.Opaque())
+		s.pix, s.stride, s.depth, s.srcBpp = src.Pix, src.Stride, 16, 8
+		if src.Opaque() {
+			s.colorType, s.pack = pngColorRGB, packRGBX16
+		}
+		return s
 	}
-	nrgba := image.NewNRGBA(image.Rect(0, 0, b.Dx(), b.Dy()))
+	nrgba := image.NewNRGBA(image.Rect(0, 0, s.width, s.height))
 	draw.Draw(nrgba, nrgba.Rect, img, b.Min, draw.Src)
 	return newPNGSource(nrgba)
 }
 
-func (s pngSource) row(y int, scratch []byte) []byte {
-	src := s.pix[y*s.stride : y*s.stride+s.width*s.srcBpp]
-	switch s.dstBpp {
-	case s.srcBpp:
-		return src
-	case 3:
-		for x := 0; x < s.width; x++ {
-			d, p := scratch[x*3:x*3+3], src[x*4:x*4+4]
-			d[0], d[1], d[2] = p[0], p[1], p[2]
-		}
+func newShmPNGSource(buf *ShmBuffer, format uint32) pngSource {
+	s := pngSource{pix: buf.Data(), stride: buf.Stride, width: buf.Width, height: buf.Height, depth: 8, colorType: pngColorRGB, srcBpp: 4}
+	switch f := PixelFormat(format); {
+	case f.Is10Bit():
+		s.depth, s.pack = 16, pack2101010(tenBitSwapRB(format))
+	case f == FormatABGR8888, f == FormatXBGR8888:
+		s.pack = packRGBX
 	default:
-		for x := 0; x < s.width; x++ {
-			copy(scratch[x*s.dstBpp:(x+1)*s.dstBpp], src[x*s.srcBpp:])
-		}
+		s.pack = packBGRX
 	}
-	return scratch
+	return s
 }
 
-func (s pngSource) filterBand(y0, y1 int) []byte {
-	rowLen := s.width * s.dstBpp
-	out := make([]byte, (rowLen+1)*(y1-y0))
-	scratch := make([]byte, rowLen)
-	bpp := s.dstBpp
-	for y := y0; y < y1; y++ {
-		line := out[(y-y0)*(rowLen+1):][:rowLen+1]
-		line[0] = pngFilterSub
-		row := s.row(y, scratch)
-		copy(line[1:1+bpp], row[:bpp])
-		cur, prev, dst := row[bpp:], row[:rowLen-bpp], line[1+bpp:]
-		prev, dst = prev[:len(cur)], dst[:len(cur)]
-		for i := range cur {
-			dst[i] = cur[i] - prev[i]
-		}
+func (s pngSource) filterRow(y, bpp int, packed, dst []byte) {
+	row := s.pix[y*s.stride:][:s.width*s.srcBpp]
+	if s.pack != nil {
+		s.pack(packed, row)
+		row = packed
+	}
+	copy(dst[:bpp], row[:bpp])
+	cur, prev, out := row[bpp:], row[:len(row)-bpp], dst[bpp:]
+	prev, out = prev[:len(cur)], out[:len(cur)]
+	for i := range cur {
+		out[i] = cur[i] - prev[i]
 	}
-	return out
 }
 
 type pngBand struct {
@@ -106,21 +143,30 @@ type pngBand struct {
 }
 
 func (s pngSource) compressBand(y0, y1 int, final bool) (pngBand, error) {
-	raw := s.filterBand(y0, y1)
+	bpp := s.dstBpp()
+	rowLen := s.width * bpp
+	line := make([]byte, rowLen+1)
+	line[0] = pngFilterSub
+	packed := make([]byte, rowLen)
 	var buf bytes.Buffer
 	fw, err := flate.NewWriter(&buf, flate.BestSpeed)
 	if err != nil {
 		return pngBand{}, err
 	}
-	if _, err := fw.Write(raw); err != nil {
-		return pngBand{}, err
+	sum := adler32.New()
+	for y := y0; y < y1; y++ {
+		s.filterRow(y, bpp, packed, line[1:])
+		sum.Write(line)
+		if _, err := fw.Write(line); err != nil {
+			return pngBand{}, err
+		}
 	}
 	if final {
 		err = fw.Close()
 	} else {
 		err = fw.Flush()
 	}
-	return pngBand{data: buf.Bytes(), adler: adler32.Checksum(raw), size: len(raw)}, err
+	return pngBand{data: buf.Bytes(), adler: sum.Sum32(), size: len(line) * (y1 - y0)}, err
 }
 
 func adlerCombine(a, b uint32, lenB int) uint32 {
@@ -175,21 +221,20 @@ func (s pngSource) ihdr() []byte {
 	return b
 }
 
-// encodePNG deflates row bands in parallel into a single zlib stream; extraChunks go right after IHDR.
-func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error {
-	src := newPNGSource(img)
-	if src.width <= 0 || src.height <= 0 {
+// encode deflates row bands in parallel into a single zlib stream; extraChunks go right after IHDR.
+func (s pngSource) encode(w io.Writer, extraChunks ...[]byte) error {
+	if s.width <= 0 || s.height <= 0 {
 		return errors.New("png: empty image")
 	}
 
-	bands := min(runtime.GOMAXPROCS(0), max(1, src.height/pngMinBandRows))
+	bands := min(runtime.GOMAXPROCS(0), max(1, s.height/pngMinBandRows))
 	results := make([]pngBand, bands)
 	errs := make([]error, bands)
 	var wg sync.WaitGroup
 	for b := range bands {
 		wg.Go(func() {
-			y0, y1 := src.height*b/bands, src.height*(b+1)/bands
-			results[b], errs[b] = src.compressBand(y0, y1, b == bands-1)
+			y0, y1 := s.height*b/bands, s.height*(b+1)/bands
+			results[b], errs[b] = s.compressBand(y0, y1, b == bands-1)
 		})
 	}
 	wg.Wait()
@@ -209,7 +254,7 @@ func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error {
 	if _, err := w.Write(pngSignature); err != nil {
 		return err
 	}
-	if err := writeChunk(w, "IHDR", src.ihdr()); err != nil {
+	if err := writeChunk(w, "IHDR", s.ihdr()); err != nil {
 		return err
 	}
 	for _, c := range extraChunks {
@@ -222,3 +267,16 @@ func encodePNG(w io.Writer, img image.Image, extraChunks ...[]byte) error {
 	}
 	return writeChunk(w, "IEND")
 }
+
+func EncodePNG(w io.Writer, img image.Image) error {
+	return newPNGSource(img).encode(w)
+}
+
+func EncodePNGTagged(w io.Writer, img image.Image, cicp *CICP) error {
+	return newPNGSource(img).encode(w, cicp.chunks()...)
+}
+
+// EncodeBufferPNG reads the capture mapping directly; 10-bit formats are kept as 16-bit channels.
+func EncodeBufferPNG(w io.Writer, buf *ShmBuffer, format uint32, cicp *CICP) error {
+	return newShmPNGSource(buf, format).encode(w, cicp.chunks()...)
+}

core/internal/screenshot/png_test.go

@@ -2,6 +2,7 @@ package screenshot
 
 import (
 	"bytes"
+	"encoding/binary"
 	"hash/adler32"
 	"image"
 	"image/color"
@@ -136,3 +137,58 @@ func TestAdlerCombine(t *testing.T) {
 		t.Error("combine with empty prefix must be identity")
 	}
 }
+
+func TestEncodeBufferPNG(t *testing.T) {
+	const w, h = 5, 3
+	const r10, g10, b10 = uint32(1023), uint32(512), uint32(0)
+	cases := []struct {
+		name   string
+		format PixelFormat
+		pixel  uint32
+		want   color.RGBA64
+	}{
+		{"xrgb8888", FormatXRGB8888, 0x00FF8000, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}},
+		{"argb8888-alpha-ignored", FormatARGB8888, 0x10FF8000, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}},
+		{"xbgr8888", FormatXBGR8888, 0x000080FF, color.RGBA64{0xFFFF, 0x8080, 0, 0xFFFF}},
+		{"xrgb2101010", FormatXRGB2101010, r10<<20 | g10<<10 | b10, color.RGBA64{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF}},
+		{"xbgr2101010", FormatXBGR2101010, b10<<20 | g10<<10 | r10, color.RGBA64{0xFFFF, 512<<6 | 512>>4, 0, 0xFFFF}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			buf, err := CreateShmBuffer(w, h, w*4+8)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer buf.Close()
+			for y := range h {
+				for x := range w {
+					binary.LittleEndian.PutUint32(buf.Data()[y*buf.Stride+x*4:], tc.pixel)
+				}
+			}
+
+			var out bytes.Buffer
+			if err := EncodeBufferPNG(&out, buf, uint32(tc.format), nil); err != nil {
+				t.Fatal(err)
+			}
+			data := out.Bytes()
+			img, err := png.Decode(bytes.NewReader(data))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if img.Bounds().Dx() != w || img.Bounds().Dy() != h {
+				t.Fatalf("bounds = %v", img.Bounds())
+			}
+			got := color.RGBA64Model.Convert(img.At(w-1, h-1)).(color.RGBA64)
+			if got != tc.want {
+				t.Errorf("pixel = %v, want %v", got, tc.want)
+			}
+			wantDepth := byte(8)
+			if tc.format.Is10Bit() {
+				wantDepth = 16
+			}
+			if depth := data[24]; depth != wantDepth {
+				t.Errorf("bit depth = %d, want %d", depth, wantDepth)
+			}
+		})
+	}
+}