screenshot: configure surfaces before requesting screen captures

core/internal/screenshot/region.go

@@ -82,8 +82,11 @@ type OutputSurface struct {
 	yInverted         bool
 
 	// Triple-buffered render slots
-	slots      [3]*RenderSlot
-	slotsReady bool
+	slots        [3]*RenderSlot
+	slotsReady   bool
+	slotFormat   uint32
+	slotPrefault sync.WaitGroup
+	captured     bool
 
 	shown           *overlay
 	framePending    bool
@@ -96,6 +99,9 @@ type PreCapture struct {
 	screenBufNoCursor *ShmBuffer
 	format            uint32
 	yInverted         bool
+	width             int
+	height            int
+	stride            int
 }
 
 type RegionSelector struct {
@@ -211,10 +217,7 @@ func (r *RegionSelector) Run() (*CaptureResult, bool, error) {
 		return nil, false, fmt.Errorf("roundtrip after protocol check: %w", err)
 	}
 
-	if err := r.preCaptureAllOutputs(); err != nil {
-		return nil, false, fmt.Errorf("pre-capture: %w", err)
-	}
-
+	// unmapped until their first buffer, so the configure round trip overlaps the captures
 	if err := r.createSurfaces(); err != nil {
 		return nil, false, fmt.Errorf("create surfaces: %w", err)
 	}
@@ -225,6 +228,16 @@ func (r *RegionSelector) Run() (*CaptureResult, bool, error) {
 		return nil, false, fmt.Errorf("roundtrip after surfaces: %w", err)
 	}
 
+	if err := r.preCaptureAllOutputs(); err != nil {
+		return nil, false, fmt.Errorf("pre-capture: %w", err)
+	}
+
+	for _, os := range r.surfaces {
+		if os.configured && !os.captured {
+			r.captureForSurface(os)
+		}
+	}
+
 	r.running = true
 	for r.running {
 		if err := r.dispatchOrTick(); err != nil {
@@ -513,6 +526,13 @@ func (r *RegionSelector) preCaptureOutput(output *WaylandOutput, pc *PreCapture,
 			log.Error("invalid stride from compositor", "stride", e.Stride, "width", e.Width, "bpp", bpp)
 			return
 		}
+		if pc.width == 0 {
+			pc.width, pc.height, pc.stride = int(e.Width), int(e.Height), int(e.Stride)
+			pc.format = uint32(capturedFormat.To8Bit())
+			if bpp == 4 && output.transform == TransformNormal {
+				r.prepareSlots(output, pc)
+			}
+		}
 		buf, err := CreateShmBuffer(int(e.Width), int(e.Height), int(e.Stride))
 		if err != nil {
 			log.Error("create screen buffer failed", "err", err)
@@ -835,25 +855,61 @@ func (r *RegionSelector) captureForSurface(os *OutputSurface) {
 	os.screenBufNoCursor = pc.screenBufNoCursor
 	os.screenFormat = pc.format
 	os.yInverted = pc.yInverted
+	if os.screenBuf == nil {
+		return
+	}
+	os.captured = true
 
-	if os.logicalW > 0 && os.screenBuf != nil {
+	if os.logicalW > 0 {
 		os.output.fractionalScale = float64(os.screenBuf.Width) / float64(os.logicalW)
 	}
 
-	r.initRenderBuffer(os)
+	r.ensureSlots(os, os.screenBuf.Width, os.screenBuf.Height, os.screenBuf.Stride, os.screenFormat)
 	r.applyPreSelection(os)
 	r.redrawSurface(os)
 }
 
-func (r *RegionSelector) initRenderBuffer(os *OutputSurface) {
-	if os.screenBuf == nil {
+// prefaulted while the copy is in flight so the first paint does not stall on page faults
+func (r *RegionSelector) prepareSlots(output *WaylandOutput, pc *PreCapture) {
+	for _, os := range r.surfaces {
+		if os.output == output && os.configured {
+			r.ensureSlots(os, pc.width, pc.height, pc.stride, pc.format)
+		}
+	}
+}
+
+func (r *RegionSelector) releaseSlots(os *OutputSurface) {
+	os.slotPrefault.Wait()
+	os.slotsReady = false
+	for i, slot := range os.slots {
+		if slot == nil {
+			continue
+		}
+		slot.wlBuf.Destroy()
+		slot.pool.Destroy()
+		slot.shm.Close()
+		os.slots[i] = nil
+	}
+}
+
+func prefault(buf *ShmBuffer) {
+	data := buf.Data()
+	for i := 0; i < len(data); i += 4096 {
+		data[i] = 0
+	}
+}
+
+func (r *RegionSelector) ensureSlots(os *OutputSurface, width, height, stride int, format uint32) {
+	if first := os.slots[0]; os.slotsReady && first != nil && first.shm.Width == width && first.shm.Height == height && first.shm.Stride == stride && os.slotFormat == format {
 		return
 	}
+	r.releaseSlots(os)
+	os.slotFormat = format
 
 	for i := range 3 {
 		slot := &RenderSlot{}
 
-		buf, err := CreateShmBuffer(os.screenBuf.Width, os.screenBuf.Height, os.screenBuf.Stride)
+		buf, err := CreateShmBuffer(width, height, stride)
 		if err != nil {
 			log.Error("create render slot buffer failed", "err", err)
 			return
@@ -870,7 +926,7 @@ func (r *RegionSelector) initRenderBuffer(os *OutputSurface) {
 
 		// niri latches surface opacity from the first buffer's format
 		// (observed), so slots are ARGB from the start with A=255 when opaque
-		wlBuf, err := pool.CreateBuffer(0, int32(buf.Width), int32(buf.Height), int32(buf.Stride), alphaFormat(os.screenFormat))
+		wlBuf, err := pool.CreateBuffer(0, int32(buf.Width), int32(buf.Height), int32(buf.Stride), alphaFormat(format))
 		if err != nil {
 			log.Error("create render slot wl_buffer failed", "err", err)
 			pool.Destroy()
@@ -890,6 +946,13 @@ func (r *RegionSelector) initRenderBuffer(os *OutputSurface) {
 		os.slots[i] = slot
 	}
 	os.slotsReady = true
+	os.slotPrefault.Add(1)
+	go func() {
+		defer os.slotPrefault.Done()
+		for _, slot := range os.slots {
+			prefault(slot.shm)
+		}
+	}()
 }
 
 func (os *OutputSurface) acquireFreeSlot() *RenderSlot {
@@ -956,6 +1019,7 @@ func (r *RegionSelector) renderSurface(os *OutputSurface) {
 		return
 	}
 	os.redrawQueued = false
+	os.slotPrefault.Wait()
 
 	fullDamage := true
 	var damage []dirtyRect
@@ -971,8 +1035,7 @@ func (r *RegionSelector) renderSurface(os *OutputSurface) {
 		shift := r.shiftHeld && r.selection.hasSelection
 		switch {
 		case !slot.cacheValid(srcBuf, r.selection.dragging, r.showCapturedCursor, r.phase, handles, shift):
-			slot.shm.CopyFrom(srcBuf)
-			r.dimBackground(slot.shm)
+			paintDimmedFrame(slot.shm, srcBuf)
 			effectiveScale := 1.0
 			if os.output != nil {
 				effectiveScale = os.output.effectiveScale()

core/internal/screenshot/region_render.go

@@ -3,9 +3,11 @@ package screenshot
 import (
 	"fmt"
 	"math"
+	"runtime"
 	"slices"
 	"strconv"
 	"strings"
+	"sync"
 )
 
 const (
@@ -237,11 +239,36 @@ func opaqueRow(p []byte) {
 	}
 }
 
-func (r *RegionSelector) dimBackground(renderBuf *ShmBuffer) {
-	data := renderBuf.Data()
-	for y := range renderBuf.Height {
-		dimRow(data[y*renderBuf.Stride:][:renderBuf.Width*4])
+func parallelRows(rows int, fn func(y1, y2 int)) {
+	workers := min(runtime.GOMAXPROCS(0), 8, rows)
+	if workers <= 1 {
+		fn(0, rows)
+		return
 	}
+	chunk := (rows + workers - 1) / workers
+	var wg sync.WaitGroup
+	for y := 0; y < rows; y += chunk {
+		wg.Add(1)
+		go func(y1, y2 int) {
+			defer wg.Done()
+			fn(y1, y2)
+		}(y, min(y+chunk, rows))
+	}
+	wg.Wait()
+}
+
+func paintDimmedFrame(dst, src *ShmBuffer) {
+	dstData, srcData := dst.Data(), src.Data()
+	rowLen := min(dst.Width, src.Width) * 4
+	parallelRows(min(dst.Height, src.Height), func(y1, y2 int) {
+		for y := y1; y < y2; y++ {
+			d := dstData[y*dst.Stride:][:rowLen]
+			s := srcData[y*src.Stride:][:rowLen]
+			for i := 0; i+3 < rowLen; i += 4 {
+				d[i], d[i+1], d[i+2], d[i+3] = dimLUT[s[i]], dimLUT[s[i+1]], dimLUT[s[i+2]], 255
+			}
+		}
+	})
 }
 
 // paintRect copies area from the capture into dst, dimmed or bright.

core/internal/screenshot/region_render_test.go

@@ -83,8 +83,7 @@ func newFrameBuffer(t testing.TB, r *RegionSelector, os *OutputSurface) *ShmBuff
 		t.Fatal(err)
 	}
 	t.Cleanup(func() { buf.Close() })
-	buf.CopyFrom(os.screenBuf)
-	r.dimBackground(buf)
+	paintDimmedFrame(buf, os.screenBuf)
 	return buf
 }
 

core/internal/screenshot/screenshot.go

@@ -192,11 +192,19 @@ func (s *Screenshoter) captureRegion() (*CaptureResult, error) {
 		}
 	}
 
-	if s.config.SelectorHook != nil {
-		s.config.SelectorHook(true)
-		defer s.config.SelectorHook(false)
+	if hook := s.config.SelectorHook; hook != nil {
+		started := make(chan struct{})
+		go func() {
+			hook(true)
+			close(started)
+		}()
+		defer func() {
+			<-started
+			hook(false)
+		}()
 	}
 
+	PrefetchOverlayStyle()
 	selector := NewRegionSelector(s)
 	result, cancelled, err := selector.Run()
 	if err != nil {

core/internal/screenshot/theme.go

@@ -5,6 +5,7 @@ import (
 	"os"
 	"path/filepath"
 	"strings"
+	"sync"
 
 	"github.com/AvengeMedia/DankMaterialShell/core/internal/log"
 	"github.com/AvengeMedia/DankMaterialShell/core/internal/utils"
@@ -25,17 +26,25 @@ type ColorsFile struct {
 	Colors ColorScheme `json:"colors"`
 }
 
-var cachedStyle *OverlayStyle
+var (
+	styleOnce   sync.Once
+	cachedStyle OverlayStyle
+)
+
+// keeps the dconf spawn off the first paint
+func PrefetchOverlayStyle() {
+	go LoadOverlayStyle()
+}
 
 func LoadOverlayStyle() OverlayStyle {
-	if cachedStyle != nil {
-		return *cachedStyle
-	}
+	styleOnce.Do(func() { cachedStyle = loadOverlayStyle() })
+	return cachedStyle
+}
 
+func loadOverlayStyle() OverlayStyle {
 	style := DefaultOverlayStyle
 	colors := loadColorsFile()
 	if colors == nil {
-		cachedStyle = &style
 		return style
 	}
 
@@ -54,7 +63,6 @@ func LoadOverlayStyle() OverlayStyle {
 		style.AccentR, style.AccentG, style.AccentB = accent[0], accent[1], accent[2]
 	}
 
-	cachedStyle = &style
 	return style
 }