diff --git a/src/capture.rs b/src/capture.rs new file mode 100644 index 0000000..59b87f6 --- /dev/null +++ b/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>; +} + +/// 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, + layout: Option, + // 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, + buffer: Option, + y_invert: bool, + result: Option>, +} + +impl Capturer for Wayland { + fn capture(&mut self, outputs: &[Output], cursor: bool) -> Result> { + 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 for State { + fn event( + state: &mut Self, + frame: &ZwlrScreencopyFrameV1, + event: zwlr_screencopy_frame_v1::Event, + &i: &usize, + _: &Connection, + qh: &QueueHandle, + ) { + 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) { + 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 { + 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(), + }) +} diff --git a/src/main.rs b/src/main.rs index ee44c9d..50fc543 100644 --- a/src/main.rs +++ b/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, + /// Wait this many seconds before capturing. + #[arg(long, value_name = "SECS", value_parser = parse_delay)] + delay: Option, + /// Include the mouse cursor. + #[arg(long)] + cursor: bool, +} + +fn parse_delay(s: &str) -> Result { + 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)>, 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)>, + clip: usize, + target: &Target, + config: &Config, + lock: Lock, +) -> Result<()> { + let now = Local::now(); + let pngs = shots + .iter() + .map(|(img, _)| output::encode_png(img)) + .collect::>>()?; + + 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()); } } diff --git a/src/wayland.rs b/src/wayland.rs index 3e95a4d..1625327 100644 --- a/src/wayland.rs +++ b/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, pub keyboard: Option, pub pointer: Option, + /// Screencopy requests in flight. + pub captures: Vec, } 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);