From 4077a85b84cc46437944ae5d7c323b61c29bf165 Mon Sep 17 00:00:00 2001 From: Access Date: Wed, 25 Oct 2023 21:57:10 +0100 Subject: [PATCH] Refactor: always screencopy full output, new FrameGuard and ScreenCapturer --- libwayshot/src/dispatch.rs | 20 +-- libwayshot/src/image_util.rs | 15 +- libwayshot/src/lib.rs | 295 ++++++++++++++--------------------- libwayshot/src/region.rs | 92 +++++++++++ libwayshot/src/screencopy.rs | 19 ++- wayshot/Cargo.toml | 2 + wayshot/src/wayshot.rs | 2 +- 7 files changed, 241 insertions(+), 204 deletions(-) create mode 100644 libwayshot/src/region.rs diff --git a/libwayshot/src/dispatch.rs b/libwayshot/src/dispatch.rs index 9b760c21..979b8824 100644 --- a/libwayshot/src/dispatch.rs +++ b/libwayshot/src/dispatch.rs @@ -25,6 +25,7 @@ use crate::{ screencopy::FrameFormat, }; +#[derive(Debug)] pub struct OutputCaptureState { pub outputs: Vec, } @@ -128,12 +129,10 @@ impl Dispatch for OutputCaptureState { zxdg_output_v1::Event::LogicalPosition { x, y } => { output_info.dimensions.x = x; output_info.dimensions.y = y; - tracing::debug!("Logical position event fired!"); } zxdg_output_v1::Event::LogicalSize { width, height } => { output_info.dimensions.width = width; output_info.dimensions.height = height; - tracing::debug!("Logical size event fired!"); } _ => {} }; @@ -156,6 +155,7 @@ pub struct CaptureFrameState { } impl Dispatch for CaptureFrameState { + #[tracing::instrument(skip(frame), level = "trace")] fn event( frame: &mut Self, _: &ZwlrScreencopyFrameV1, @@ -171,7 +171,6 @@ impl Dispatch for CaptureFrameState { height, stride, } => { - tracing::debug!("Received Buffer event"); if let Value(f) = format { frame.formats.push(FrameFormat { format: f, @@ -184,30 +183,27 @@ impl Dispatch for CaptureFrameState { exit(1); } } - zwlr_screencopy_frame_v1::Event::Flags { .. } => { - tracing::debug!("Received Flags event"); - } zwlr_screencopy_frame_v1::Event::Ready { .. } => { // If the frame is successfully copied, a “flags” and a “ready” events are sent. Otherwise, a “failed” event is sent. // This is useful when we call .copy on the frame object. - tracing::debug!("Received Ready event"); + tracing::trace!("Received Ready event"); frame.state.replace(FrameState::Finished); } zwlr_screencopy_frame_v1::Event::Failed => { - tracing::debug!("Received Failed event"); + tracing::trace!("Received Failed event"); frame.state.replace(FrameState::Failed); } zwlr_screencopy_frame_v1::Event::Damage { .. } => { - tracing::debug!("Received Damage event"); + tracing::trace!("Received Damage event"); } zwlr_screencopy_frame_v1::Event::LinuxDmabuf { .. } => { - tracing::debug!("Received LinuxDmaBuf event"); + tracing::trace!("Received LinuxDmaBuf event"); } zwlr_screencopy_frame_v1::Event::BufferDone => { - tracing::debug!("Received bufferdone event"); + tracing::trace!("Received bufferdone event"); frame.buffer_done.store(true, Ordering::SeqCst); } - _ => unreachable!(), + _ => {} }; } } diff --git a/libwayshot/src/image_util.rs b/libwayshot/src/image_util.rs index 1099395b..68c9619e 100644 --- a/libwayshot/src/image_util.rs +++ b/libwayshot/src/image_util.rs @@ -1,13 +1,15 @@ use image::RgbaImage; use wayland_client::protocol::wl_output::Transform; +#[tracing::instrument(skip(image))] pub(crate) fn rotate_image_buffer( image: RgbaImage, transform: Transform, width: u32, height: u32, ) -> RgbaImage { - let final_buffer = match transform { + tracing::debug!("Rotating image buffer"); + match transform { Transform::_90 => image::imageops::rotate90(&image), Transform::_180 => image::imageops::rotate180(&image), Transform::_270 => image::imageops::rotate270(&image), @@ -25,16 +27,5 @@ pub(crate) fn rotate_image_buffer( image::imageops::rotate270(&flipped_buffer) } _ => image, - }; - - if final_buffer.dimensions() == (width, height) { - return final_buffer; } - - image::imageops::resize( - &final_buffer, - width, - height, - image::imageops::FilterType::Gaussian, - ) } diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 67c29a38..ef55818b 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -8,10 +8,10 @@ mod dispatch; mod error; mod image_util; pub mod output; +mod region; mod screencopy; use std::{ - cmp, fs::File, os::fd::AsFd, process::exit, @@ -19,12 +19,15 @@ use std::{ thread, }; -use image::{imageops::overlay, RgbaImage}; +use image::{imageops::replace, Rgba, RgbaImage}; use memmap2::MmapMut; +use region::RegionCapturer; +use screencopy::FrameGuard; +use tracing::debug; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ - wl_output::{Transform, WlOutput}, + wl_output::WlOutput, wl_shm::{self, WlShm}, }, Connection, EventQueue, @@ -44,35 +47,16 @@ use crate::{ screencopy::{create_shm_fd, FrameCopy, FrameFormat}, }; -pub use crate::error::{Error, Result}; +pub use crate::{ + error::{Error, Result}, + region::CaptureRegion, +}; pub mod reexport { use wayland_client::protocol::wl_output; pub use wl_output::{Transform, WlOutput}; } -type Frame = (Vec, (i32, i32)); - -/// Struct to store region capture details. -#[derive(Debug, Copy, Clone, PartialEq, Eq)] -pub struct CaptureRegion { - /// X coordinate of the area to capture. - pub x_coordinate: i32, - /// y coordinate of the area to capture. - pub y_coordinate: i32, - /// Width of the capture area. - pub width: i32, - /// Height of the capture area. - pub height: i32, -} - -#[derive(Debug)] -struct IntersectingOutput { - output: WlOutput, - region: CaptureRegion, - transform: Transform, -} - /// Struct to store wayland connection and globals list. /// # Example usage /// @@ -161,7 +145,7 @@ impl WayshotConnection { tracing::error!("Compositor did not advertise any wl_output devices!"); exit(1); } - tracing::debug!("Outputs detected: {:#?}", state.outputs); + tracing::trace!("Outputs detected: {:#?}", state.outputs); self.output_infos = state.outputs; Ok(()) @@ -174,18 +158,19 @@ impl WayshotConnection { cursor_overlay: i32, output: &WlOutput, fd: T, - capture_region: Option, - ) -> Result { + ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = - self.capture_output_frame_get_state(cursor_overlay, output, capture_region)?; - self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd) + self.capture_output_frame_get_state(cursor_overlay, output)?; + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, fd)?; + + Ok((frame_format, frame_guard)) } fn capture_output_frame_get_state( &self, cursor_overlay: i32, output: &WlOutput, - capture_region: Option, ) -> Result<( CaptureFrameState, EventQueue, @@ -216,21 +201,8 @@ impl WayshotConnection { } }; - // Capture output. - let frame: ZwlrScreencopyFrameV1 = if let Some(region) = capture_region { - screencopy_manager.capture_output_region( - cursor_overlay, - output, - region.x_coordinate, - region.y_coordinate, - region.width, - region.height, - &qh, - (), - ) - } else { - screencopy_manager.capture_output(cursor_overlay, output, &qh, ()) - }; + debug!("Capturing output..."); + let frame = screencopy_manager.capture_output(cursor_overlay, output, &qh, ()); // Empty internal event buffer until buffer_done is set to true which is when the Buffer done // event is fired, aka the capture from the compositor is succesful. @@ -238,7 +210,7 @@ impl WayshotConnection { event_queue.blocking_dispatch(&mut state)?; } - tracing::debug!( + tracing::trace!( "Received compositor frame buffer formats: {:#?}", state.formats ); @@ -257,7 +229,7 @@ impl WayshotConnection { ) }) .copied(); - tracing::debug!("Selected frame buffer format: {:#?}", frame_format); + tracing::trace!("Selected frame buffer format: {:#?}", frame_format); // Check if frame format exists. let frame_format = match frame_format { @@ -277,7 +249,7 @@ impl WayshotConnection { frame: ZwlrScreencopyFrameV1, frame_format: FrameFormat, fd: T, - ) -> Result { + ) -> Result { // Connecting to wayland environment. let qh = event_queue.handle(); @@ -309,9 +281,7 @@ impl WayshotConnection { return Err(Error::FramecopyFailed); } FrameState::Finished => { - buffer.destroy(); - shm_pool.destroy(); - return Ok(frame_format); + return Ok(FrameGuard { buffer, shm_pool }); } } } @@ -325,36 +295,36 @@ impl WayshotConnection { cursor_overlay: bool, output: &WlOutput, file: &File, - capture_region: Option, - ) -> Result { + ) -> Result<(FrameFormat, FrameGuard)> { let (state, event_queue, frame, frame_format) = - self.capture_output_frame_get_state(cursor_overlay as i32, output, capture_region)?; + self.capture_output_frame_get_state(cursor_overlay as i32, output)?; // Bytes of data in the frame = stride * height. let frame_bytes = frame_format.stride * frame_format.height; file.set_len(frame_bytes as u64)?; - self.capture_output_frame_inner(state, event_queue, frame, frame_format, file) + let frame_guard = + self.capture_output_frame_inner(state, event_queue, frame, frame_format, file)?; + + Ok((frame_format, frame_guard)) } /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. - fn capture_output_frame( + #[tracing::instrument(skip_all, fields(output = output_info.name))] + fn capture_frame_copy( &self, cursor_overlay: bool, - output: &WlOutput, - transform: Transform, - capture_region: Option, - ) -> Result { + output_info: &OutputInfo, + ) -> Result<(FrameCopy, FrameGuard)> { // Create an in memory file and return it's file descriptor. let fd = create_shm_fd()?; // Create a writeable memory map backed by a mem_file. let mem_file = File::from(fd); - let frame_format = self.capture_output_frame_shm_from_file( + let (frame_format, frame_guard) = self.capture_output_frame_shm_from_file( cursor_overlay, - output, + &output_info.wl_output, &mem_file, - capture_region, )?; let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? }; @@ -366,63 +336,35 @@ impl WayshotConnection { tracing::error!("You can send a feature request for the above format to the mailing list for wayshot over at https://sr.ht/~shinyzenith/wayshot."); return Err(Error::NoSupportedBufferFormat); }; - Ok(FrameCopy { - frame_format, - frame_color_type, - frame_mmap, - transform, - }) + Ok(( + FrameCopy { + frame_format, + frame_color_type, + frame_mmap, + transform: output_info.transform, + position: ( + output_info.dimensions.x as i64, + output_info.dimensions.y as i64, + ), + }, + frame_guard, + )) } - fn create_frame_copy( + pub fn capture_frame_copies( &self, - capture_region: CaptureRegion, + outputs: &Vec, cursor_overlay: bool, - ) -> Result { + ) -> Result> { let frame_copies = thread::scope(|scope| -> Result<_> { - let join_handles = self - .get_all_outputs() + let join_handles = outputs .into_iter() - .filter_map(|output| { - let x1: i32 = cmp::max(output.dimensions.x, capture_region.x_coordinate); - let y1: i32 = cmp::max(output.dimensions.y, capture_region.y_coordinate); - let x2: i32 = cmp::min( - output.dimensions.x + output.dimensions.width, - capture_region.x_coordinate + capture_region.width, - ); - let y2: i32 = cmp::min( - output.dimensions.y + output.dimensions.height, - capture_region.y_coordinate + capture_region.height, - ); - - let width = x2 - x1; - let height = y2 - y1; - - if width <= 0 || height <= 0 { - return None; - } - - let true_x = capture_region.x_coordinate - output.dimensions.x; - let true_y = capture_region.y_coordinate - output.dimensions.y; - let true_region = CaptureRegion { - x_coordinate: true_x, - y_coordinate: true_y, - width: capture_region.width, - height: capture_region.height, - }; - Some(IntersectingOutput { - output: output.wl_output.clone(), - region: true_region, - transform: output.transform, - }) - }) - .map(|intersecting_output| { + .map(|output_info| { scope.spawn(move || { - self.capture_output_frame( - cursor_overlay, - &intersecting_output.output, - intersecting_output.transform, - Some(intersecting_output.region), + self.capture_frame_copy(cursor_overlay, &output_info).map( + |(frame_copy, frame_guard)| { + (frame_copy, frame_guard, output_info.clone()) + }, ) }) }) @@ -435,30 +377,43 @@ impl WayshotConnection { .collect::>() })?; - Ok((frame_copies, (capture_region.width, capture_region.height))) + Ok(frame_copies) } /// Take a screenshot from the specified region. - pub fn screenshot( + fn screenshot_region_capturer( &self, - capture_region: CaptureRegion, + region_capturer: RegionCapturer, cursor_overlay: bool, ) -> Result { - let (frame_copies, (width, height)) = - self.create_frame_copy(capture_region, cursor_overlay)?; + let outputs = if let RegionCapturer::Outputs(ref outputs) = region_capturer { + outputs + } else { + &self.get_all_outputs() + }; + let frames = self.capture_frame_copies(outputs, cursor_overlay)?; + + let capture_region: CaptureRegion = match region_capturer { + RegionCapturer::Outputs(ref outputs) => outputs.try_into()?, + RegionCapturer::Region(region) => region, + }; thread::scope(|scope| { - let rotate_join_handles = frame_copies + let rotate_join_handles = frames .into_iter() - .map(|frame_copy| { + // Filter out the frames that do not contain the capture region. + .filter(|(frame_copy, _, _)| capture_region.overlaps(&frame_copy.into())) + .map(|(frame_copy, _, _)| { scope.spawn(move || { - let transform = frame_copy.transform; - let image = frame_copy.try_into()?; - Ok(image_util::rotate_image_buffer( - image, - transform, - width as u32, - height as u32, + let image = (&frame_copy).try_into()?; + Ok(( + image_util::rotate_image_buffer( + image, + frame_copy.transform, + frame_copy.frame_format.width, + frame_copy.frame_format.height, + ), + frame_copy, )) }) }) @@ -470,21 +425,27 @@ impl WayshotConnection { .flatten() .fold( None, - |possible_overlayed_image_or_error: Option>, image: Result<_>| { - if let Some(overlayed_image_or_error) = possible_overlayed_image_or_error { - if let Ok(mut overlayed_image) = overlayed_image_or_error { - if let Ok(image) = image { - overlay(&mut overlayed_image, &image, 0, 0); - Some(Ok(overlayed_image)) - } else { - Some(image) - } - } else { - Some(image) - } - } else { - Some(image) - } + |composite_image: Option>, image: Result<_>| { + // Default to a transparent image. + let composite_image = composite_image.unwrap_or_else(|| { + Ok(RgbaImage::from_pixel( + capture_region.width as u32, + capture_region.height as u32, + Rgba([0 as u8, 0 as u8, 0 as u8, 255 as u8]), + )) + }); + + Some(|| -> Result<_> { + let mut composite_image = composite_image?; + let (image, frame_copy) = image?; + replace( + &mut composite_image, + &image, + frame_copy.position.0 - capture_region.x_coordinate as i64, + frame_copy.position.1 - capture_region.y_coordinate as i64, + ); + Ok(composite_image) + }()) }, ) .ok_or_else(|| { @@ -494,19 +455,23 @@ impl WayshotConnection { }) } + /// Take a screenshot from the specified region. + pub fn screenshot( + &self, + capture_region: CaptureRegion, + cursor_overlay: bool, + ) -> Result { + self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay) + } + /// shot one ouput pub fn screenshot_single_output( &self, output_info: &OutputInfo, cursor_overlay: bool, ) -> Result { - let frame_copy = self.capture_output_frame( - cursor_overlay, - &output_info.wl_output, - output_info.transform, - None, - )?; - frame_copy.try_into() + let (frame_copy, _) = self.capture_frame_copy(cursor_overlay, output_info)?; + (&frame_copy).try_into() } /// Take a screenshot from all of the specified outputs. @@ -519,33 +484,7 @@ impl WayshotConnection { return Err(Error::NoOutputs); } - let x1 = outputs - .iter() - .map(|output| output.dimensions.x) - .min() - .unwrap(); - let y1 = outputs - .iter() - .map(|output| output.dimensions.y) - .min() - .unwrap(); - let x2 = outputs - .iter() - .map(|output| output.dimensions.x + output.dimensions.width) - .max() - .unwrap(); - let y2 = outputs - .iter() - .map(|output| output.dimensions.y + output.dimensions.height) - .max() - .unwrap(); - let capture_region = CaptureRegion { - x_coordinate: x1, - y_coordinate: y1, - width: x2 - x1, - height: y2 - y1, - }; - self.screenshot(capture_region, cursor_overlay) + self.screenshot_region_capturer(RegionCapturer::Outputs(outputs.clone()), cursor_overlay) } /// Take a screenshot from all accessible outputs. diff --git a/libwayshot/src/region.rs b/libwayshot/src/region.rs new file mode 100644 index 00000000..4a012263 --- /dev/null +++ b/libwayshot/src/region.rs @@ -0,0 +1,92 @@ +use wayland_client::protocol::wl_output::Transform; + +use crate::error::Error; +use crate::output::OutputInfo; +use crate::screencopy::FrameCopy; + +pub enum RegionCapturer { + Outputs(Vec), + Region(CaptureRegion), +} + +/// Struct to store region capture details. +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +pub struct CaptureRegion { + /// X coordinate of the area to capture. + pub x_coordinate: i32, + /// y coordinate of the area to capture. + pub y_coordinate: i32, + /// Width of the capture area. + pub width: i32, + /// Height of the capture area. + pub height: i32, +} + +impl CaptureRegion { + #[tracing::instrument(ret, level = "debug")] + pub fn overlaps(&self, other: &CaptureRegion) -> bool { + let left = self.x_coordinate; + let bottom = self.y_coordinate; + let right = self.x_coordinate + self.width; + let top = self.y_coordinate + self.height; + + let other_left = other.x_coordinate; + let other_bottom: i32 = other.y_coordinate; + let other_right = other.x_coordinate + other.width; + let other_top: i32 = other.y_coordinate + other.height; + + left < other_right && other_left < right && bottom < other_top && other_bottom < top + } +} + +impl TryFrom<&Vec> for CaptureRegion { + type Error = Error; + + fn try_from(value: &Vec) -> std::result::Result { + let x1 = value + .iter() + .map(|output| output.dimensions.x) + .min() + .unwrap(); + let y1 = value + .iter() + .map(|output| output.dimensions.y) + .min() + .unwrap(); + let x2 = value + .iter() + .map(|output| output.dimensions.x + output.dimensions.width) + .max() + .unwrap(); + let y2 = value + .iter() + .map(|output| output.dimensions.y + output.dimensions.height) + .max() + .unwrap(); + Ok(CaptureRegion { + x_coordinate: x1, + y_coordinate: y1, + width: x2 - x1, + height: y2 - y1, + }) + } +} + +impl From<&FrameCopy> for CaptureRegion { + fn from(value: &FrameCopy) -> Self { + let (width, height) = ( + value.frame_format.width as i32, + value.frame_format.height as i32, + ); + let is_portait = match value.transform { + Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => true, + _ => false, + }; + CaptureRegion { + x_coordinate: value.position.0 as i32, + y_coordinate: value.position.1 as i32, + width: if is_portait { height } else { width }, + height: if is_portait { width } else { height }, + } + } +} diff --git a/libwayshot/src/screencopy.rs b/libwayshot/src/screencopy.rs index 86cf86cd..25b7074c 100644 --- a/libwayshot/src/screencopy.rs +++ b/libwayshot/src/screencopy.rs @@ -11,10 +11,24 @@ use nix::{ sys::{memfd, mman, stat}, unistd, }; -use wayland_client::protocol::{wl_output, wl_shm::Format}; +use wayland_client::protocol::{ + wl_buffer::WlBuffer, wl_output, wl_shm::Format, wl_shm_pool::WlShmPool, +}; use crate::{Error, Result}; +pub struct FrameGuard { + pub buffer: WlBuffer, + pub shm_pool: WlShmPool, +} + +impl Drop for FrameGuard { + fn drop(&mut self) { + self.buffer.destroy(); + self.shm_pool.destroy(); + } +} + /// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and /// Xbgr8888. #[derive(Debug, Copy, Clone, PartialEq)] @@ -25,6 +39,7 @@ pub struct FrameFormat { pub stride: u32, } +#[tracing::instrument(skip(frame_mmap))] fn create_image_buffer

