feat: add screen freezing support

Author: dy-tea

README.md

@@ -27,6 +27,7 @@ v .              # build
 - Include or exclude cursor
 - PNG, JXL, QOI and PPM support
 - HDR support
+- Screen freezing support
 
 ### Usage
 ```sh
@@ -39,6 +40,8 @@ mrpenishot -g "100,200 300x400"                 # geometry
 mrpenishot -g "$(slurp)"                        # geometry from slurp
 mrpenishot -t "$(awmsg t f | jq -j '.foreign')" # capture toplevel in awm
 mrpenishot - | wl-copy                          # copy image to clipboard
+mrpenishot -F 'sleep 2'                         # freeze for 2 seconds and screenshot
+mrpenishot -F '' -g '$(slurp)'                  # screenshot and freeze until slurp exits, geometry from slurp
 ```
 
 ### References

build.vsh

@@ -18,15 +18,23 @@ fn program_installed(name string) {
 	}
 }
 
-fn pkg_installed(name string, version ?f64) {
-	ins := sh('pkg-config --modversion ${name}').f64()
-	if ins == 0 {
+fn pkg_installed(name string, version ?string) {
+	ins_str := sh('pkg-config --modversion ${name}')
+	ins := ins_str.trim_space().split('.').map(|x| x.int())
+	if ins.len == 0 {
 		err('package `${name}` not found')
 	} else {
 		if ver := version {
-			if ins < ver {
-				err('package `${name}` version must be >= ${ver}')
+		  	exp := ver.split('.').map(|x| x.int())
+			if exp.len != ins.len {
+			  err('package version lengths `${ver}` and `${ins_str}` differ')
 			}
+		  	for i in 0..ins.len {
+				if ins[i] > exp[i] {
+				  	return
+				}
+			}
+			err('package `${name}` version must be >= ${ver}')
 		}
 	}
 }
@@ -45,11 +53,11 @@ if arguments().contains('clean') {
 program_installed('pkg-config')
 program_installed('wayland-scanner')
 
-pkg_installed('wayland-protocols', 1.41)
+pkg_installed('wayland-protocols', '1.41')
 pkg_installed('wayland-client', none)
 pkg_installed('pixman-1', none)
 pkg_installed('libjxl', none)
-pkg_installed('libpng16', none)
+pkg_installed('libpng16', '2.5.0')
 
 // build vscanner
 sh('v ${vscanner_dir}')
@@ -59,11 +67,14 @@ wl_proto_dir := sh('pkg-config --variable=pkgdatadir wayland-protocols').trim_sp
 
 protocols := [
 	wl_dir + '/wayland.xml',
+	wl_proto_dir + '/stable/xdg-shell/xdg-shell.xml',
+	wl_proto_dir + '/stable/viewporter/viewporter.xml',
 	wl_proto_dir + '/staging/color-management/color-management-v1.xml',
 	wl_proto_dir + '/staging/ext-image-capture-source/ext-image-capture-source-v1.xml',
 	wl_proto_dir + '/staging/ext-image-copy-capture/ext-image-copy-capture-v1.xml',
 	wl_proto_dir + '/staging/ext-foreign-toplevel-list/ext-foreign-toplevel-list-v1.xml',
 	wl_proto_dir + '/unstable/xdg-output/xdg-output-unstable-v1.xml',
+	'protocols/wlr-layer-shell-unstable-v1.xml'
 ]
 
 sh('mkdir -p ${protocols_dir}')

main.v

@@ -71,6 +71,7 @@ fn main() {
 	passed_geometry := fp.string('geometry', `g`, '', 'geometry in the format "400,500 200x300"')
 	output_name := fp.string('output', `o`, '', 'name of output to screenshot')
 	toplevel_identifier := fp.string('toplevel', `t`, '', 'use a toplevel as the screenshot source by its identifier')
+	freeze_screen_cmd := fp.string('freeze', `F`, '', 'freeze the screen until passed command finishes or until the -g command finishes')
 	additional_args := fp.finalize() or {
 		eprintln(err)
 		println(fp.usage())
@@ -95,12 +96,16 @@ fn main() {
 		name
 	}
 
-	// get geometry if passed
-	mut geometry := if passed_geometry == '' {
-		Geometry{}
-	} else {
-		Geometry.new(passed_geometry) or { panic('invalid geometry') }
+	// check if geometry needs command execution (prefix with $)
+	geometry_is_cmd := passed_geometry.starts_with('$') && !passed_geometry.starts_with('$-')
+	mut geometry_cmd := ''
+	if geometry_is_cmd {
+		geometry_cmd = passed_geometry.trim_left('$')
+		if freeze_screen_cmd != '' {
+			panic('ERROR: cannot use both --freeze and command prefix in -g')
+		}
 	}
+	mut geometry := Geometry{}
 
 	// mutual exclusion
 	if output_name != '' && geometry != Geometry{} {
@@ -153,6 +158,18 @@ fn main() {
 	if state.outputs.len == 0 {
 		panic('no outputs found')
 	}
+	needs_freeze := freeze_screen_cmd != '' || geometry_is_cmd
+	if needs_freeze {
+		if state.compositor == none {
+			panic('wl_compositor not supported by compositor')
+		}
+		if state.wp_viewporter == none {
+			panic('wp_viewporter not supported by compositor')
+		}
+		if state.wlr_layer_shell_v1 == none {
+			panic('zwlr_layer_shell_v1 not supported by compositor')
+		}
+	}
 
 	// init output manager
 	if mut manager := state.zxdg_output_manager_v1 {
@@ -180,10 +197,10 @@ fn main() {
 			mut cm_output := color_manager.get_output(output.wl_output.proxy)
 			cm_output.add_listener(&cm_output_listener, state)
 			mut description := cm_output.get_image_description()
-      description.add_listener(&cm_image_description_listener, output.state)
-      if C.wl_display_roundtrip(display_proxy) < 0 {
-        panic('wl_display_roundtrip failed')
-      }
+		    description.add_listener(&cm_image_description_listener, output.state)
+		    if C.wl_display_roundtrip(display_proxy) < 0 {
+			    panic('wl_display_roundtrip failed')
+		    }
 		}
 	}
 
@@ -211,7 +228,7 @@ fn main() {
 	} else {
 		// capture output
 		for output in state.outputs {
-			if geometry != Geometry{} && !geometry.intersect(output.logical_geometry) {
+			if !geometry_is_cmd && geometry != Geometry{} && !geometry.intersect(output.logical_geometry) {
 				continue
 			}
 			if output.logical_scale > scale {
@@ -224,12 +241,98 @@ fn main() {
 		panic('no captures found')
 	}
 
-	// dispatch captures
+	// dispatch initial captures to get buffer data for overlays
 	mut done := false
 	expected_cm := if state.wp_color_manager_v1 != none { state.outputs.len } else { 0 }
 	for !done && C.wl_display_dispatch(display_proxy) != -1{
 		done = state.n_done == state.captures.len && state.n_cm_done >= expected_cm
 	}
+
+	// run geometry command with freeze
+	if geometry_is_cmd {
+		mut geometry_cmd_overlays := []&Overlay{}
+		for capture in state.captures {
+			overlay := Overlay.new(capture)
+			geometry_cmd_overlays << overlay
+		}
+
+		// run command in background thread
+		result_ch := chan string{cap: 1}
+		spawn fn (cmd string, ch chan string) {
+			result := os.execute(cmd)
+			ch <- result.output
+		}(geometry_cmd, result_ch)
+
+		// process Wayland events while waiting for command
+		mut cmd_output := ''
+		for {
+			C.wl_display_dispatch(display_proxy)
+			C.wl_display_flush(display_proxy)
+			select {
+				output := <-result_ch {
+					cmd_output = output
+					break
+				}
+			}
+		}
+		for mut overlay in geometry_cmd_overlays {
+			overlay.destroy()
+		}
+		geometry = Geometry.new(cmd_output.trim('\n')) or { panic('invalid geometry from command') }
+		state.captures = []
+		state.n_done = 0
+		state.n_cm_done = 0
+		for output in state.outputs {
+			if geometry != Geometry{} && !geometry.intersect(output.logical_geometry) {
+				continue
+			}
+			if output.logical_scale > scale {
+				scale = output.logical_scale
+			}
+			state.capture_output(output, include_cursor)
+		}
+		if state.captures.len == 0 {
+			panic('no captures found after geometry command')
+		}
+		// re-dispatch for new captures
+		done = false
+		for !done && C.wl_display_dispatch(display_proxy) != -1{
+			done = state.n_done == state.captures.len && state.n_cm_done >= expected_cm
+		}
+	}
+
+	// freeze screen overlay
+	mut overlays := []&Overlay{}
+	if freeze_screen_cmd != '' {
+		for capture in state.captures {
+			overlay := Overlay.new(capture)
+			overlays << overlay
+		}
+
+		// run command in background thread
+		ch := chan int{cap: 1}
+		spawn fn (cmd string, result_ch chan int) {
+			result := os.execute(cmd)
+			result_ch <- result.exit_code
+		}(freeze_screen_cmd, ch)
+
+		// process events while waiting for command
+		for {
+			C.wl_display_dispatch(display_proxy)
+			C.wl_display_flush(display_proxy)
+			select {
+				exit_code := <-ch {
+					for mut overlay in overlays {
+						overlay.destroy()
+					}
+					if exit_code != 0 {
+						eprintln('freeze command exited with code ${exit_code}')
+					}
+					break
+				}
+			}
+		}
+	}
 	if geometry == Geometry{0, 0, 0, 0} {
 		geometry = state.get_extents()
 	}
@@ -307,5 +410,8 @@ fn main() {
 	if mut manager := state.zxdg_output_manager_v1 {
 		manager.destroy()
 	}
+	if mut manager := state.wp_viewporter {
+		manager.destroy()
+	}
 	C.wl_display_disconnect(display_proxy)
 }

overlay.v

@@ -0,0 +1,107 @@
+module main
+
+import wl
+import protocols.wlr_layer_shell_unstable_v1 as ls
+
+// needed for xdg_popup interface
+import protocols.xdg_shell as _
+
+fn layer_surface_v1_configure(mut overlay Overlay, _ voidptr, serial u32, _ u32, _ u32) {
+	if mut layer_surface := overlay.layer_surface_v1 {
+		layer_surface.ack_configure(serial)
+	}
+	if mut surface := overlay.surface {
+		surface.commit()
+	}
+}
+
+fn layer_surface_v1_closed(mut overlay Overlay, _ voidptr) {
+	panic('Layer surface died unexpectedly')
+}
+
+const layer_surface_listener = C.zwlr_layer_surface_v1_listener{
+	configure: layer_surface_v1_configure
+	closed:    layer_surface_v1_closed
+}
+
+const surface_listener = C.wl_surface_listener{
+	enter:                      fn (_ voidptr, _ voidptr, _ voidptr) {}
+	leave:                      fn (_ voidptr, _ voidptr, _ voidptr) {}
+	preferred_buffer_scale:     fn (_ voidptr, _ voidptr, _ int) {}
+	preferred_buffer_transform: fn (_ voidptr, _ voidptr, _ u32) {}
+}
+
+fn Overlay.new(capture &Capture) &Overlay {
+	mut overlay := &Overlay{
+		capture:          capture
+		layer_surface_v1: none
+		surface:          none
+	}
+	mut compositor := capture.state.compositor or { panic('Failed to get compositor') }
+
+	mut surface := compositor.create_surface()
+	if surface.proxy == unsafe { nil } {
+		panic('Failed to create surface')
+	}
+	overlay.surface = surface
+	surface.add_listener(&surface_listener, overlay)
+
+	mut wp_viewporter := capture.state.wp_viewporter or { panic('No viewporter init') }
+
+	mut wp_viewport := wp_viewporter.get_viewport(surface.proxy)
+	overlay.wp_viewport = wp_viewport
+	if wp_viewport == unsafe { nil } {
+		panic('Failed to create viewport')
+	}
+	wp_viewport.set_destination(capture.logical_geometry.width, capture.logical_geometry.height)
+	wp_viewport.set_source(wl.wl_fixed_from_int(0), wl.wl_fixed_from_int(0), wl.wl_fixed_from_int(int(capture.buffer_width)),
+		wl.wl_fixed_from_int(int(capture.buffer_height)))
+
+	if output := capture.output {
+		mut wlr_layer_shell := capture.state.wlr_layer_shell_v1 or { panic('No layer shell init') }
+
+		mut layer_surface := wlr_layer_shell.get_layer_surface(surface.proxy, output.wl_output.proxy,
+			u32(ls.ZwlrLayerShellV1_Layer.overlay), 'mrpenishot')
+		overlay.layer_surface_v1 = layer_surface
+		if layer_surface == unsafe { nil } {
+			panic('Failed to get layer surface')
+		}
+		layer_surface.add_listener(&layer_surface_listener, overlay)
+
+		layer_surface.set_size(u32(output.logical_geometry.width), u32(output.logical_geometry.height))
+		layer_surface.set_anchor(u32(int(ls.ZwlrLayerSurfaceV1_Anchor.top) | int(ls.ZwlrLayerSurfaceV1_Anchor.bottom) | int(ls.ZwlrLayerSurfaceV1_Anchor.left) | int(ls.ZwlrLayerSurfaceV1_Anchor.right)))
+		layer_surface.set_exclusive_zone(-1)
+
+		surface.commit()
+		C.wl_display_roundtrip(overlay.capture.state.display.proxy)
+
+		if buffer := capture.buffer {
+			mut shm := capture.state.shm or { panic('No shm init') }
+			stride := int(get_min_stride(buffer.shm_format, u32(buffer.width)))
+			mut overlay_buffer := Buffer.new(mut shm, buffer.shm_format, buffer.width, buffer.height, stride)
+			unsafe {
+				C.memcpy(overlay_buffer.data, buffer.data, buffer.size)
+			}
+			overlay.buffer = overlay_buffer
+			surface.attach(overlay_buffer.wl_buffer.proxy, 0, 0)
+			surface.commit()
+		}
+	}
+
+	return overlay
+}
+
+fn (mut overlay Overlay) destroy() {
+	if mut layer_surface := overlay.layer_surface_v1 {
+		layer_surface.destroy()
+	}
+	if mut surface := overlay.surface {
+		surface.destroy()
+	}
+	if mut wp_viewport := overlay.wp_viewport {
+		wp_viewport.destroy()
+	}
+	if mut buffer := overlay.buffer {
+		buffer.destroy()
+	}
+}

protocols/.gitignore

@@ -0,0 +1,3 @@
+*
+!*.xml
+!.gitignore