From 1d26367ce5ba963955139332e9e8fd1b56bae28e Mon Sep 17 00:00:00 2001 From: elwrcl Date: Tue, 29 Sep 2026 23:26:38 +0300 Subject: [PATCH] feat: screencopy frames, pixel formats and output transforms --- src/frame.rs | 260 +++++++++++++++++++++++++++++++++++++++++++++++++++ src/main.rs | 1 + 2 files changed, 261 insertions(+) create mode 100644 src/frame.rs diff --git a/src/frame.rs b/src/frame.rs new file mode 100644 index 0000000..3d2b87a --- /dev/null +++ b/src/frame.rs @@ -0,0 +1,260 @@ +use std::sync::Arc; + +use anyhow::{Result, bail, ensure}; +use image::{RgbaImage, imageops}; +use wayland_client::protocol::{wl_output::Transform, wl_shm::Format}; + +/// An output's place in the global logical coordinate space. +#[derive(Debug, Clone, PartialEq)] +pub struct OutputGeom { + pub name: String, + pub x: i32, + pub y: i32, + pub width: i32, + pub height: i32, +} + +/// A rectangle in a frame's physical pixels. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct PixelRect { + pub x: u32, + pub y: u32, + pub width: u32, + pub height: u32, +} + +/// Pixels in wl_shm XRGB8888 memory order: B, G, R, X. This is what niri +/// hands out and what the overlay shows, so nothing on the way to the first +/// overlay frame has to reorder bytes. `RgbaImage` is only the container; +/// its channels do not mean R, G, B, A here. +pub type Bgrx = RgbaImage; + +/// One captured output: upright pixels at physical resolution. +#[derive(Debug, Clone)] +pub struct Frame { + pub output: OutputGeom, + /// Shared with the overlay, which shows it as the undimmed selection. + pub pixels: Arc, +} + +impl Frame { + /// The whole frame as a rectangle. + pub fn full(&self) -> PixelRect { + PixelRect { + x: 0, + y: 0, + width: self.pixels.width(), + height: self.pixels.height(), + } + } + + /// The `rect` part as opaque RGBA, for saving. + /// + /// Alpha is always 255: a screenshot of what's on screen is opaque, even + /// where the compositor's buffer has alpha 0. + pub fn to_rgba(&self, rect: PixelRect) -> RgbaImage { + let stride = self.pixels.width() as usize * 4; + let raw = self.pixels.as_raw(); + let mut out = RgbaImage::new(rect.width, rect.height); + for (y, row) in out.rows_mut().enumerate() { + let start = (rect.y as usize + y) * stride + rect.x as usize * 4; + let (src, _) = raw[start..start + rect.width as usize * 4].as_chunks::<4>(); + for (px, s) in row.zip(src) { + px.0 = [s[2], s[1], s[0], 255]; + } + } + out + } +} + +/// Shm formats valw can read, in order of preference. +pub const SUPPORTED: [Format; 4] = [ + Format::Xrgb8888, + Format::Argb8888, + Format::Xbgr8888, + Format::Abgr8888, +]; + +/// Copies a screencopy buffer into B, G, R, X order, dropping stride padding. +pub fn to_bgrx( + format: Format, + width: u32, + height: u32, + stride: u32, + data: &[u8], + y_invert: bool, +) -> Result { + let swap = match format { + Format::Xrgb8888 | Format::Argb8888 => false, + Format::Xbgr8888 | Format::Abgr8888 => true, + other => bail!("unsupported buffer format {other:?}"), + }; + ensure!( + stride >= width * 4, + "stride {stride} is too small for width {width}" + ); + ensure!( + data.len() >= (stride * height) as usize, + "buffer is {} bytes, expected at least {}", + data.len(), + stride * height + ); + + let row_len = width as usize * 4; + let mut out = Bgrx::new(width, height); + for (y, dst) in out.chunks_exact_mut(row_len).enumerate() { + let src_y = if y_invert { height as usize - 1 - y } else { y }; + let src = &data[src_y * stride as usize..][..row_len]; + dst.copy_from_slice(src); + if swap { + for px in dst.as_chunks_mut::<4>().0 { + px.swap(0, 2); + } + } + } + Ok(out) +} + +/// Turns a buffer in the output's native orientation upright, the same way +/// grim does: rotate clockwise by the transform's angle, then mirror. +pub fn apply_transform(img: Bgrx, transform: Transform) -> Bgrx { + let rotated = match transform { + Transform::_90 | Transform::Flipped90 => imageops::rotate90(&img), + Transform::_180 | Transform::Flipped180 => imageops::rotate180(&img), + Transform::_270 | Transform::Flipped270 => imageops::rotate270(&img), + _ => img, + }; + match transform { + Transform::Flipped + | Transform::Flipped90 + | Transform::Flipped180 + | Transform::Flipped270 => imageops::flip_horizontal(&rotated), + _ => rotated, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// 2x1 buffer: a red pixel then a blue one, in the given byte order. + fn two_pixels(format: Format) -> Vec { + match format { + Format::Xrgb8888 | Format::Argb8888 => vec![0, 0, 255, 0, 255, 0, 0, 0], + _ => vec![255, 0, 0, 0, 0, 0, 255, 0], + } + } + + fn bgr(img: &Bgrx, x: u32, y: u32) -> [u8; 3] { + let p = img.get_pixel(x, y).0; + [p[0], p[1], p[2]] + } + + #[test] + fn reads_all_supported_formats_as_bgrx() { + for format in SUPPORTED { + let img = to_bgrx(format, 2, 1, 8, &two_pixels(format), false).unwrap(); + assert_eq!(bgr(&img, 0, 0), [0, 0, 255], "red, {format:?}"); + assert_eq!(bgr(&img, 1, 0), [255, 0, 0], "blue, {format:?}"); + } + } + + #[test] + fn honours_stride_padding() { + // One pixel per row, 8-byte stride. + let data = [0, 0, 255, 0, 9, 9, 9, 9, 255, 0, 0, 0, 9, 9, 9, 9]; + let img = to_bgrx(Format::Xrgb8888, 1, 2, 8, &data, false).unwrap(); + assert_eq!(img.as_raw().len(), 8); + assert_eq!(bgr(&img, 0, 0), [0, 0, 255]); + assert_eq!(bgr(&img, 0, 1), [255, 0, 0]); + } + + #[test] + fn y_invert_flips_rows() { + let data = [0, 0, 255, 0, 255, 0, 0, 0]; + let img = to_bgrx(Format::Xrgb8888, 1, 2, 4, &data, true).unwrap(); + assert_eq!(bgr(&img, 0, 0), [255, 0, 0]); + assert_eq!(bgr(&img, 0, 1), [0, 0, 255]); + } + + #[test] + fn rejects_unknown_format_by_name() { + let err = to_bgrx(Format::Xbgr2101010, 1, 1, 4, &[0; 4], false).unwrap_err(); + assert_eq!(err.to_string(), "unsupported buffer format Xbgr2101010"); + } + + #[test] + fn rejects_short_buffer() { + assert!(to_bgrx(Format::Xrgb8888, 2, 2, 8, &[0; 8], false).is_err()); + } + + fn frame(width: u32, height: u32, data: Vec) -> Frame { + Frame { + output: OutputGeom { + name: "A".into(), + x: 0, + y: 0, + width: width as i32, + height: height as i32, + }, + pixels: Arc::new(Bgrx::from_raw(width, height, data).unwrap()), + } + } + + #[test] + fn to_rgba_swaps_channels_and_makes_opaque() { + // Red with alpha 0, as an ARGB8888 buffer may deliver it. + let f = frame(1, 1, vec![0, 0, 255, 0]); + assert_eq!(f.to_rgba(f.full()).get_pixel(0, 0).0, [255, 0, 0, 255]); + } + + #[test] + fn to_rgba_crops() { + // 3x2: pixel value = its index, in the blue byte. + let data = (0u8..6).flat_map(|i| [i, 0, 0, 0]).collect(); + let f = frame(3, 2, data); + let rect = PixelRect { + x: 1, + y: 1, + width: 2, + height: 1, + }; + let out = f.to_rgba(rect); + assert_eq!(out.dimensions(), (2, 1)); + assert_eq!(out.get_pixel(0, 0).0, [0, 0, 4, 255]); + assert_eq!(out.get_pixel(1, 0).0, [0, 0, 5, 255]); + } + + fn marked() -> Bgrx { + // 3x2, top-left pixel marked. + let mut img = Bgrx::new(3, 2); + img.put_pixel(0, 0, image::Rgba([255, 255, 255, 255])); + img + } + + fn marked_at(img: &Bgrx) -> (u32, u32, u32, u32) { + let (x, y, _) = img + .enumerate_pixels() + .find(|(_, _, p)| p.0[0] == 255) + .unwrap(); + (img.width(), img.height(), x, y) + } + + #[test] + fn transforms_match_grim() { + use Transform::*; + let cases = [ + (Normal, (3, 2, 0, 0)), + (_90, (2, 3, 1, 0)), + (_180, (3, 2, 2, 1)), + (_270, (2, 3, 0, 2)), + (Flipped, (3, 2, 2, 0)), + (Flipped90, (2, 3, 0, 0)), + (Flipped180, (3, 2, 0, 1)), + (Flipped270, (2, 3, 1, 2)), + ]; + for (t, want) in cases { + assert_eq!(marked_at(&apply_transform(marked(), t)), want, "{t:?}"); + } + } +} diff --git a/src/main.rs b/src/main.rs index bc7e635..e44d1bc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ mod config; mod detach; mod error; +mod frame; mod lock; mod log;