diff --git a/Cargo.toml b/Cargo.toml index 6fd230a..62d6954 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ edition = "2021" [dependencies] anyhow = "1.0.93" bytemuck = "1.20.0" -clap = "4.5.21" +clap = { version = "4.5.21", features = ["cargo", "derive"] } env_logger = "0.11.5" image = "0.25.5" libc = "0.2.165" diff --git a/src/main.rs b/src/main.rs index df53762..4e8ea8d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,14 @@ use std::{ collections::HashMap, + default, ffi::CString, io::BufWriter, os::fd::{AsFd, AsRawFd, FromRawFd}, }; +use clap::{ + arg, builder::TypedValueParser, command, value_parser, ArgAction, Parser, +}; use image::{math::Rect, GenericImage, GenericImageView}; use rand::Rng; use wayland_client::{ @@ -20,8 +24,12 @@ use wayland_client::{ Connection, Dispatch, Proxy, QueueHandle, WEnum, }; -use wayland_protocols::xdg::xdg_output::zv1::client::{ - zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, +use wayland_protocols::{ + wp::input_timestamps::zv1::client::__interfaces::zwp_input_timestamps_manager_v1_interface, + xdg::xdg_output::zv1::client::{ + zxdg_output_manager_v1::ZxdgOutputManagerV1, + zxdg_output_v1::ZxdgOutputV1, + }, }; use wayland_protocols_wlr::screencopy::v1::client::{ zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1}, @@ -43,11 +51,15 @@ enum FrameData { Copying { buffer: WlBuffer, info: BufferInfo, + rect: Rect, offset: usize, + scale: f32, }, Copied { info: BufferInfo, + rect: Rect, offset: usize, + scale: f32, }, Done, } @@ -70,6 +82,7 @@ struct Output { logical_position: Option<(i32, i32)>, physical_position: Option<(i32, i32)>, transform: Option, + name: Option, scale: Option, frame: Frame, } @@ -89,10 +102,16 @@ impl Output { output_data: U, xdg_manager: &ZxdgOutputManagerV1, screencopy: &ZwlrScreencopyManagerV1, + with_cursor: bool, ) -> Self { let output = registry.bind(name, version, &qh, output_data); let xdg_output = xdg_manager.get_xdg_output(&output, &qh, output.id()); - let frame = screencopy.capture_output(0, &output, &qh, output.id()); + let frame = screencopy.capture_output( + with_cursor as i32, + &output, + &qh, + output.id(), + ); Output { output, xdg_output, @@ -102,6 +121,7 @@ impl Output { scale: Some(1), logical_position: None, physical_size: None, + name: None, frame: Frame { frame, inverted_y: false, @@ -113,6 +133,7 @@ impl Output { #[derive(Debug)] struct App { + with_cursor: bool, screencopy_mgr: ZwlrScreencopyManagerV1, xdg_manager: ZxdgOutputManagerV1, outputs: HashMap, @@ -170,11 +191,19 @@ impl Dispatch for App { if let FrameData::Copying { buffer, info, + rect, offset, - } = core::mem::replace(&mut output.frame.data, FrameData::Done) + scale, + } = + core::mem::replace(&mut output.frame.data, FrameData::Done) { buffer.destroy(); - output.frame.data = FrameData::Copied { info, offset }; + output.frame.data = FrameData::Copied { + info, + offset, + rect, + scale, + }; } } zwlr_screencopy_frame_v1::Event::Failed => { @@ -272,6 +301,9 @@ impl Dispatch for App { wl_output::Event::Scale { factor } => { output.scale = Some(factor); } + wl_output::Event::Name { name } => { + output.name = Some(name); + } _ => {} } } @@ -332,21 +364,25 @@ impl Dispatch for App { name, &state.xdg_manager, &state.screencopy_mgr, + state.with_cursor, ); state.outputs.insert(output.output.id(), output); } - wl_registry::Event::GlobalRemove { name } => globals.with_list(|globals| { - if globals - .iter() - .find(|global| global.name == name) - .map(|global| &global.interface == wl_output::WlOutput::interface().name) - == Some(true) - { - state - .outputs - .retain(|_, output| output.output.data() != Some(&name)); - } - }), + wl_registry::Event::GlobalRemove { name } => { + globals.with_list(|globals| { + if globals.iter().find(|global| global.name == name).map( + |global| { + &global.interface + == wl_output::WlOutput::interface().name + }, + ) == Some(true) + { + state.outputs.retain(|_, output| { + output.output.data() != Some(&name) + }); + } + }) + } _ => {} } } @@ -400,7 +436,9 @@ impl Shm { ) } { libc::MAP_FAILED => Err(std::io::Error::last_os_error()), - ptr => Ok(unsafe { core::slice::from_raw_parts_mut(ptr as *mut _, self.size) }), + ptr => Ok(unsafe { + core::slice::from_raw_parts_mut(ptr as *mut _, self.size) + }), } } } @@ -422,7 +460,154 @@ impl Drop for Shm { // wl_output gives a transform which is applied to buffer contents. // +#[derive(clap::Parser)] +#[command(name = "grisly")] +#[command(version = "1.0")] +#[command( + about = "grim compatible screenshot tool for wlroots-based wayland compositors." +)] +struct Config { + /// Specify the region to capture. + #[arg(short = 'g', value_parser = RectParser)] + geometry: Option, + /// Use the physical size of the buffers, not the logical size of the outputs. + /// + /// Will in most cases be the same as the physical size of the monitor, ignoring any fractional scaling via xdg protocols. + #[arg(long)] + physical_size: bool, + /// Ignore any transforms (flipping, rotating) applied by the compositors. + /// + /// This means that the resulting screenshot will look as if all the monitors were all in their normal orientation. + #[arg(long)] + ignore_transforms: bool, + /// Include the cursor in screenshots. + #[arg(short = 'c')] + with_cursor: bool, + /// Specify the wayland output (monitor) to capture from. + #[arg(long, short = 'o')] + output: Option, + /// Specify the output file format. One of png, jpg, ppm, webp, farbfeld. Defaults to png. + #[arg(short = 't', value_parser = ImageTypeParser, default_value = "png")] + kind: ImageType, + /// Specify the quality of the output file from 0 to 100. Defaults to 80. + #[arg(short = 'q', default_value = "80")] + jpeg_quality: u32, + /// For PNG: specify the compression level: fast, best, default. + #[arg(short = 'l', default_value = "default", value_parser = CompressionTypeParser)] + png_level: image::codecs::png::CompressionType, +} + +#[derive(Clone, Copy, Debug, Default)] +enum ImageType { + #[default] + Png, + Jpeg, + Ppm, + WebP, + Farbfeld, +} + +#[derive(Clone, Copy, Debug)] +struct CompressionTypeParser; + +impl TypedValueParser for CompressionTypeParser { + type Value = image::codecs::png::CompressionType; + + fn parse_ref( + &self, + _cmd: &clap::Command, + _arg: Option<&clap::Arg>, + value: &std::ffi::OsStr, + ) -> Result { + let value = value + .to_str() + .ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?; + match value { + "best" => Ok(image::codecs::png::CompressionType::Best), + "fast" => Ok(image::codecs::png::CompressionType::Fast), + "default" => Ok(image::codecs::png::CompressionType::Default), + _ => Err(clap::Error::new(clap::error::ErrorKind::ValueValidation)), + } + } +} + +#[derive(Clone, Copy, Debug)] +struct ImageTypeParser; + +impl TypedValueParser for ImageTypeParser { + type Value = ImageType; + + fn parse_ref( + &self, + _cmd: &clap::Command, + _arg: Option<&clap::Arg>, + value: &std::ffi::OsStr, + ) -> Result { + let value = value + .to_str() + .ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?; + match value { + "png" => Ok(ImageType::Png), + "jpg" => Ok(ImageType::Jpeg), + "ppm" => Ok(ImageType::Ppm), + "webp" => Ok(ImageType::WebP), + "farbfeld" => Ok(ImageType::Farbfeld), + _ => Err(clap::Error::new(clap::error::ErrorKind::ValueValidation)), + } + } +} + +#[derive(Clone, Copy, Debug)] +struct RectParser; + +impl TypedValueParser for RectParser { + type Value = image::math::Rect; + + fn parse_ref( + &self, + _cmd: &clap::Command, + _arg: Option<&clap::Arg>, + value: &std::ffi::OsStr, + ) -> Result { + let value = value + .to_str() + .ok_or(clap::Error::new(clap::error::ErrorKind::InvalidUtf8))?; + let (pos, size) = value + .split_once(' ') + .ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?; + + let (x, y) = pos + .split_once(',') + .ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?; + let (width, height) = size + .split_once('x') + .ok_or(clap::Error::new(clap::error::ErrorKind::ValueValidation))?; + + let x = x.parse::().map_err(|_| { + clap::Error::new(clap::error::ErrorKind::ValueValidation) + })?; + let y = y.parse::().map_err(|_| { + clap::Error::new(clap::error::ErrorKind::ValueValidation) + })?; + let width = width.parse::().map_err(|_| { + clap::Error::new(clap::error::ErrorKind::ValueValidation) + })?; + let height = height.parse::().map_err(|_| { + clap::Error::new(clap::error::ErrorKind::ValueValidation) + })?; + + Ok(Rect { + x, + y, + width, + height, + }) + } +} + fn main() { + let config = Config::parse(); + let conn = Connection::connect_to_env().unwrap(); let (globals, mut queue) = registry_queue_init(&conn).unwrap(); @@ -438,7 +623,9 @@ fn main() { let outputs = globals.contents().with_list(|list| { list.iter() - .filter(|global| &global.interface == wl_output::WlOutput::interface().name) + .filter(|global| { + &global.interface == wl_output::WlOutput::interface().name + }) .map(|global| { let output = Output::bind( globals.registry(), @@ -448,6 +635,7 @@ fn main() { global.name, &xdg_manager, &screencopy_mgr, + config.with_cursor, ); (output.output.id(), output) @@ -462,6 +650,7 @@ fn main() { xdg_manager, screencopy_mgr, shm, + with_cursor: config.with_cursor, }; while !app @@ -473,20 +662,67 @@ fn main() { } let mut total_size = 0; + let mut union_rect = Rect { + x: 0, + y: 0, + width: 0, + height: 0, + }; let buffers = app .outputs .iter_mut() + .filter(|(_, output)| { + if let Some(output_name) = config.output.as_ref() { + output.name.as_ref() == Some(output_name) + } else { + true + } + }) .map(|(_, output)| { + let (px, py) = output.physical_position.unwrap_or((0, 0)); + let (lx, ly) = output.logical_position.unwrap_or((0, 0)); + // image destination + let (x, y) = (lx - px, ly - py); + let FrameData::Buffer(info) = core::mem::replace(&mut output.frame.data, FrameData::Done) else { unreachable!(); }; + + let (width, height) = if config.physical_size { + (info.width as i32, info.height as i32) + } else { + output + .logical_size + .unwrap_or((info.width as i32, info.height as i32)) + }; + + let scale = width as f32 / info.width as f32; + + let rect = image::math::Rect { + x: x as u32, + y: y as u32, + width: width as u32, + height: height as u32, + }; + + (rect, info, output, scale) + }) + .filter(|(rect, _, _, _)| { + if let Some(ref geom) = config.geometry { + rect_overlaps(*rect, *geom) + } else { + true + } + }) + .map(|(rect, info, output, scale)| { let size = info.height * info.stride; let offset = total_size; total_size += size; + union_rect = rect_union(union_rect, rect); - (output, info, offset) + (output, info, rect, offset, scale) }) .collect::>(); @@ -497,7 +733,7 @@ fn main() { .shm .create_pool(shm.fd.as_fd(), total_size as i32, &qh, ()); - for (output, info, offset) in buffers { + for (output, info, rect, offset, scale) in buffers { let buffer = pool.create_buffer( offset as i32, info.width as i32, @@ -513,25 +749,61 @@ fn main() { output.frame.data = FrameData::Copying { buffer, info, + rect, + scale, offset: offset as usize, }; } while !app .outputs - .iter() - .all(|(_, output)| matches!(output.frame.data, FrameData::Copied { .. })) + .values() + .all(|output| !matches!(output.frame.data, FrameData::Copying { .. })) { _ = queue.roundtrip(&mut app).unwrap(); } - fn rect_union(lhs: Rect, rhs: Rect) -> Rect { - struct AABB { - ax: u32, - ay: u32, - bx: u32, - by: u32, + struct AABB { + ax: u32, + ay: u32, + bx: u32, + by: u32, + } + + fn rect_overlaps(lhs: Rect, rhs: Rect) -> bool { + let intersection = rect_intersection(lhs, rhs); + !(intersection.width == 0 || intersection.height == 0) + } + + fn rect_intersection(lhs: Rect, rhs: Rect) -> Rect { + let lhs = AABB { + ax: lhs.x, + ay: lhs.y, + bx: lhs.x + lhs.width, + by: lhs.y + lhs.height, + }; + let rhs = AABB { + ax: rhs.x, + ay: rhs.y, + bx: rhs.x + rhs.width, + by: rhs.y + rhs.height, + }; + let intersection = AABB { + ax: lhs.ax.max(rhs.ax), + ay: lhs.ay.max(rhs.ay), + bx: lhs.bx.min(rhs.bx), + by: lhs.by.min(rhs.by), + }; + + Rect { + x: intersection.ax, + y: intersection.ay, + width: intersection.bx.saturating_sub(intersection.ax), + height: intersection.by.saturating_sub(intersection.ay), } + } + + fn rect_union(lhs: Rect, rhs: Rect) -> Rect { let lhs = AABB { ax: lhs.x, ay: lhs.y, @@ -559,77 +831,99 @@ fn main() { } } - let mut image_rect = Rect { - x: 0, - y: 0, - width: 0, - height: 0, - }; - let buffers = app - .outputs - .values_mut() - .map(|output| { - let FrameData::Copied { info, offset, .. } = - core::mem::replace(&mut output.frame.data, FrameData::Done) - else { - unreachable!(); - }; - - let (px, py) = output.physical_position.unwrap_or((0, 0)); - let (lx, ly) = output.logical_position.unwrap_or((0, 0)); - // image destination - let (x, y) = (lx - px, ly - py); - - let (width, height) = output - .logical_size - .unwrap_or((info.width as i32, info.height as i32)); - - let rect = image::math::Rect { - x: x as u32, - y: y as u32, - width: width as u32, - height: height as u32, - }; - - println!("buffer: {:?}", info); - println!(" - invert-y: {}", output.frame.inverted_y); - println!("destination rect: {:?}", rect); - println!(" - transform: {:?}", output.transform); - - image_rect = rect_union(image_rect, rect); - - ( - info, - offset, - rect, - output.frame.inverted_y, - output.transform.unwrap_or(wl_output::Transform::Normal), - ) - }) - .collect::>(); - + let image_rect = config.geometry.unwrap_or(union_rect); let mut image = image::RgbaImage::new(image_rect.width, image_rect.height); + let geom = config.geometry.unwrap_or(image_rect); + println!("image: {:?}", image_rect); + println!(" - geom: {:?}", config.geometry); let now = std::time::Instant::now(); - for (info, offset, rect, inverted_y, transform) in buffers { - let size = info.stride * info.height; - let view = BufferView { - info, - inverted_y, - transform, - bytes: &buffer[offset..][..size as usize], - }; - - if rect.width != view.width() || rect.height != view.height() { - let sampled = SampledView { - view: &view, - dimensions: (rect.width, rect.height), + app.outputs + .values_mut() + .filter_map(|output| { + match core::mem::replace(&mut output.frame.data, FrameData::Done) { + FrameData::Copied { + info, + rect, + offset, + scale, + } => Some((output, info, rect, offset, scale)), + _ => None, + } + }) + .for_each(|(output, info, rect, offset, scale)| { + let transform = if config.ignore_transforms { + wl_output::Transform::Normal + } else { + output.transform.unwrap_or(wl_output::Transform::Normal) }; - image.copy_from(&sampled, rect.x, rect.y).unwrap(); - } else { - image.copy_from(&view, rect.x, rect.y).unwrap(); - } - } + + let size = info.stride * info.height; + + // get intersection of selected geometry and this output. + // output rect `rect` is already in 'global space' + // view() needs buffer-local space coordinates + // which are the diff between rect and section_rect + + let section_rect = rect_intersection(geom, rect); + let section = Rect { + x: section_rect.x - rect.x, + y: section_rect.y - rect.y, + width: section_rect.width, + height: section_rect.height, + }; + + let view = BufferView { + info, + inverted_y: output.frame.inverted_y, + transform, + crop: section, + bytes: &buffer[offset..][..size as usize], + }; + + // println!("buffer: {:?}", info); + // println!(" - invert-y: {}", output.frame.inverted_y); + // println!("destination rect: {:?}", rect); + // println!(" - section_rect: {:?}", section_rect); + // println!(" - section: {:?}", section); + // println!(" - transform: {:?}", output.transform); + // println!("view: {:?}", view.crop); + + if scale != 1.0 { + let view = ScaledView { + view: &BufferView { + info, + inverted_y: output.frame.inverted_y, + transform, + crop: Rect { + x: 0, + y: 0, + width: info.width, + height: info.height, + }, + bytes: &buffer[offset..][..size as usize], + }, + crop: section, + dimensions: (rect.width, rect.height), + }; + + image + .copy_from( + &view, + section_rect.x - image_rect.x, + section_rect.y - image_rect.y, + ) + .unwrap(); + } else { + image + .copy_from( + &view, + section_rect.x - image_rect.x, + section_rect.y - image_rect.y, + ) + .unwrap(); + } + }); println!("copying: {}s", now.elapsed().as_secs_f64()); let file = std::fs::File::create("out.png").unwrap(); @@ -645,21 +939,22 @@ fn main() { println!("encoding: {}s", now.elapsed().as_secs_f64()); } -struct SampledView<'a, V: GenericImageView> { +struct ScaledView<'a, V: GenericImageView> { view: &'a V, + crop: Rect, dimensions: (u32, u32), } -impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> { +impl<'a, V: GenericImageView> GenericImageView for ScaledView<'a, V> { type Pixel = V::Pixel; fn dimensions(&self) -> (u32, u32) { - self.dimensions + (self.crop.width, self.crop.height) } fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel { - let u = x as f32 / self.dimensions.0 as f32; - let v = y as f32 / self.dimensions.1 as f32; + let u = (x + self.crop.x) as f32 / self.dimensions.0 as f32; + let v = (y + self.crop.y) as f32 / self.dimensions.1 as f32; image::imageops::sample_bilinear(self.view, u, v).unwrap() } @@ -668,6 +963,7 @@ impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> { struct BufferView<'a> { info: BufferInfo, inverted_y: bool, + crop: Rect, transform: wl_output::Transform, bytes: &'a [u8], } @@ -678,6 +974,8 @@ impl BufferView<'_> { let width = width - 1; let height = height - 1; + let (x, y) = (x + self.crop.x, y + self.crop.y); + let (x, y) = match self.transform { Transform::Normal => (x, y), Transform::_90 => (y, x), @@ -701,12 +999,14 @@ impl BufferView<'_> { } fn transformed_dimensions(&self) -> (u32, u32) { match self.transform { - Transform::Flipped180 | Transform::_180 | Transform::Flipped | Transform::Normal => { - (self.info.width, self.info.height) - } - Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => { - (self.info.height, self.info.width) - } + Transform::Flipped180 + | Transform::_180 + | Transform::Flipped + | Transform::Normal => (self.crop.width, self.crop.height), + Transform::_90 + | Transform::_270 + | Transform::Flipped90 + | Transform::Flipped270 => (self.crop.height, self.crop.width), _ => unreachable!(), } } @@ -714,7 +1014,8 @@ impl BufferView<'_> { // apply transform let (x, y) = self.transformed_pixel_position(x, y); - self.info.stride as usize * y as usize + (x as usize * self.pixel_stride() as usize) + self.info.stride as usize * y as usize + + (x as usize * self.pixel_stride() as usize) } fn pixel_stride(&self) -> u32 {