Skip to content

Commit

Permalink
chore: a little tidy up
Browse files Browse the repository at this point in the history
  • Loading branch information
@Decodetalkers
Decodetalkers committed Aug 18, 2023
1 parent 45abae3 commit 03d33bc
Showing 1 changed file with 31 additions and 138 deletions.
169 changes: 31 additions & 138 deletions libwayshot/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -168,10 +169,10 @@ impl WayshotConnection {
&self,
cursor_overlay: i32,
output: &WlOutput,
_transform: Transform,
fd: T,
file: Option<&File>,
capture_region: Option<CaptureRegion>,
) -> Result<()> {
) -> Result<FrameFormat> {
// Connecting to wayland environment.
let mut state = CaptureFrameState {
formats: Vec::new(),
Expand Down Expand Up @@ -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.
Expand All @@ -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...
Expand All @@ -280,18 +282,17 @@ 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);
}
}
}

event_queue.blocking_dispatch(&mut state)?;
}
}

/// Get a FrameCopy instance with screenshot pixel data for any wl_output object.
fn capture_output_frame(
&self,
Expand All @@ -300,142 +301,34 @@ impl WayshotConnection {
transform: Transform,
capture_region: Option<CaptureRegion>,
) -> Result<FrameCopy> {
// 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::<CaptureFrameState>();
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::<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(),
));
}
};
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::<WlShm, _, _>(&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(
Expand Down

0 comments on commit 03d33bc

Please sign in to comment.