commit b647d95713ed118cc14a52bd5ad64cdb01dc22fa Author: Janis Date: Tue Nov 26 20:35:38 2024 +0100 initial commit diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..6fd230a --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,18 @@ +[package] +name = "grisly" +version = "0.1.0" +edition = "2021" + +[dependencies] +anyhow = "1.0.93" +bytemuck = "1.20.0" +clap = "4.5.21" +env_logger = "0.11.5" +image = "0.25.5" +libc = "0.2.165" +log = "0.4.22" +rand = "0.8.5" +thiserror = "2.0.3" +wayland-client = "0.31.7" +wayland-protocols = { version = "0.32.5", features = ["client", "unstable"] } +wayland-protocols-wlr = {version = "0.3.5", features = ["client"]} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..81c0e97 --- /dev/null +++ b/src/main.rs @@ -0,0 +1,940 @@ +use std::{ + collections::HashMap, + ffi::CString, + io::BufWriter, + os::fd::{AsFd, AsRawFd, FromRawFd}, +}; + +use image::{math::Rect, GenericImage, GenericImageView}; +use rand::Rng; +use wayland_client::{ + backend::ObjectId, + globals::{registry_queue_init, GlobalListContents}, + protocol::{ + wl_buffer::WlBuffer, + wl_output::{self, Mode, Transform, WlOutput}, + wl_registry::{self, WlRegistry}, + wl_shm::{self, WlShm}, + wl_shm_pool::WlShmPool, + }, + Connection, Dispatch, Proxy, QueueHandle, WEnum, +}; + +use wayland_protocols::xdg::xdg_output::zv1::client::{ + zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1::ZxdgOutputV1, +}; +use wayland_protocols_wlr::screencopy::v1::client::{ + zwlr_screencopy_frame_v1::{self, ZwlrScreencopyFrameV1}, + zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1, +}; + +#[derive(Debug, Clone, Copy)] +struct BufferInfo { + format: wl_shm::Format, + width: u32, + height: u32, + stride: u32, +} + +#[derive(Debug, PartialEq, Eq)] +enum FrameState { + Pre, + BufferDone, + Copied, +} + +#[derive(Debug)] +enum FrameData { + Pre(Vec), + Buffer(BufferInfo), + Copying { + buffer: WlBuffer, + info: BufferInfo, + offset: usize, + }, + Copied { + buffer: WlBuffer, + info: BufferInfo, + offset: usize, + }, + Done, +} + +impl FrameData { + fn select_buffer(&mut self) { + if let Self::Pre(infos) = self { + *self = Self::Buffer(infos.pop().unwrap()); + } + } +} + +#[derive(Debug)] +struct Frame { + frame: ZwlrScreencopyFrameV1, + buffer_infos: Vec, + inverted_y: bool, + data: FrameData, +} + +#[derive(Debug)] +struct Output { + output: WlOutput, + xdg_output: ZxdgOutputV1, + logical_size: Option<(i32, i32)>, + physical_size: Option<(i32, i32)>, + logical_position: Option<(i32, i32)>, + physical_position: Option<(i32, i32)>, + transform: Option, + scale: Option, + frame: Frame, +} + +impl Output { + fn bind< + U: Send + Sync + 'static, + D: Dispatch + + Dispatch + + Dispatch + + 'static, + >( + registry: &WlRegistry, + name: u32, + version: u32, + qh: &QueueHandle, + output_data: U, + xdg_manager: &ZxdgOutputManagerV1, + screencopy: &ZwlrScreencopyManagerV1, + ) -> Self { + let output = registry.bind(name, version, &qh, output_data); + let xdg_output = xdg_manager.get_xdg_output(&output, &qh, output.id()); + let frame = screencopy.capture_output(0, &output, &qh, output.id()); + Output { + output, + xdg_output, + logical_size: None, + physical_position: None, + transform: None, + scale: Some(1), + logical_position: None, + physical_size: None, + frame: Frame { + frame, + inverted_y: false, + buffer_infos: Vec::new(), + data: FrameData::Pre(vec![]), + }, + } + } +} + +#[derive(Debug)] +struct App { + screencopy_mgr: ZwlrScreencopyManagerV1, + xdg_manager: ZxdgOutputManagerV1, + outputs: HashMap, + shm: WlShm, +} + +impl Dispatch for App { + fn event( + _state: &mut Self, + _proxy: &ZwlrScreencopyManagerV1, + _event: ::Event, + _data: &(), + _conn: &Connection, + _qhandle: &wayland_client::QueueHandle, + ) { + // this does nothing + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + proxy: &ZwlrScreencopyFrameV1, + event: ::Event, + data: &ObjectId, + _conn: &Connection, + _qhandle: &wayland_client::QueueHandle, + ) { + let Some(output) = state.outputs.get_mut(data) else { + return; + }; + match event { + zwlr_screencopy_frame_v1::Event::Buffer { + format, + width, + height, + stride, + } => { + if let FrameData::Pre(ref mut infos) = output.frame.data { + infos.push(BufferInfo { + format: format.into_result().unwrap(), + width, + height, + stride, + }); + } + } + zwlr_screencopy_frame_v1::Event::Flags { flags } => { + output.frame.inverted_y = flags + .into_result() + .unwrap() + .contains(zwlr_screencopy_frame_v1::Flags::YInvert); + } + zwlr_screencopy_frame_v1::Event::Ready { .. } => { + if let FrameData::Copying { + buffer, + info, + offset, + } = core::mem::replace(&mut output.frame.data, FrameData::Done) + { + output.frame.data = FrameData::Copied { + buffer, + info, + offset, + }; + } + } + zwlr_screencopy_frame_v1::Event::Failed => { + output.frame.data = FrameData::Done; + panic!("screencopy failed.") + } + zwlr_screencopy_frame_v1::Event::BufferDone => { + // FIXME: pick best buffer format + // proxy.copy(buffer) + if let FrameData::Pre(infos) = + core::mem::replace(&mut output.frame.data, FrameData::Done) + { + let info = infos + .iter() + .max_by_key(|info| match info.format { + wl_shm::Format::Xrgb8888 => 20, + wl_shm::Format::Argb8888 => 10, + _ => 0, + }) + .unwrap() + .clone(); + output.frame.data = FrameData::Buffer(info); + } + } + _ => {} + } + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + proxy: &ZxdgOutputManagerV1, + event: ::Event, + data: &(), + conn: &Connection, + qhandle: &wayland_client::QueueHandle, + ) { + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + proxy: &ZxdgOutputV1, + event: ::Event, + data: &ObjectId, + conn: &Connection, + qhandle: &wayland_client::QueueHandle, + ) { + let Some(output) = state.outputs.get_mut(data) else { + return; + }; + match event { + wayland_protocols::xdg::xdg_output::zv1::client::zxdg_output_v1::Event::LogicalPosition { x, y } => { + output.logical_position = Some((x,y)); + }, + wayland_protocols::xdg::xdg_output::zv1::client::zxdg_output_v1::Event::LogicalSize { width, height } => { + output.logical_size = Some((width, height)); + }, + _ => {}, + } + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + proxy: &wl_output::WlOutput, + event: ::Event, + _data: &u32, + _conn: &Connection, + _qhandle: &wayland_client::QueueHandle, + ) { + let Some(output) = state.outputs.get_mut(&proxy.id()) else { + return; + }; + match event { + wl_output::Event::Geometry { + x, y, transform, .. + } => { + output.physical_position = Some((x, y)); + output.transform = Some(transform.into_result().unwrap()); + } + wl_output::Event::Mode { + flags, + width, + height, + .. + } => { + if flags == WEnum::Value(Mode::Current) { + output.physical_size = Some((width, height)); + } + } + wl_output::Event::Scale { factor } => { + output.scale = Some(factor); + } + _ => {} + } + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + proxy: &WlShm, + event: ::Event, + data: &(), + conn: &Connection, + qhandle: &wayland_client::QueueHandle, + ) { + } +} + +impl Dispatch for App { + fn event( + _state: &mut Self, + _proxy: &WlShmPool, + event: ::Event, + _data: &(), + _conn: &Connection, + _qhandle: &wayland_client::QueueHandle, + ) { + } +} + +impl Dispatch for App { + fn event( + _state: &mut Self, + _proxy: &WlBuffer, + event: ::Event, + _data: &(), + _conn: &Connection, + _qhandle: &wayland_client::QueueHandle, + ) { + } +} + +impl Dispatch for App { + fn event( + state: &mut Self, + registry: &wl_registry::WlRegistry, + event: ::Event, + globals: &GlobalListContents, + _conn: &Connection, + qh: &wayland_client::QueueHandle, + ) { + match event { + wl_registry::Event::Global { name, version, .. } => { + let output = Output::bind( + registry, + name, + version, + qh, + name, + &state.xdg_manager, + &state.screencopy_mgr, + ); + state.outputs.insert(output.output.id(), output); + } + wl_registry::Event::GlobalRemove { name } => globals.with_list(|globals| { + if globals + .iter() + .find(|global| global.name == name) + .map(|global| &global.interface == wl_output::WlOutput::interface().name) + == Some(true) + { + state + .outputs + .retain(|_, output| output.output.data() != Some(&name)); + } + }), + _ => {} + } + } +} + +struct Shm { + fd: std::os::unix::io::OwnedFd, + name: CString, + size: usize, +} + +impl Shm { + fn new(size: usize) -> std::io::Result { + let name = CString::new(format!( + "grisly-{}", + rand::thread_rng() + .sample_iter(&rand::distributions::Alphanumeric) + .take(8) + .map(char::from) + .collect::() + )) + .unwrap(); + let fd = match unsafe { + libc::shm_open( + name.as_ptr(), + libc::O_RDWR | libc::O_CREAT | libc::O_EXCL, + 0600, + ) + } { + -1 => Err(std::io::Error::last_os_error()), + fd => Ok(unsafe { std::os::unix::io::OwnedFd::from_raw_fd(fd) }), + }?; + + match unsafe { libc::ftruncate(fd.as_raw_fd(), size as i64) } { + -1 => Err(std::io::Error::last_os_error()), + _ => Ok(()), + }?; + + Ok(Self { fd, name, size }) + } + + fn map<'a>(&'a self) -> std::io::Result<&'a mut [u8]> { + match unsafe { + libc::mmap( + core::ptr::null_mut(), + self.size, + libc::PROT_READ | libc::PROT_WRITE, + libc::MAP_SHARED, + self.fd.as_raw_fd(), + 0, + ) + } { + libc::MAP_FAILED => Err(std::io::Error::last_os_error()), + ptr => Ok(unsafe { core::slice::from_raw_parts_mut(ptr as *mut _, self.size) }), + } + } +} + +impl Drop for Shm { + fn drop(&mut self) { + unsafe { + libc::shm_unlink(self.name.as_ptr()); + } + } +} + +// +// Post-Processing transforms & scalings: +// wl_output gives an integer scale to apply to buffer contents. +// - question: does this apply after the xdg_output logical size? +// - we use a sampled image to copy to a differently sized logical size anyway +// +// wl_output gives a transform which is applied to buffer contents. +// + +fn main() { + let conn = Connection::connect_to_env().unwrap(); + + let (globals, mut queue) = registry_queue_init(&conn).unwrap(); + + let qh = queue.handle(); + + let screencopy_mgr = globals + .bind::(&queue.handle(), 1..=3, ()) + .unwrap(); + let xdg_manager = globals + .bind::(&queue.handle(), 3..=3, ()) + .unwrap(); + + let outputs = globals.contents().with_list(|list| { + list.iter() + .filter(|global| &global.interface == wl_output::WlOutput::interface().name) + .map(|global| { + let output = Output::bind( + globals.registry(), + global.name, + global.version, + &qh, + global.name, + &xdg_manager, + &screencopy_mgr, + ); + + (output.output.id(), output) + }) + .collect::>() + }); + + let shm = globals.bind(&queue.handle(), 0..=2, ()).unwrap(); + + let mut app = App { + outputs, + xdg_manager, + screencopy_mgr, + shm, + }; + + while !app + .outputs + .iter() + .all(|(_, output)| matches!(output.frame.data, FrameData::Buffer(_))) + { + _ = queue.roundtrip(&mut app).unwrap(); + } + + let mut total_size = 0; + let buffers = app + .outputs + .iter_mut() + .map(|(_, output)| { + let FrameData::Buffer(info) = + core::mem::replace(&mut output.frame.data, FrameData::Done) + else { + unreachable!(); + }; + let size = info.height * info.stride; + let offset = total_size; + total_size += size; + + (output, info, offset) + }) + .collect::>(); + + let shm = Shm::new(total_size as usize).unwrap(); + let buffer = shm.map().unwrap(); + + let pool = app + .shm + .create_pool(shm.fd.as_fd(), total_size as i32, &qh, ()); + + for (output, info, offset) in buffers { + let buffer = pool.create_buffer( + offset as i32, + info.width as i32, + info.height as i32, + info.stride as i32, + info.format, + &qh, + (), + ); + + output.frame.frame.copy(&buffer); + + output.frame.data = FrameData::Copying { + buffer, + info, + offset: offset as usize, + }; + } + + while !app + .outputs + .iter() + .all(|(_, output)| matches!(output.frame.data, FrameData::Copied { .. })) + { + _ = queue.roundtrip(&mut app).unwrap(); + } + + fn rect_union(lhs: Rect, rhs: Rect) -> Rect { + struct AABB { + ax: u32, + ay: u32, + bx: u32, + by: u32, + } + let lhs = AABB { + ax: lhs.x, + ay: lhs.y, + bx: lhs.x + lhs.width, + by: lhs.y + lhs.height, + }; + let rhs = AABB { + ax: rhs.x, + ay: rhs.y, + bx: rhs.x + rhs.width, + by: rhs.y + rhs.height, + }; + let union = AABB { + ax: lhs.ax.min(rhs.ax), + ay: lhs.ay.min(rhs.ay), + bx: lhs.bx.max(rhs.bx), + by: lhs.by.max(rhs.by), + }; + + Rect { + x: union.ax, + y: union.ay, + width: union.bx - union.ax, + height: union.by - union.ay, + } + } + + let mut image_rect = Rect { + x: 0, + y: 0, + width: 0, + height: 0, + }; + let buffers = app + .outputs + .values_mut() + .map(|output| { + let FrameData::Copied { info, offset, .. } = + core::mem::replace(&mut output.frame.data, FrameData::Done) + else { + unreachable!(); + }; + + let (px, py) = output.physical_position.unwrap_or((0, 0)); + let (lx, ly) = output.logical_position.unwrap_or((0, 0)); + // image destination + let (x, y) = (lx - px, ly - py); + + let (width, height) = output + .logical_size + .unwrap_or((info.width as i32, info.height as i32)); + + let rect = image::math::Rect { + x: x as u32, + y: y as u32, + width: width as u32, + height: height as u32, + }; + + println!("buffer: {:?}", info); + println!(" - invert-y: {}", output.frame.inverted_y); + println!("destination rect: {:?}", rect); + println!(" - transform: {:?}", output.transform); + + image_rect = rect_union(image_rect, rect); + + (info, offset, rect, output.frame.inverted_y) + }) + .collect::>(); + + let mut image = image::RgbImage::new(image_rect.width, image_rect.height); + + let now = std::time::Instant::now(); + for (info, offset, rect, inverted_y) in buffers { + let size = info.stride * info.height; + let view = BufferView { + info, + inverted_y, + bytes: &buffer[offset..][..size as usize], + }; + + if rect.width != info.width || rect.height != info.height { + let sampled = SampledView { + view: &view, + dimensions: (rect.width, rect.height), + }; + image.copy_from(&sampled, rect.x, rect.y).unwrap(); + } else { + image.copy_from(&view, rect.x, rect.y).unwrap(); + } + } + println!("copying: {}s", now.elapsed().as_secs_f64()); + + let file = std::fs::File::create("out.png").unwrap(); + let writer = BufWriter::new(file); + let now = std::time::Instant::now(); + image + .write_with_encoder(image::codecs::png::PngEncoder::new_with_quality( + writer, + image::codecs::png::CompressionType::Default, + image::codecs::png::FilterType::Adaptive, + )) + .unwrap(); + println!("encoding: {}s", now.elapsed().as_secs_f64()); +} + +struct SampledView<'a, V: GenericImageView> { + view: &'a V, + dimensions: (u32, u32), +} + +impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> { + type Pixel = V::Pixel; + + fn dimensions(&self) -> (u32, u32) { + self.dimensions + } + + fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel { + let u = x as f32 / self.dimensions.0 as f32; + let v = y as f32 / self.dimensions.1 as f32; + + image::imageops::sample_bilinear(self.view, u, v).unwrap() + } +} + +struct BufferView<'a> { + info: BufferInfo, + inverted_y: bool, + bytes: &'a [u8], +} + +impl BufferView<'_> { + fn pixel_offset(&self, x: u32, y: u32) -> usize { + let y = if self.inverted_y { + self.info.height - y + } else { + y + }; + + self.info.stride as usize * y as usize + (x as usize * self.pixel_stride() as usize) + } + + fn pixel_stride(&self) -> u32 { + (match self.info.format { + wl_shm::Format::Argb8888 => 32, + wl_shm::Format::Xrgb8888 => 32, + wl_shm::Format::C8 => 8, + wl_shm::Format::Rgb332 => 8, + wl_shm::Format::Bgr233 => 8, + wl_shm::Format::Xrgb4444 => 16, + wl_shm::Format::Xbgr4444 => 16, + wl_shm::Format::Rgbx4444 => 16, + wl_shm::Format::Bgrx4444 => 16, + wl_shm::Format::Argb4444 => 16, + wl_shm::Format::Abgr4444 => 16, + wl_shm::Format::Rgba4444 => 16, + wl_shm::Format::Bgra4444 => 16, + wl_shm::Format::Xrgb1555 => 16, + wl_shm::Format::Xbgr1555 => 16, + wl_shm::Format::Rgbx5551 => 16, + wl_shm::Format::Bgrx5551 => 16, + wl_shm::Format::Argb1555 => 16, + wl_shm::Format::Abgr1555 => 16, + wl_shm::Format::Rgba5551 => 16, + wl_shm::Format::Bgra5551 => 16, + wl_shm::Format::Rgb565 => 16, + wl_shm::Format::Bgr565 => 16, + wl_shm::Format::Rgb888 => 24, + wl_shm::Format::Bgr888 => 24, + wl_shm::Format::Xbgr8888 => 32, + wl_shm::Format::Rgbx8888 => 32, + wl_shm::Format::Bgrx8888 => 32, + wl_shm::Format::Abgr8888 => 32, + wl_shm::Format::Rgba8888 => 32, + wl_shm::Format::Bgra8888 => 32, + wl_shm::Format::Xrgb2101010 => 32, + wl_shm::Format::Xbgr2101010 => 32, + wl_shm::Format::Rgbx1010102 => 32, + wl_shm::Format::Bgrx1010102 => 32, + wl_shm::Format::Argb2101010 => 32, + wl_shm::Format::Abgr2101010 => 32, + wl_shm::Format::Rgba1010102 => 32, + wl_shm::Format::Bgra1010102 => 32, + wl_shm::Format::R8 => 8, + wl_shm::Format::R16 => 16, + wl_shm::Format::Rg88 => 16, + wl_shm::Format::Gr88 => 16, + wl_shm::Format::Rg1616 => 32, + wl_shm::Format::Gr1616 => 32, + wl_shm::Format::Xrgb16161616f => 64, + wl_shm::Format::Xbgr16161616f => 64, + wl_shm::Format::Argb16161616f => 64, + wl_shm::Format::Abgr16161616f => 64, + wl_shm::Format::Xrgb8888A8 => 40, + wl_shm::Format::Xbgr8888A8 => 40, + wl_shm::Format::Rgbx8888A8 => 40, + wl_shm::Format::Bgrx8888A8 => 40, + wl_shm::Format::Rgb888A8 => 32, + wl_shm::Format::Bgr888A8 => 32, + wl_shm::Format::Rgb565A8 => 24, + wl_shm::Format::Bgr565A8 => 24, + wl_shm::Format::Axbxgxrx106106106106 => 64, + wl_shm::Format::Xrgb16161616 => 64, + wl_shm::Format::Xbgr16161616 => 64, + wl_shm::Format::Argb16161616 => 64, + wl_shm::Format::Abgr16161616 => 64, + _ => 1, + }) / 8 + } + + fn sample_r(&self, x: u32, y: u32) -> u8 { + let offset = self.pixel_offset(x, y); + let (shift, width) = match self.info.format { + // (shift, width) + wl_shm::Format::Argb8888 => (8 * 3, u8::BITS), + wl_shm::Format::Xrgb8888 => (8 * 3, u8::BITS), + wl_shm::Format::Xbgr8888 => (8 * 0, u8::BITS), + wl_shm::Format::Abgr8888 => (8 * 0, u8::BITS), + wl_shm::Format::Bgr888 => (8 * 2, u8::BITS), + wl_shm::Format::Rgb888 => (0, u8::BITS), + wl_shm::Format::Bgra8888 => (8 * 2, u8::BITS), + wl_shm::Format::Bgrx8888 => (8 * 2, u8::BITS), + wl_shm::Format::Rgba4444 => (0, u8::BITS / 2), + wl_shm::Format::Rgbx4444 => (0, u8::BITS / 2), + wl_shm::Format::Bgra4444 => (8, u8::BITS / 2), + wl_shm::Format::Bgrx4444 => (8, u8::BITS / 2), + wl_shm::Format::Argb4444 => (4, u8::BITS / 2), + wl_shm::Format::Xrgb4444 => (4, u8::BITS / 2), + wl_shm::Format::Rgba5551 => (0, 5), + wl_shm::Format::Rgbx5551 => (0, 5), + wl_shm::Format::Xrgb2101010 => (2, 10), + wl_shm::Format::Argb2101010 => (2, 10), + wl_shm::Format::Xrgb16161616 => (16, 16), + wl_shm::Format::Argb16161616 => (16, 16), + wl_shm::Format::Rgb332 => (0, 3), + wl_shm::Format::Bgr233 => (5, 3), + _ => unimplemented!(), + }; + + let offset = offset + (shift / 8) as usize; + let shift = shift % 8; + + let acc = (self.bytes[offset] >> shift) & !(!1u8 << (width - 1).min(7)); + let bits = (1..(width / 8)).fold(acc as u64, |acc, i| { + // byte chunks: + let mask = !(!1u8 << (width - 1 - i * 8).min(8)); + + let bits = self.bytes[offset + i as usize] & mask; + (acc << 8) & bits as u64 + }); + + let bits = if width > 8 { + bits >> (width - 8) + } else if width < 8 { + bits << (8 - width) + } else { + bits + }; + + bits as u8 + } + fn sample_g(&self, x: u32, y: u32) -> u8 { + let offset = self.pixel_offset(x, y); + let (shift, width) = match self.info.format { + // (shift, width) + wl_shm::Format::Argb8888 => (8 * 1, u8::BITS), + wl_shm::Format::Xrgb8888 => (8 * 1, u8::BITS), + wl_shm::Format::Xbgr8888 => (8 * 1, u8::BITS), + wl_shm::Format::Abgr8888 => (8 * 1, u8::BITS), + wl_shm::Format::Bgr888 => (8 * 1, u8::BITS), + wl_shm::Format::Rgb888 => (8, u8::BITS), + wl_shm::Format::Bgra8888 => (8 * 1, u8::BITS), + wl_shm::Format::Bgrx8888 => (8 * 1, u8::BITS), + wl_shm::Format::Rgba4444 => (4, u8::BITS / 2), + wl_shm::Format::Rgbx4444 => (4, u8::BITS / 2), + wl_shm::Format::Bgra4444 => (4, u8::BITS / 2), + wl_shm::Format::Bgrx4444 => (4, u8::BITS / 2), + wl_shm::Format::Argb4444 => (8, u8::BITS / 2), + wl_shm::Format::Xrgb4444 => (8, u8::BITS / 2), + wl_shm::Format::Rgba5551 => (5, 5), + wl_shm::Format::Rgbx5551 => (5, 5), + wl_shm::Format::Xrgb2101010 => (12, 10), + wl_shm::Format::Argb2101010 => (12, 10), + wl_shm::Format::Xrgb16161616 => (32, 16), + wl_shm::Format::Argb16161616 => (32, 16), + wl_shm::Format::Rgb332 => (3, 3), + wl_shm::Format::Bgr233 => (2, 3), + _ => unimplemented!(), + }; + + let offset = offset + (shift / 8) as usize; + let shift = shift % 8; + + let acc = (self.bytes[offset] >> shift) & !(!1u8 << (width - 1).min(7)); + let bits = (1..(width / 8)).fold(acc as u64, |acc, i| { + // byte chunks: + let mask = !(!1u8 << (width - 1 - i * 8).min(8)); + + let bits = self.bytes[offset + i as usize] & mask; + (acc << 8) & bits as u64 + }); + + let bits = if width > 8 { + bits >> (width - 8) + } else if width < 8 { + bits << (8 - width) + } else { + bits + }; + + bits as u8 + } + fn sample_b(&self, x: u32, y: u32) -> u8 { + let offset = self.pixel_offset(x, y); + let (shift, width) = match self.info.format { + // (shift, width) + wl_shm::Format::Argb8888 => (8 * 0, u8::BITS), + wl_shm::Format::Xrgb8888 => (8 * 0, u8::BITS), + wl_shm::Format::Xbgr8888 => (8 * 2, u8::BITS), + wl_shm::Format::Abgr8888 => (8 * 2, u8::BITS), + wl_shm::Format::Bgr888 => (8 * 0, u8::BITS), + wl_shm::Format::Rgb888 => (8 * 2, u8::BITS), + wl_shm::Format::Bgra8888 => (8 * 0, u8::BITS), + wl_shm::Format::Bgrx8888 => (8 * 0, u8::BITS), + wl_shm::Format::Rgba4444 => (8, u8::BITS / 2), + wl_shm::Format::Rgbx4444 => (8, u8::BITS / 2), + wl_shm::Format::Bgra4444 => (0, u8::BITS / 2), + wl_shm::Format::Bgrx4444 => (0, u8::BITS / 2), + wl_shm::Format::Argb4444 => (12, u8::BITS / 2), + wl_shm::Format::Xrgb4444 => (12, u8::BITS / 2), + wl_shm::Format::Rgba5551 => (10, 5), + wl_shm::Format::Rgbx5551 => (10, 5), + wl_shm::Format::Xrgb2101010 => (22, 10), + wl_shm::Format::Argb2101010 => (22, 10), + wl_shm::Format::Xrgb16161616 => (48, 16), + wl_shm::Format::Argb16161616 => (48, 16), + wl_shm::Format::Rgb332 => (6, 2), + wl_shm::Format::Bgr233 => (0, 2), + _ => unimplemented!(), + }; + + let offset = offset + (shift / 8) as usize; + let shift = shift % 8; + + let acc = (self.bytes[offset] >> shift) & !(!1u8 << (width - 1).min(7)); + let bits = (1..(width / 8)).fold(acc as u64, |acc, i| { + // byte chunks: + let mask = !(!1u8 << (width - 1 - i * 8).min(8)); + + let bits = self.bytes[offset + i as usize] & mask; + (acc << 8) & bits as u64 + }); + + let bits = if width > 8 { + bits >> (width - 8) + } else if width < 8 { + bits << (8 - width) + } else { + bits + }; + + bits as u8 + } +} + +impl GenericImageView for BufferView<'_> { + type Pixel = image::Rgb; + + fn dimensions(&self) -> (u32, u32) { + (self.info.width, self.info.height) + } + + fn get_pixel(&self, x: u32, y: u32) -> Self::Pixel { + match self.info.format { + wl_shm::Format::Xbgr8888 => Self::Pixel::from( + <[u8; 3] as TryFrom<&[u8]>>::try_from(&self.bytes[self.pixel_offset(x, y)..][..3]) + .unwrap(), + ), + _ => Self::Pixel::from([ + self.sample_r(x, y), + self.sample_g(x, y), + self.sample_b(x, y), + ]), + } + } +}