From 2f7b20fde798f0675c8288983de5f0dba408c862 Mon Sep 17 00:00:00 2001
From: Janis <janis@nirgendwo.xyz>
Date: Wed, 27 Nov 2024 01:23:13 +0100
Subject: [PATCH] alpha image, transforms

---
 src/main.rs | 69 ++++++++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 58 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 81c0e97..25b670b 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -617,22 +617,29 @@ fn main() {
 
             image_rect = rect_union(image_rect, rect);
 
-            (info, offset, rect, output.frame.inverted_y)
+            (
+                info,
+                offset,
+                rect,
+                output.frame.inverted_y,
+                output.transform.unwrap_or(wl_output::Transform::Normal),
+            )
         })
         .collect::<Vec<_>>();
 
-    let mut image = image::RgbImage::new(image_rect.width, image_rect.height);
+    let mut image = image::RgbaImage::new(image_rect.width, image_rect.height);
 
     let now = std::time::Instant::now();
-    for (info, offset, rect, inverted_y) in buffers {
+    for (info, offset, rect, inverted_y, transform) in buffers {
         let size = info.stride * info.height;
         let view = BufferView {
             info,
             inverted_y,
+            transform,
             bytes: &buffer[offset..][..size as usize],
         };
 
-        if rect.width != info.width || rect.height != info.height {
+        if rect.width != view.width() || rect.height != view.height() {
             let sampled = SampledView {
                 view: &view,
                 dimensions: (rect.width, rect.height),
@@ -680,17 +687,52 @@ impl<'a, V: GenericImageView> GenericImageView for SampledView<'a, V> {
 struct BufferView<'a> {
     info: BufferInfo,
     inverted_y: bool,
+    transform: wl_output::Transform,
     bytes: &'a [u8],
 }
 
 impl BufferView<'_> {
-    fn pixel_offset(&self, x: u32, y: u32) -> usize {
+    fn transformed_pixel_position(&self, x: u32, y: u32) -> (u32, u32) {
+        let (width, height) = self.transformed_dimensions();
+        let width = width - 1;
+        let height = height - 1;
+
+        let (x, y) = match self.transform {
+            Transform::Normal => (x, y),
+            Transform::_90 => (y, x),
+            Transform::_180 => (width - x, height - y),
+            Transform::_270 => (height - y, x),
+            Transform::Flipped => (width - x, y),
+            Transform::Flipped90 => (y, width - x),
+            Transform::Flipped180 => (x, height - y),
+            Transform::Flipped270 => (height - y, width - x),
+            _ => unreachable!(),
+        };
+
+        // apply buffer-local y flip
         let y = if self.inverted_y {
             self.info.height - y
         } else {
             y
         };
 
+        (x, y)
+    }
+    fn transformed_dimensions(&self) -> (u32, u32) {
+        match self.transform {
+            Transform::Flipped180 | Transform::_180 | Transform::Flipped | Transform::Normal => {
+                (self.info.width, self.info.height)
+            }
+            Transform::_90 | Transform::_270 | Transform::Flipped90 | Transform::Flipped270 => {
+                (self.info.height, self.info.width)
+            }
+            _ => unreachable!(),
+        }
+    }
+    fn pixel_offset(&self, x: u32, y: u32) -> usize {
+        // apply transform
+        let (x, y) = self.transformed_pixel_position(x, y);
+
         self.info.stride as usize * y as usize + (x as usize * self.pixel_stride() as usize)
     }
 
@@ -918,22 +960,27 @@ impl BufferView<'_> {
 }
 
 impl GenericImageView for BufferView<'_> {
-    type Pixel = image::Rgb<u8>;
+    type Pixel = image::Rgba<u8>;
 
     fn dimensions(&self) -> (u32, u32) {
-        (self.info.width, self.info.height)
+        self.transformed_dimensions()
     }
 
     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(),
-            ),
+            wl_shm::Format::Xbgr8888 => {
+                let [r, g, b] = <[u8; 3] as TryFrom<&[u8]>>::try_from(
+                    &self.bytes[self.pixel_offset(x, y)..][..3],
+                )
+                .unwrap();
+
+                Self::Pixel::from([r, g, b, 255])
+            }
             _ => Self::Pixel::from([
                 self.sample_r(x, y),
                 self.sample_g(x, y),
                 self.sample_b(x, y),
+                255,
             ]),
         }
     }