From 03d33bcbed55842fb86e4200130f3320f4d9402d Mon Sep 17 00:00:00 2001 From: ShootingStarDragons Date: Fri, 18 Aug 2023 11:22:58 +0800 Subject: [PATCH] chore: a little tidy up --- libwayshot/src/lib.rs | 169 ++++++++---------------------------------- 1 file changed, 31 insertions(+), 138 deletions(-) diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 0ae87100..73b8c8c2 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -20,6 +20,7 @@ use std::{ use image::{imageops::overlay, DynamicImage, RgbaImage}; use memmap2::MmapMut; +use screencopy::FrameFormat; use wayland_client::{ globals::{registry_queue_init, GlobalList}, protocol::{ @@ -168,10 +169,10 @@ impl WayshotConnection { &self, cursor_overlay: i32, output: &WlOutput, - _transform: Transform, fd: T, + file: Option<&File>, capture_region: Option, - ) -> Result<()> { + ) -> Result { // Connecting to wayland environment. let mut state = CaptureFrameState { formats: Vec::new(), @@ -251,7 +252,9 @@ impl WayshotConnection { // Bytes of data in the frame = stride * height. let frame_bytes = frame_format.stride * frame_format.height; - + if let Some(file) = file { + file.set_len(frame_bytes as u64)?; + } // Create an in memory file and return it's file descriptor. // Instantiate shm global. @@ -269,7 +272,6 @@ impl WayshotConnection { // Copy the pixel data advertised by the compositor into the buffer we just created. frame.copy(&buffer); - // On copy the Ready / Failed events are fired by the frame object, so here we check for them. loop { // Basically reads, if frame state is not None then... @@ -280,11 +282,9 @@ impl WayshotConnection { return Err(Error::FramecopyFailed); } FrameState::Finished => { - // Create a writeable memory map backed by a mem_file. - buffer.destroy(); shm_pool.destroy(); - return Ok(()); + return Ok(frame_format); } } } @@ -292,6 +292,7 @@ impl WayshotConnection { event_queue.blocking_dispatch(&mut state)?; } } + /// Get a FrameCopy instance with screenshot pixel data for any wl_output object. fn capture_output_frame( &self, @@ -300,142 +301,34 @@ impl WayshotConnection { transform: Transform, capture_region: Option, ) -> Result { + // Create an in memory file and return it's file descriptor. let fd = create_shm_fd()?; - let mut state = CaptureFrameState { - formats: Vec::new(), - state: None, - buffer_done: AtomicBool::new(false), - }; - let mut event_queue = self.conn.new_event_queue::(); - let qh = event_queue.handle(); + // Create a writeable memory map backed by a mem_file. + let mem_file = unsafe { File::from_raw_fd(fd.as_raw_fd()) }; - // Instantiating screencopy manager. - let screencopy_manager = match self.globals.bind::( - &qh, - 3..=3, - (), - ) { - Ok(x) => x, - Err(e) => { - log::error!("Failed to create screencopy manager. Does your compositor implement ZwlrScreencopy?"); - log::error!("err: {e}"); - return Err(Error::ProtocolNotFound( - "ZwlrScreencopy Manager not found".to_string(), - )); - } - }; + let frame_format = self.capture_output_frame_shm_fd( + cursor_overlay, + output, + fd, + Some(&mem_file), + capture_region, + )?; - // 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, - (), - ) + let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? }; + let data = &mut *frame_mmap; + let frame_color_type = if let Some(converter) = create_converter(frame_format.format) { + converter.convert_inplace(data) } else { - 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. - while !state.buffer_done.load(Ordering::SeqCst) { - event_queue.blocking_dispatch(&mut state)?; - } - - log::debug!( - "Received compositor frame buffer formats: {:#?}", - state.formats - ); - // Filter advertised wl_shm formats and select the first one that matches. - let frame_format = state - .formats - .iter() - .find(|frame| { - matches!( - frame.format, - wl_shm::Format::Xbgr2101010 - | wl_shm::Format::Abgr2101010 - | wl_shm::Format::Argb8888 - | wl_shm::Format::Xrgb8888 - | wl_shm::Format::Xbgr8888 - ) - }) - .copied(); - log::debug!("Selected frame buffer format: {:#?}", frame_format); - - // Check if frame format exists. - let frame_format = match frame_format { - Some(format) => format, - None => { - log::error!("No suitable frame format found"); - return Err(Error::NoSupportedBufferFormat); - } + log::error!("Unsupported buffer format: {:?}", frame_format.format); + log::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); }; - - // Bytes of data in the frame = stride * height. - let frame_bytes = frame_format.stride * frame_format.height; - - // Create an in memory file and return it's file descriptor. - let mem_file = unsafe { File::from_raw_fd(fd.as_raw_fd()) }; - mem_file.set_len(frame_bytes as u64)?; - - // Instantiate shm global. - let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); - let shm_pool = shm.create_pool(fd.as_raw_fd(), frame_bytes as i32, &qh, ()); - let buffer = shm_pool.create_buffer( - 0, - frame_format.width as i32, - frame_format.height as i32, - frame_format.stride as i32, - frame_format.format, - &qh, - (), - ); - - // Copy the pixel data advertised by the compositor into the buffer we just created. - frame.copy(&buffer); - - // On copy the Ready / Failed events are fired by the frame object, so here we check for them. - loop { - // Basically reads, if frame state is not None then... - if let Some(state) = state.state { - match state { - FrameState::Failed => { - log::error!("Frame copy failed"); - return Err(Error::FramecopyFailed); - } - FrameState::Finished => { - // Create a writeable memory map backed by a mem_file. - let mut frame_mmap = unsafe { MmapMut::map_mut(&mem_file)? }; - let data = &mut *frame_mmap; - let frame_color_type = if let Some(converter) = - create_converter(frame_format.format) - { - converter.convert_inplace(data) - } else { - log::error!("Unsupported buffer format: {:?}", frame_format.format); - log::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); - }; - buffer.destroy(); - shm_pool.destroy(); - return Ok(FrameCopy { - frame_format, - frame_color_type, - frame_mmap, - transform, - }); - } - } - } - - event_queue.blocking_dispatch(&mut state)?; - } + Ok(FrameCopy { + frame_format, + frame_color_type, + frame_mmap, + transform, + }) } fn create_frame_copy(