feat: wlr-screencopy capture and valw screen

src/capture.rs

@@ -0,0 +1,218 @@
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use anyhow::{Context, Result, anyhow};
+use smithay_client_toolkit::shm::raw::RawPool;
+use wayland_client::{
+    Connection, Dispatch, Proxy, QueueHandle, WEnum,
+    protocol::{wl_buffer::WlBuffer, wl_shm::Format},
+};
+use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_frame_v1::{
+    self, ZwlrScreencopyFrameV1,
+};
+
+use crate::error::HintExt;
+use crate::frame::{self, Frame, SUPPORTED};
+use crate::wayland::{Output, State, Wayland};
+
+/// How long the compositor gets to deliver all frames.
+const TIMEOUT: Duration = Duration::from_secs(5);
+
+/// Something that captures outputs. Only wlr-screencopy exists today;
+/// ext-image-copy-capture would be a second implementation.
+pub trait Capturer {
+    fn capture(&mut self, outputs: &[Output], cursor: bool) -> Result<Vec<Frame>>;
+}
+
+/// A buffer layout the compositor offered.
+#[derive(Debug, Clone, Copy)]
+struct Layout {
+    format: Format,
+    width: u32,
+    height: u32,
+    stride: u32,
+}
+
+/// One screencopy request in flight.
+pub struct Pending {
+    output: Output,
+    offered: Vec<Format>,
+    layout: Option<Layout>,
+    // niri rejects a screencopy buffer unless its pool is exactly the
+    // buffer's size, so every capture gets a pool of its own.
+    pool: Option<RawPool>,
+    buffer: Option<WlBuffer>,
+    y_invert: bool,
+    result: Option<Result<Frame>>,
+}
+
+impl Capturer for Wayland {
+    fn capture(&mut self, outputs: &[Output], cursor: bool) -> Result<Vec<Frame>> {
+        let manager = self
+            .state
+            .screencopy
+            .clone()
+            .context("compositor does not support wlr-screencopy")
+            .hint("run `valw doctor` to see what the compositor supports")?;
+        let qh = self.queue.handle();
+
+        self.state.captures = outputs
+            .iter()
+            .map(|o| Pending {
+                output: o.clone(),
+                offered: Vec::new(),
+                layout: None,
+                pool: None,
+                buffer: None,
+                y_invert: false,
+                result: None,
+            })
+            .collect();
+        for (i, o) in outputs.iter().enumerate() {
+            manager.capture_output(cursor as i32, &o.wl, &qh, i);
+        }
+
+        let started = Instant::now();
+        self.dispatch_until(started + TIMEOUT, |s| {
+            s.captures.iter().all(|c| c.result.is_some())
+        })
+        .context("screen capture failed")?;
+        tracing::info!(
+            "captured {} output(s) in {:?} ({:?} after start)",
+            outputs.len(),
+            started.elapsed(),
+            crate::log::since_start()
+        );
+
+        std::mem::take(&mut self.state.captures)
+            .into_iter()
+            .map(|c| {
+                let name = c.output.geom.name.clone();
+                c.result
+                    .expect("checked above")
+                    .with_context(|| format!("could not capture output {name}"))
+            })
+            .collect()
+    }
+}
+
+impl Dispatch<ZwlrScreencopyFrameV1, usize> for State {
+    fn event(
+        state: &mut Self,
+        frame: &ZwlrScreencopyFrameV1,
+        event: zwlr_screencopy_frame_v1::Event,
+        &i: &usize,
+        _: &Connection,
+        qh: &QueueHandle<Self>,
+    ) {
+        use zwlr_screencopy_frame_v1::Event;
+        match event {
+            Event::Buffer {
+                format: WEnum::Value(format),
+                width,
+                height,
+                stride,
+            } => {
+                let c = &mut state.captures[i];
+                c.offered.push(format);
+                let better = c.layout.is_none_or(|l| rank(format) < rank(l.format));
+                if SUPPORTED.contains(&format) && better {
+                    c.layout = Some(Layout {
+                        format,
+                        width,
+                        height,
+                        stride,
+                    });
+                }
+                // Before v3 there is no buffer_done; the first buffer event is all we get.
+                if frame.version() < 3 {
+                    start_copy(state, frame, i, qh);
+                }
+            }
+            Event::BufferDone => start_copy(state, frame, i, qh),
+            Event::Flags { flags } => {
+                state.captures[i].y_invert = flags
+                    .into_result()
+                    .is_ok_and(|f| f.contains(zwlr_screencopy_frame_v1::Flags::YInvert));
+            }
+            Event::Ready { .. } => {
+                let c = &mut state.captures[i];
+                c.result = Some(finish(c));
+                if let Some(buffer) = c.buffer.take() {
+                    buffer.destroy();
+                }
+                frame.destroy();
+            }
+            Event::Failed => {
+                state.captures[i].result = Some(Err(anyhow!("the compositor refused the capture")));
+                frame.destroy();
+            }
+            _ => {}
+        }
+    }
+}
+
+fn rank(format: Format) -> usize {
+    SUPPORTED
+        .iter()
+        .position(|&f| f == format)
+        .unwrap_or(usize::MAX)
+}
+
+fn start_copy(state: &mut State, frame: &ZwlrScreencopyFrameV1, i: usize, qh: &QueueHandle<State>) {
+    let c = &mut state.captures[i];
+    if c.buffer.is_some() || c.result.is_some() {
+        return;
+    }
+    let Some(l) = c.layout else {
+        c.result = Some(
+            Err(anyhow!(
+                "compositor offered only {:?}, which valw can't convert yet",
+                c.offered
+            ))
+            .hint("run `valw doctor` and include its output in a bug report"),
+        );
+        frame.destroy();
+        return;
+    };
+    tracing::info!(
+        "{}: offered {:?}, using {:?} {}x{}",
+        c.output.geom.name,
+        c.offered,
+        l.format,
+        l.width,
+        l.height
+    );
+    let len = (l.stride * l.height) as usize;
+    match RawPool::new(len, &state.shm) {
+        Ok(mut pool) => {
+            let buffer = pool.create_buffer(
+                0,
+                l.width as i32,
+                l.height as i32,
+                l.stride as i32,
+                l.format,
+                (),
+                qh,
+            );
+            frame.copy(&buffer);
+            c.pool = Some(pool);
+            c.buffer = Some(buffer);
+        }
+        Err(e) => {
+            c.result = Some(Err(e).context("could not allocate a capture buffer"));
+            frame.destroy();
+        }
+    }
+}
+
+fn finish(c: &mut Pending) -> Result<Frame> {
+    let l = c.layout.context("ready without a buffer")?;
+    let pool = c.pool.as_mut().context("ready without a buffer")?;
+    let data = &pool.mmap()[..];
+    let pixels = frame::to_bgrx(l.format, l.width, l.height, l.stride, data, c.y_invert)?;
+    Ok(Frame {
+        pixels: Arc::new(frame::apply_transform(pixels, c.output.transform)),
+        output: c.output.geom.clone(),
+    })
+}

