[feature] added 10-bit back buffer support

- Allow user to make screenshots if back buffer uses 10 bits per color channel. - Image conversion moved to convert.rs
waycrate · Sep 30, 2022 · 27cf882 · 27cf882
1 parent 72101c6
commit 27cf882
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 18 deletions.
diff --git a/src/backend.rs b/src/backend.rs
@@ -24,9 +24,7 @@ use image::{
         png::PngEncoder,
         pnm::{self, PnmEncoder},
     },
-    ColorType,
-    ColorType::Rgba8,
-    ImageEncoder,
+    ColorType, ImageEncoder,
 };
 
 use memmap2::MmapMut;
@@ -40,6 +38,8 @@ use wayland_protocols::wlr::unstable::screencopy::v1::client::{
     zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
 };
 
+use crate::convert::create_converter;
+
 /// Type of frame supported by the compositor. For now we only support Argb8888, Xrgb8888, and
 /// Xbgr8888.
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -201,7 +201,11 @@ pub fn capture_output_frame(
         .find(|frame| {
             matches!(
                 frame.format,
-                wl_shm::Format::Argb8888 | wl_shm::Format::Xrgb8888 | wl_shm::Format::Xbgr8888
+                wl_shm::Format::Xbgr2101010
+                    | wl_shm::Format::Abgr2101010
+                    | wl_shm::Format::Argb8888
+                    | wl_shm::Format::Xrgb8888
+                    | wl_shm::Format::Xbgr8888
             )
         })
         .copied();
@@ -253,20 +257,14 @@ pub fn capture_output_frame(
                     // 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 = match frame_format.format {
-                        wl_shm::Format::Argb8888 | wl_shm::Format::Xrgb8888 => {
-                            // Swap out b with r as these formats are in little endian notation.
-                            for chunk in data.chunks_exact_mut(4) {
-                                chunk.swap(0, 2);
-                            }
-                            Rgba8
-                        }
-                        wl_shm::Format::Xbgr8888 => Rgba8,
-                        unsupported_format => {
-                            log::error!("Unsupported buffer format: {:?}", unsupported_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.");
-                            exit(1);
-                        }
+                    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.");
+                        exit(1);
                     };
                     return Ok(FrameCopy {
                         frame_format,

diff --git a/src/convert.rs b/src/convert.rs
@@ -0,0 +1,75 @@
+use image::ColorType;
+use wayland_client::protocol::wl_shm;
+
+pub trait Convert {
+    /// Convert raw image data into output type, return said type
+    fn convert_inplace(&self, data: &mut [u8]) -> ColorType;
+}
+
+#[derive(Default)]
+struct ConvertBGR10 {}
+
+#[derive(Default)]
+struct ConvertNone {}
+
+#[derive(Default)]
+struct ConvertRGB8 {}
+
+const SHIFT10BITS_1: u32 = 20;
+const SHIFT10BITS_2: u32 = 10;
+
+/// Creates format converter based of input format, return None if conversion
+/// isn't possible. Conversion is happening inplace.
+pub fn create_converter(format: wl_shm::Format) -> Option<Box<dyn Convert>> {
+    match format {
+        wl_shm::Format::Xbgr8888 | wl_shm::Format::Abgr8888 => {
+            Some(Box::new(ConvertNone::default()))
+        }
+        wl_shm::Format::Xrgb8888 | wl_shm::Format::Argb8888 => {
+            Some(Box::new(ConvertRGB8::default()))
+        }
+        wl_shm::Format::Xbgr2101010 | wl_shm::Format::Abgr2101010 => {
+            Some(Box::new(ConvertBGR10::default()))
+        }
+        _ => None,
+    }
+}
+
+impl Convert for ConvertNone {
+    fn convert_inplace(&self, _data: &mut [u8]) -> ColorType {
+        ColorType::Rgba8
+    }
+}
+
+impl Convert for ConvertRGB8 {
+    fn convert_inplace(&self, data: &mut [u8]) -> ColorType {
+        for chunk in data.chunks_exact_mut(4) {
+            chunk.swap(0, 2);
+        }
+        ColorType::Rgba8
+    }
+}
+
+/// Simple conversion from 10 to 8 bits for one channel
+fn convert10_to_8(color: u32) -> u8 {
+    ((color >> 2) & 255) as u8
+}
+
+impl Convert for ConvertBGR10 {
+    fn convert_inplace(&self, data: &mut [u8]) -> ColorType {
+        for chunk in data.chunks_exact_mut(4) {
+            let pixel = ((chunk[3] as u32) << 24)
+                | ((chunk[2] as u32) << 16)
+                | ((chunk[1] as u32) << 8)
+                | chunk[0] as u32;
+            let r = convert10_to_8(pixel >> SHIFT10BITS_1);
+            let g = convert10_to_8(pixel >> SHIFT10BITS_2);
+            let b = convert10_to_8(pixel);
+            chunk[0] = b;
+            chunk[1] = g;
+            chunk[2] = r;
+            chunk[3] = 255;
+        }
+        ColorType::Rgba8
+    }
+}
diff --git a/src/wayshot.rs b/src/wayshot.rs
@@ -11,6 +11,7 @@ use wayland_client::{protocol::wl_output::WlOutput, Display};
 
 mod backend;
 mod clap;
+mod convert;
 mod output;
 
 // TODO: Create a xdg-shell surface, check for the enter event, grab the output from it.