Skip to content

Commit

Permalink
chore: reduce repeat code
Browse files Browse the repository at this point in the history
  • Loading branch information
@Decodetalkers
Decodetalkers committed Oct 8, 2023
1 parent b290beb commit a5cb923
Showing 1 changed file with 35 additions and 116 deletions.
151 changes: 35 additions & 116 deletions libwayshot/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -172,7 +172,22 @@ impl WayshotConnection {
fd: T,
capture_region: Option<CaptureRegion>,
) -> Result<FrameFormat> {
// 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<CaptureRegion>,
) -> Result<(
CaptureFrameState,
EventQueue<CaptureFrameState>,
ZwlrScreencopyFrameV1,
FrameFormat,
)> {
let mut state = CaptureFrameState {
formats: Vec::new(),
state: None,
Expand Down Expand Up @@ -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<T: AsFd>(
&self,
mut state: CaptureFrameState,
mut event_queue: EventQueue<CaptureFrameState>,
frame: ZwlrScreencopyFrameV1,
frame_format: FrameFormat,
fd: T,
) -> Result<FrameFormat> {
// 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::<WlShm, _, _>(&qh, 1..=1, ()).unwrap();
let shm_pool = shm.create_pool(fd.as_fd(), frame_bytes as i32, &qh, ());
Expand Down Expand Up @@ -297,123 +324,15 @@ impl WayshotConnection {
file: &File,
capture_region: Option<CaptureRegion>,
) -> Result<FrameFormat> {
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::<CaptureFrameState>();
let qh = event_queue.handle();

// Instantiating screencopy manager.
let screencopy_manager = match self.globals.bind::<ZwlrScreencopyManagerV1, _, _>(
&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::<WlShm, _, _>(&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.
Expand Down

0 comments on commit a5cb923

Please sign in to comment.