src/main.rs

@@ -1,6 +1,7 @@
 // Modules land before the CLI uses them; Task 11 removes this.
 #![allow(dead_code)]
 
+mod capture;
 mod config;
 mod detach;
 mod doctor;
@@ -14,13 +15,21 @@ mod render;
 mod selection;
 mod wayland;
 
+use std::path::PathBuf;
 use std::process::ExitCode;
+use std::time::Duration;
 
-use anyhow::Result;
+use anyhow::{Context, Result};
 use chrono::Local;
-use clap::{Parser, Subcommand};
+use clap::{Args, Parser, Subcommand};
+use image::RgbaImage;
 
+use crate::capture::Capturer;
+use crate::config::Config;
 use crate::error::Cancelled;
+use crate::lock::Lock;
+use crate::output::Target;
+use crate::wayland::Wayland;
 
 #[derive(Parser)]
 #[command(version, about = "macOS-style screenshots for niri")]
@@ -31,10 +40,44 @@ struct Cli {
 
 #[derive(Subcommand)]
 enum Command {
+    /// Capture the focused output (Cmd+Shift+3).
+    Screen {
+        /// Capture every output, one file each.
+        #[arg(long, conflicts_with_all = ["output", "clipboard_only"])]
+        all: bool,
+        #[command(flatten)]
+        common: Common,
+    },
     /// Report what the compositor and system support.
     Doctor,
 }
 
+#[derive(Args)]
+struct Common {
+    /// Copy to the clipboard without saving a file.
+    #[arg(long, conflicts_with = "output")]
+    clipboard_only: bool,
+    /// Write to this path instead of the configured directory; "-" for stdout.
+    #[arg(short, long, value_name = "PATH")]
+    output: Option<PathBuf>,
+    /// Wait this many seconds before capturing.
+    #[arg(long, value_name = "SECS", value_parser = parse_delay)]
+    delay: Option<Duration>,
+    /// Include the mouse cursor.
+    #[arg(long)]
+    cursor: bool,
+}
+
+fn parse_delay(s: &str) -> Result<Duration, String> {
+    let secs: f64 = s.parse().map_err(|_| format!("not a number: {s}"))?;
+    Duration::try_from_secs_f64(secs)
+        .map_err(|_| format!("must be a non-negative number of seconds: {s}"))
+}
+
+enum Mode {
+    Screen { all: bool },
+}
+
 fn main() -> ExitCode {
     let cli = Cli::parse();
     let log = log::init(&log::default_dir(), Local::now());
@@ -62,5 +105,137 @@ fn main() -> ExitCode {
 fn run(cli: Cli) -> Result<()> {
     match cli.command {
         Command::Doctor => doctor::run(),
+        Command::Screen { all, common } => capture(Mode::Screen { all }, common),
+    }
+}
+
+fn capture(mode: Mode, common: Common) -> Result<()> {
+    let config = config::load(&config::default_path())?;
+    let lock = Lock::acquire(&lock::default_path())?;
+    if let Some(delay) = common.delay {
+        std::thread::sleep(delay);
+    }
+    let cursor = common.cursor || config.capture.show_cursor;
+    let mut wl = Wayland::connect()?;
+    if let Ok(version) = niri::version() {
+        tracing::info!("niri {version}");
+    }
+    let outputs = wl.outputs();
+    anyhow::ensure!(!outputs.is_empty(), "the compositor reported no outputs");
+
+    // (image, output name for the file name) and which one goes to the clipboard.
+    let (shots, clip): (Vec<(RgbaImage, Option<String>)>, usize) = match mode {
+        Mode::Screen { all } => {
+            let focused = focused_output(&outputs);
+            if all {
+                let frames = wl.capture(&outputs, cursor)?;
+                let shots = frames
+                    .into_iter()
+                    .map(|f| (f.to_rgba(f.full()), Some(f.output.name)))
+                    .collect();
+                (shots, focused)
+            } else {
+                let frame = wl.capture(&outputs[focused..=focused], cursor)?.remove(0);
+                (vec![(frame.to_rgba(frame.full()), None)], 0)
+            }
+        }
+    };
+    drop(wl);
+
+    let target = Target::from_args(common.output, common.clipboard_only);
+    deliver(shots, clip, &target, &config, lock)
+}
+
+/// Index of niri's focused output, or 0 if niri can't tell us.
+fn focused_output(outputs: &[wayland::Output]) -> usize {
+    let name = match niri::focused_output() {
+        Ok(Some(name)) => name,
+        Ok(None) => return 0,
+        Err(e) => {
+            tracing::warn!("using the first output: {e:#}");
+            return 0;
+        }
+    };
+    outputs
+        .iter()
+        .position(|o| o.geom.name == name)
+        .unwrap_or_else(|| {
+            tracing::warn!("niri's focused output {name} is not among the Wayland outputs");
+            0
+        })
+}
+
+fn deliver(
+    shots: Vec<(RgbaImage, Option<String>)>,
+    clip: usize,
+    target: &Target,
+    config: &Config,
+    lock: Lock,
+) -> Result<()> {
+    let now = Local::now();
+    let pngs = shots
+        .iter()
+        .map(|(img, _)| output::encode_png(img))
+        .collect::<Result<Vec<_>>>()?;
+
+    match target {
+        Target::Default => {
+            let dir = config.save_dir();
+            for (png, (_, name)) in pngs.iter().zip(&shots) {
+                let file = output::file_name(&config.save.filename, &now, name.as_deref())?;
+                let path = output::unique_path(&dir, &file);
+                output::write_atomic(&path, png)?;
+                println!("{}", path.display());
+            }
+        }
+        Target::Path(path) => {
+            output::write_atomic(path, &pngs[0])?;
+            println!("{}", path.display());
+        }
+        Target::Stdout => output::write_stdout(&pngs[0])?,
+        Target::ClipboardOnly => {}
+    }
+
+    if *target == Target::ClipboardOnly || config.save.copy_to_clipboard {
+        let png = pngs
+            .into_iter()
+            .nth(clip)
+            .context("no image for the clipboard")?;
+        output::copy_to_clipboard(png, lock)?;
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use clap::CommandFactory;
+
+    fn parses(args: &[&str]) -> bool {
+        Cli::try_parse_from(std::iter::once("valw").chain(args.iter().copied())).is_ok()
+    }
+
+    #[test]
+    fn cli_definition_is_valid() {
+        Cli::command().debug_assert();
+    }
+
+    #[test]
+    fn flag_conflicts() {
+        assert!(parses(&["screen", "--all"]));
+        assert!(parses(&["screen", "-o", "-"]));
+        assert!(parses(&["screen", "--clipboard-only", "--cursor"]));
+        assert!(!parses(&["screen", "--all", "-o", "-"]));
+        assert!(!parses(&["screen", "--all", "--clipboard-only"]));
+        assert!(!parses(&["screen", "--clipboard-only", "-o", "a.png"]));
+    }
+
+    #[test]
+    fn delay_values() {
+        assert_eq!(parse_delay("1.5"), Ok(Duration::from_millis(1500)));
+        assert_eq!(parse_delay("0"), Ok(Duration::ZERO));
+        assert!(parse_delay("-1").is_err());
+        assert!(parse_delay("soon").is_err());
+        assert!(parse_delay("NaN").is_err());
     }
 }

src/wayland.rs

@@ -19,7 +19,7 @@ use smithay_client_toolkit::{
 use wayland_client::{
     Connection, EventQueue, QueueHandle,
     globals::{GlobalList, registry_queue_init},
-    protocol::{wl_keyboard, wl_output, wl_pointer, wl_seat, wl_surface},
+    protocol::{wl_buffer::WlBuffer, wl_keyboard, wl_output, wl_pointer, wl_seat, wl_surface},
 };
 use wayland_protocols::wp::cursor_shape::v1::client::wp_cursor_shape_device_v1::WpCursorShapeDeviceV1;
 use wayland_protocols::wp::viewporter::client::{
@@ -27,6 +27,7 @@ use wayland_protocols::wp::viewporter::client::{
 };
 use wayland_protocols_wlr::screencopy::v1::client::zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1;
 
+use crate::capture::Pending;
 use crate::error::HintExt;
 use crate::frame::OutputGeom;
 
@@ -58,6 +59,8 @@ pub struct State {
     pub cursor_device: Option<WpCursorShapeDeviceV1>,
     pub keyboard: Option<wl_keyboard::WlKeyboard>,
     pub pointer: Option<wl_pointer::WlPointer>,
+    /// Screencopy requests in flight.
+    pub captures: Vec<Pending>,
 }
 
 impl Wayland {
@@ -83,6 +86,7 @@ impl Wayland {
             cursor_device: None,
             keyboard: None,
             pointer: None,
+            captures: Vec::new(),
         };
         // Two round trips: one for wl_output, one for the xdg-output details.
         queue
@@ -425,3 +429,4 @@ smithay_client_toolkit::delegate_dispatch2!(State);
 wayland_client::delegate_noop!(State: ZwlrScreencopyManagerV1);
 wayland_client::delegate_noop!(State: WpViewporter);
 wayland_client::delegate_noop!(State: WpViewport);
+wayland_client::delegate_noop!(State: ignore WlBuffer);