From a5cb9231ff2591957c83aa85e75632a42b86a28c Mon Sep 17 00:00:00 2001 From: ShootingStarDragons Date: Sun, 8 Oct 2023 12:49:11 +0800 Subject: [PATCH] chore: reduce repeat code --- libwayshot/src/lib.rs | 151 ++++++++++-------------------------------- 1 file changed, 35 insertions(+), 116 deletions(-) diff --git a/libwayshot/src/lib.rs b/libwayshot/src/lib.rs index 7507ae9a..2517eb22 100644 --- a/libwayshot/src/lib.rs +++ b/libwayshot/src/lib.rs @@ -26,7 +26,7 @@ use wayland_client::{ wl_output::{Transform, WlOutput}, wl_shm::{self, WlShm}, }, - Connection, + Connection, EventQueue, }; use wayland_protocols::xdg::xdg_output::zv1::client::{ zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, @@ -172,7 +172,22 @@ impl WayshotConnection { fd: T, capture_region: Option, ) -> Result { - // Connecting to wayland environment. + 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) + } + + fn capture_output_frame_get_state( + &self, + cursor_overlay: i32, + output: &WlOutput, + capture_region: Option, + ) -> Result<( + CaptureFrameState, + EventQueue, + ZwlrScreencopyFrameV1, + FrameFormat, + )> { let mut state = CaptureFrameState { formats: Vec::new(), state: None, @@ -248,12 +263,24 @@ impl WayshotConnection { return Err(Error::NoSupportedBufferFormat); } }; + Ok((state, event_queue, frame, frame_format)) + // Bytes of data in the frame = stride * height. + } + + fn capture_output_frame_inner( + &self, + mut state: CaptureFrameState, + mut event_queue: EventQueue, + frame: ZwlrScreencopyFrameV1, + frame_format: FrameFormat, + fd: T, + ) -> Result { + // Connecting to wayland environment. + let qh = event_queue.handle(); // 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. - // Instantiate shm global. let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ()); @@ -297,123 +324,15 @@ impl WayshotConnection { file: &File, capture_region: Option, ) -> Result { - let fd = file.as_fd(); - // Connecting to wayland environment. - 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(); - - // 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(), - )); - } - }; - - // 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, ()) - }; - - // 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); - } - }; + let (state, event_queue, frame, frame_format) = + self.capture_output_frame_get_state(cursor_overlay, output, capture_region)?; // Bytes of data in the frame = stride * height. let frame_bytes = frame_format.stride * frame_format.height; file.set_len(frame_bytes as u64)?; - // Create an in memory file and return it's file descriptor. - - // Instantiate shm global. - let shm = self.globals.bind::(&qh, 1..=1, ()).unwrap(); - let shm_pool = shm.create_pool(fd.as_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 => { - buffer.destroy(); - shm_pool.destroy(); - return Ok(frame_format); - } - } - } - event_queue.blocking_dispatch(&mut state)?; - } + // Create an in memory file and return it's file descriptor. + self.capture_output_frame_inner(state, event_queue, frame, frame_format, file) } /// Get a FrameCopy instance with screenshot pixel data for any wl_output object.