( frame_format: &FrameFormat, frame_mmap: &MmapMut, @@ -32,6 +47,7 @@ fn create_image_buffer

( where P: Pixel, { + tracing::debug!("Creating image buffer"); ImageBuffer::from_vec(frame_format.width, frame_format.height, frame_mmap.to_vec()) .ok_or(Error::BufferTooSmall) } @@ -44,6 +60,7 @@ pub struct FrameCopy { pub frame_color_type: ColorType, pub frame_mmap: MmapMut, pub transform: wl_output::Transform, + pub position: (i64, i64), } impl TryFrom for RgbaImage { diff --git a/wayshot/Cargo.toml b/wayshot/Cargo.toml index 4b75d05e..66d1eb41 100644 --- a/wayshot/Cargo.toml +++ b/wayshot/Cargo.toml @@ -24,6 +24,8 @@ libwayshot = { version="0.3.0", path = "../libwayshot" } image = { version = "0.24", default-features = false, features = ["jpeg", "png", "pnm", "qoi"] } dialoguer = { version = "0.11.0", features = ["fuzzy-select"] } +tracing-subscriber = "0.3.17" +tracing = "0.1.40" [[bin]] name = "wayshot" diff --git a/wayshot/src/wayshot.rs b/wayshot/src/wayshot.rs index b4b3c01f..b8ee5a6c 100644 --- a/wayshot/src/wayshot.rs +++ b/wayshot/src/wayshot.rs @@ -32,7 +32,7 @@ where fn main() -> Result<(), Box> { let args = clap::set_flags().get_matches(); let level = if args.get_flag("debug") { - Level::TRACE + Level::DEBUG } else { Level::INFO };