grim commandline option support

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"

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::{
+use wayland_protocols::{
-    zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1,
+    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<Transform>,
+    name: Option<String>,
     scale: Option<i32>,
     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<ObjectId, Output>,
@@ -170,11 +191,19 @@ impl Dispatch<ZwlrScreencopyFrameV1, ObjectId> 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<wl_output::WlOutput, u32> 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<wl_registry::WlRegistry, GlobalListContents> 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| {
+            wl_registry::Event::GlobalRemove { name } => {
-                if globals
+                globals.with_list(|globals| {
-                    .iter()
+                    if globals.iter().find(|global| global.name == name).map(
-                    .find(|global| global.name == name)
+                        |global| {
-                    .map(|global| &global.interface == wl_output::WlOutput::interface().name)
+                            &global.interface
-                    == Some(true)
+                                == wl_output::WlOutput::interface().name
-                {
+                        },
-                    state
+                    ) == Some(true)
-                        .outputs
+                    {
-                        .retain(|_, output| output.output.data() != Some(&name));
+                        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<Rect>,
+    /// 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<String>,
+    /// 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<Self::Value, clap::Error> {
+        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<Self::Value, clap::Error> {
+        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<Self::Value, clap::Error> {
+        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::<u32>().map_err(|_| {
+            clap::Error::new(clap::error::ErrorKind::ValueValidation)
+        })?;
+        let y = y.parse::<u32>().map_err(|_| {
+            clap::Error::new(clap::error::ErrorKind::ValueValidation)
+        })?;
+        let width = width.parse::<u32>().map_err(|_| {
+            clap::Error::new(clap::error::ErrorKind::ValueValidation)
+        })?;
+        let height = height.parse::<u32>().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::<Vec<_>>();
 
@@ -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()
+        .values()
-        .all(|(_, output)| matches!(output.frame.data, FrameData::Copied { .. }))
+        .all(|output| !matches!(output.frame.data, FrameData::Copying { .. }))
     {
         _ = queue.roundtrip(&mut app).unwrap();
     }
 
-    fn rect_union(lhs: Rect, rhs: Rect) -> Rect {
+    struct AABB {
-        struct AABB {
+        ax: u32,
-            ax: u32,
+        ay: u32,
-            ay: u32,
+        bx: u32,
-            bx: u32,
+        by: 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 {
+    let image_rect = config.geometry.unwrap_or(union_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::<Vec<_>>();
-
     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 {
+    app.outputs
-        let size = info.stride * info.height;
+        .values_mut()
-        let view = BufferView {
+        .filter_map(|output| {
-            info,
+            match core::mem::replace(&mut output.frame.data, FrameData::Done) {
-            inverted_y,
+                FrameData::Copied {
-            transform,
+                    info,
-            bytes: &buffer[offset..][..size as usize],
+                    rect,
-        };
+                    offset,
-
+                    scale,
-        if rect.width != view.width() || rect.height != view.height() {
+                } => Some((output, info, rect, offset, scale)),
-            let sampled = SampledView {
+                _ => None,
-                view: &view,
+            }
-                dimensions: (rect.width, rect.height),
+        })
+        .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 {
+            let size = info.stride * info.height;
-            image.copy_from(&view, rect.x, rect.y).unwrap();
+
-        }
+            // 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 u = (x + self.crop.x) as f32 / self.dimensions.0 as f32;
-        let v = y as f32 / self.dimensions.1 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 => {
+            Transform::Flipped180
-                (self.info.width, self.info.height)
+            | Transform::_180
-            }
+            | Transform::Flipped
-            Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => {
+            | Transform::Normal => (self.crop.width, self.crop.height),
-                (self.info.height, self.info.width)
+            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 {