colorpick/screenshot: make color-format aware

core/cmd/dms/commands_screenshot.go

@@ -178,7 +178,7 @@ func runScreenshot(config screenshot.Config) {
 	}
 
 	if config.Stdout {
-		if err := writeImageToStdout(result.Buffer, config.Format, config.Quality); err != nil {
+		if err := writeImageToStdout(result.Buffer, config.Format, config.Quality, result.Format); err != nil {
 			fmt.Fprintf(os.Stderr, "Error writing to stdout: %v\n", err)
 			os.Exit(1)
 		}
@@ -199,7 +199,7 @@ func runScreenshot(config screenshot.Config) {
 		}
 
 		filePath = filepath.Join(outputDir, filename)
-		if err := screenshot.WriteToFile(result.Buffer, filePath, config.Format, config.Quality); err != nil {
+		if err := screenshot.WriteToFileWithFormat(result.Buffer, filePath, config.Format, config.Quality, result.Format); err != nil {
 			fmt.Fprintf(os.Stderr, "Error writing file: %v\n", err)
 			os.Exit(1)
 		}
@@ -207,7 +207,7 @@ func runScreenshot(config screenshot.Config) {
 	}
 
 	if config.Clipboard {
-		if err := copyImageToClipboard(result.Buffer, config.Format, config.Quality); err != nil {
+		if err := copyImageToClipboard(result.Buffer, config.Format, config.Quality, result.Format); err != nil {
 			fmt.Fprintf(os.Stderr, "Error copying to clipboard: %v\n", err)
 			os.Exit(1)
 		}
@@ -217,7 +217,7 @@ func runScreenshot(config screenshot.Config) {
 	}
 
 	if config.Notify {
-		thumbData, thumbW, thumbH := bufferToRGBThumbnail(result.Buffer, 256)
+		thumbData, thumbW, thumbH := bufferToRGBThumbnail(result.Buffer, 256, result.Format)
 		screenshot.SendNotification(screenshot.NotifyResult{
 			FilePath:  filePath,
 			Clipboard: config.Clipboard,
@@ -228,11 +228,11 @@ func runScreenshot(config screenshot.Config) {
 	}
 }
 
-func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, quality int) error {
+func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, quality int, pixelFormat uint32) error {
 	var mimeType string
 	var data bytes.Buffer
 
-	img := screenshot.BufferToImage(buf)
+	img := screenshot.BufferToImageWithFormat(buf, pixelFormat)
 
 	switch format {
 	case screenshot.FormatJPEG:
@@ -252,8 +252,8 @@ func copyImageToClipboard(buf *screenshot.ShmBuffer, format screenshot.Format, q
 	return cmd.Run()
 }
 
-func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, quality int) error {
+func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, quality int, pixelFormat uint32) error {
-	img := screenshot.BufferToImage(buf)
+	img := screenshot.BufferToImageWithFormat(buf, pixelFormat)
 
 	switch format {
 	case screenshot.FormatJPEG:
@@ -263,7 +263,7 @@ func writeImageToStdout(buf *screenshot.ShmBuffer, format screenshot.Format, qua
 	}
 }
 
-func bufferToRGBThumbnail(buf *screenshot.ShmBuffer, maxSize int) ([]byte, int, int) {
+func bufferToRGBThumbnail(buf *screenshot.ShmBuffer, maxSize int, pixelFormat uint32) ([]byte, int, int) {
 	srcW, srcH := buf.Width, buf.Height
 	scale := 1.0
 	if srcW > maxSize || srcH > maxSize {
@@ -285,6 +285,7 @@ func bufferToRGBThumbnail(buf *screenshot.ShmBuffer, maxSize int) ([]byte, int,
 
 	data := buf.Data()
 	rgb := make([]byte, dstW*dstH*3)
+	swapRB := pixelFormat == uint32(screenshot.FormatARGB8888) || pixelFormat == uint32(screenshot.FormatXRGB8888) || pixelFormat == 0
 
 	for y := 0; y < dstH; y++ {
 		srcY := int(float64(y) / scale)
@@ -299,9 +300,15 @@ func bufferToRGBThumbnail(buf *screenshot.ShmBuffer, maxSize int) ([]byte, int,
 			si := srcY*buf.Stride + srcX*4
 			di := (y*dstW + x) * 3
 			if si+2 < len(data) {
-				rgb[di+0] = data[si+2] // R
+				if swapRB {
-				rgb[di+1] = data[si+1] // G
+					rgb[di+0] = data[si+2]
-				rgb[di+2] = data[si+0] // B
+					rgb[di+1] = data[si+1]
+					rgb[di+2] = data[si+0]
+				} else {
+					rgb[di+0] = data[si+0]
+					rgb[di+1] = data[si+1]
+					rgb[di+2] = data[si+2]
+				}
 			}
 		}
 	}

core/internal/colorpicker/picker.go

@@ -34,15 +34,19 @@ type Output struct {
 }
 
 type LayerSurface struct {
-	output     *Output
+	output      *Output
-	state      *SurfaceState
+	state       *SurfaceState
-	wlSurface  *client.Surface
+	wlSurface   *client.Surface
-	layerSurf  *wlr_layer_shell.ZwlrLayerSurfaceV1
+	layerSurf   *wlr_layer_shell.ZwlrLayerSurfaceV1
-	viewport   *wp_viewporter.WpViewport
+	viewport    *wp_viewporter.WpViewport
-	wlPool     *client.ShmPool
+	wlPool      *client.ShmPool
-	wlBuffer   *client.Buffer
+	wlBuffer    *client.Buffer
-	configured bool
+	bufferBusy  bool
-	hidden     bool
+	oldPool     *client.ShmPool
+	oldBuffer   *client.Buffer
+	scopyBuffer *client.Buffer
+	configured  bool
+	hidden      bool
 }
 
 type Picker struct {
@@ -463,6 +467,12 @@ func (p *Picker) captureForSurface(ls *LayerSurface) {
 			return
 		}
 
+		if ls.scopyBuffer != nil {
+			ls.scopyBuffer.Destroy()
+		}
+		ls.scopyBuffer = wlBuffer
+		wlBuffer.SetReleaseHandler(func(e client.BufferReleaseEvent) {})
+
 		if err := frame.Copy(wlBuffer); err != nil {
 			log.Error("failed to copy frame", "err", err)
 		}
@@ -489,7 +499,6 @@ func (p *Picker) captureForSurface(ls *LayerSurface) {
 func (p *Picker) redrawSurface(ls *LayerSurface) {
 	var renderBuf *ShmBuffer
 	if ls.hidden {
-		// When hidden, just show the screenshot without overlay
 		renderBuf = ls.state.RedrawScreenOnly()
 	} else {
 		renderBuf = ls.state.Redraw()
@@ -498,27 +507,38 @@ func (p *Picker) redrawSurface(ls *LayerSurface) {
 		return
 	}
 
-	if ls.wlPool != nil {
+	if ls.oldBuffer != nil {
-		ls.wlPool.Destroy()
+		ls.oldBuffer.Destroy()
-		ls.wlPool = nil
+		ls.oldBuffer = nil
 	}
-	if ls.wlBuffer != nil {
+	if ls.oldPool != nil {
-		ls.wlBuffer.Destroy()
+		ls.oldPool.Destroy()
-		ls.wlBuffer = nil
+		ls.oldPool = nil
 	}
 
+	ls.oldPool = ls.wlPool
+	ls.oldBuffer = ls.wlBuffer
+	ls.wlPool = nil
+	ls.wlBuffer = nil
+
 	pool, err := p.shm.CreatePool(renderBuf.Fd(), int32(renderBuf.Size()))
 	if err != nil {
 		return
 	}
 	ls.wlPool = pool
 
-	wlBuffer, err := pool.CreateBuffer(0, int32(renderBuf.Width), int32(renderBuf.Height), int32(renderBuf.Stride), uint32(FormatARGB8888))
+	wlBuffer, err := pool.CreateBuffer(0, int32(renderBuf.Width), int32(renderBuf.Height), int32(renderBuf.Stride), uint32(ls.state.ScreenFormat()))
 	if err != nil {
 		return
 	}
 	ls.wlBuffer = wlBuffer
 
+	lsRef := ls
+	wlBuffer.SetReleaseHandler(func(e client.BufferReleaseEvent) {
+		lsRef.bufferBusy = false
+	})
+	ls.bufferBusy = true
+
 	logicalW, logicalH := ls.state.LogicalSize()
 	if logicalW == 0 || logicalH == 0 {
 		logicalW = int(ls.output.width)
@@ -533,30 +553,13 @@ func (p *Picker) redrawSurface(ls *LayerSurface) {
 	if ls.viewport != nil {
 		srcW := float64(renderBuf.Width) / float64(scale)
 		srcH := float64(renderBuf.Height) / float64(scale)
-		if err := ls.viewport.SetSource(0, 0, srcW, srcH); err != nil {
+		_ = ls.viewport.SetSource(0, 0, srcW, srcH)
-			log.Warn("failed to set viewport source", "err", err)
+		_ = ls.viewport.SetDestination(int32(logicalW), int32(logicalH))
-		}
-		if err := ls.viewport.SetDestination(int32(logicalW), int32(logicalH)); err != nil {
-			log.Warn("failed to set viewport destination", "err", err)
-		}
-		if err := ls.wlSurface.SetBufferScale(scale); err != nil {
-			log.Warn("failed to set buffer scale", "err", err)
-		}
-	} else {
-		if err := ls.wlSurface.SetBufferScale(scale); err != nil {
-			log.Warn("failed to set buffer scale", "err", err)
-		}
-	}
-
-	if err := ls.wlSurface.Attach(wlBuffer, 0, 0); err != nil {
-		log.Warn("failed to attach buffer", "err", err)
-	}
-	if err := ls.wlSurface.Damage(0, 0, int32(logicalW), int32(logicalH)); err != nil {
-		log.Warn("failed to damage surface", "err", err)
-	}
-	if err := ls.wlSurface.Commit(); err != nil {
-		log.Warn("failed to commit surface", "err", err)
 	}
+	_ = ls.wlSurface.SetBufferScale(scale)
+	_ = ls.wlSurface.Attach(wlBuffer, 0, 0)
+	_ = ls.wlSurface.Damage(0, 0, int32(logicalW), int32(logicalH))
+	_ = ls.wlSurface.Commit()
 
 	ls.state.SwapBuffers()
 }
@@ -599,9 +602,14 @@ func (p *Picker) setupPointerHandlers() {
 			log.Debug("failed to hide cursor", "err", err)
 		}
 
+		if e.Surface == nil {
+			return
+		}
+
 		p.activeSurface = nil
+		surfaceID := e.Surface.ID()
 		for _, ls := range p.surfaces {
-			if ls.wlSurface.ID() == e.Surface.ID() {
+			if ls.wlSurface.ID() == surfaceID {
 				p.activeSurface = ls
 				break
 			}
@@ -610,7 +618,6 @@ func (p *Picker) setupPointerHandlers() {
 			return
 		}
 
-		// If surface was hidden, mark it as visible again
 		if p.activeSurface.hidden {
 			p.activeSurface.hidden = false
 		}
@@ -620,8 +627,12 @@ func (p *Picker) setupPointerHandlers() {
 	})
 
 	p.pointer.SetLeaveHandler(func(e client.PointerLeaveEvent) {
+		if e.Surface == nil {
+			return
+		}
+		surfaceID := e.Surface.ID()
 		for _, ls := range p.surfaces {
-			if ls.wlSurface.ID() == e.Surface.ID() {
+			if ls.wlSurface.ID() == surfaceID {
 				p.hideSurface(ls)
 				break
 			}
@@ -654,6 +665,15 @@ func (p *Picker) setupKeyboardHandlers() {
 
 func (p *Picker) cleanup() {
 	for _, ls := range p.surfaces {
+		if ls.scopyBuffer != nil {
+			ls.scopyBuffer.Destroy()
+		}
+		if ls.oldBuffer != nil {
+			ls.oldBuffer.Destroy()
+		}
+		if ls.oldPool != nil {
+			ls.oldPool.Destroy()
+		}
 		if ls.wlBuffer != nil {
 			ls.wlBuffer.Destroy()
 		}

core/internal/colorpicker/shm.go

@@ -9,6 +9,10 @@ func CreateShmBuffer(width, height, stride int) (*ShmBuffer, error) {
 }
 
 func GetPixelColor(buf *ShmBuffer, x, y int) Color {
+	return GetPixelColorWithFormat(buf, x, y, FormatARGB8888)
+}
+
+func GetPixelColorWithFormat(buf *ShmBuffer, x, y int, format PixelFormat) Color {
 	if x < 0 || x >= buf.Width || y < 0 || y >= buf.Height {
 		return Color{}
 	}
@@ -19,6 +23,14 @@ func GetPixelColor(buf *ShmBuffer, x, y int) Color {
 		return Color{}
 	}
 
+	if format == FormatABGR8888 || format == FormatXBGR8888 {
+		return Color{
+			R: data[offset],
+			G: data[offset+1],
+			B: data[offset+2],
+			A: data[offset+3],
+		}
+	}
 	return Color{
 		B: data[offset],
 		G: data[offset+1],

core/internal/colorpicker/state.go

@@ -100,6 +100,12 @@ func (s *SurfaceState) ScreenBuffer() *ShmBuffer {
 	return s.screenBuf
 }
 
+func (s *SurfaceState) ScreenFormat() PixelFormat {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.screenFormat
+}
+
 func (s *SurfaceState) OnScreencopyFlags(flags uint32) {
 	s.mu.Lock()
 	s.yInverted = (flags & 1) != 0
@@ -263,7 +269,7 @@ func (s *SurfaceState) Redraw() *ShmBuffer {
 	px = clamp(px, 0, dst.Width-1)
 	py = clamp(py, 0, dst.Height-1)
 
-	picked := GetPixelColor(s.screenBuf, px, py)
+	picked := GetPixelColorWithFormat(s.screenBuf, px, py, s.screenFormat)
 
 	drawMagnifier(
 		dst.Data(), dst.Stride, dst.Width, dst.Height,
@@ -313,7 +319,7 @@ func (s *SurfaceState) PickColor() (Color, bool) {
 		sy = s.screenBuf.Height - 1 - sy
 	}
 
-	return GetPixelColor(s.screenBuf, sx, sy), true
+	return GetPixelColorWithFormat(s.screenBuf, sx, sy, s.screenFormat), true
 }
 
 func (s *SurfaceState) Destroy() {

core/internal/screenshot/encode.go

@@ -13,8 +13,15 @@ import (
 )
 
 func BufferToImage(buf *ShmBuffer) *image.RGBA {
+	return BufferToImageWithFormat(buf, uint32(FormatARGB8888))
+}
+
+func BufferToImageWithFormat(buf *ShmBuffer, format uint32) *image.RGBA {
 	img := image.NewRGBA(image.Rect(0, 0, buf.Width, buf.Height))
 	data := buf.Data()
+
+	swapRB := format == uint32(FormatARGB8888) || format == uint32(FormatXRGB8888) || format == 0
+
 	for y := 0; y < buf.Height; y++ {
 		srcOff := y * buf.Stride
 		dstOff := y * img.Stride
@@ -24,10 +31,16 @@ func BufferToImage(buf *ShmBuffer) *image.RGBA {
 			if si+3 >= len(data) || di+3 >= len(img.Pix) {
 				continue
 			}
-			img.Pix[di+0] = data[si+2] // R
+			if swapRB {
-			img.Pix[di+1] = data[si+1] // G
+				img.Pix[di+0] = data[si+2]
-			img.Pix[di+2] = data[si+0] // B
+				img.Pix[di+1] = data[si+1]
-			img.Pix[di+3] = 255        // A
+				img.Pix[di+2] = data[si+0]
+			} else {
+				img.Pix[di+0] = data[si+0]
+				img.Pix[di+1] = data[si+1]
+				img.Pix[di+2] = data[si+2]
+			}
+			img.Pix[di+3] = 255
 		}
 	}
 	return img
@@ -162,13 +175,17 @@ func expandHome(path string) string {
 }
 
 func WriteToFile(buf *ShmBuffer, path string, format Format, quality int) error {
+	return WriteToFileWithFormat(buf, path, format, quality, uint32(FormatARGB8888))
+}
+
+func WriteToFileWithFormat(buf *ShmBuffer, path string, format Format, quality int, pixelFormat uint32) error {
 	f, err := os.Create(path)
 	if err != nil {
 		return err
 	}
 	defer f.Close()
 
-	img := BufferToImage(buf)
+	img := BufferToImageWithFormat(buf, pixelFormat)
 	switch format {
 	case FormatJPEG:
 		return EncodeJPEG(f, img, quality)

core/internal/screenshot/region.go

@@ -160,14 +160,17 @@ func (r *RegionSelector) Run() (*CaptureResult, bool, error) {
 	}
 
 	yInverted := false
+	var format uint32
 	if r.selection.surface != nil {
 		yInverted = r.selection.surface.yInverted
+		format = r.selection.surface.screenFormat
 	}
 
 	return &CaptureResult{
 		Buffer:    r.capturedBuffer,
-		Region:    r.result, // Global coords for saving last region
+		Region:    r.result,
 		YInverted: yInverted,
+		Format:    format,
 	}, false, nil
 }
 

core/internal/screenshot/region_input.go

@@ -98,9 +98,10 @@ func (r *RegionSelector) setupPointerHandlers() {
 		case 0x110: // BTN_LEFT
 			switch e.State {
 			case 1: // pressed
+				r.preSelect = Region{}
 				r.selection.hasSelection = true
 				r.selection.dragging = true
-				r.selection.surface = r.activeSurface // Lock to this surface
+				r.selection.surface = r.activeSurface
 				r.selection.anchorX = r.pointerX
 				r.selection.anchorY = r.pointerY
 				r.selection.currentX = r.pointerX

core/internal/screenshot/screenshot.go

@@ -25,6 +25,7 @@ type CaptureResult struct {
 	Buffer    *ShmBuffer
 	Region    Region
 	YInverted bool
+	Format    uint32
 }
 
 type Screenshoter struct {
@@ -206,6 +207,7 @@ func (s *Screenshoter) captureAllScreens() (*CaptureResult, error) {
 
 	composite.Clear()
 
+	var format uint32
 	for _, output := range outputs {
 		result, err := s.captureWholeOutput(output)
 		if err != nil {
@@ -213,6 +215,9 @@ func (s *Screenshoter) captureAllScreens() (*CaptureResult, error) {
 			continue
 		}
 
+		if format == 0 {
+			format = result.Format
+		}
 		s.blitBuffer(composite, result.Buffer, int(output.x-minX), int(output.y-minY), result.YInverted)
 		result.Buffer.Close()
 	}
@@ -220,6 +225,7 @@ func (s *Screenshoter) captureAllScreens() (*CaptureResult, error) {
 	return &CaptureResult{
 		Buffer: composite,
 		Region: Region{X: minX, Y: minY, Width: totalW, Height: totalH},
+		Format: format,
 	}, nil
 }
 
@@ -379,6 +385,7 @@ func (s *Screenshoter) processFrame(frame *wlr_screencopy.ZwlrScreencopyFrameV1,
 		Buffer:    buf,
 		Region:    region,
 		YInverted: yInverted,
+		Format:    uint32(format),
 	}, nil
 }