feat: update screencopy to version 3

1 files changed, 156 insertions, 72 deletions

diff --git a/src/protocols/screencopy.rs b/src/protocols/screencopy.rs
index fd6bda68..74c891d0 100644
--- a/src/protocols/screencopy.rs
+++ b/src/protocols/screencopy.rs
@@ -1,7 +1,11 @@
+use std::collections::HashMap;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 use std::time::UNIX_EPOCH;
 
+use smithay::backend::allocator::dmabuf::Dmabuf;
+use smithay::backend::allocator::{Buffer, Fourcc};
+use smithay::backend::renderer::damage::OutputDamageTracker;
 use smithay::output::Output;
 use smithay::reexports::wayland_protocols_wlr::screencopy::v1::server::zwlr_screencopy_frame_v1::{
     Flags, ZwlrScreencopyFrameV1,
@@ -11,17 +15,60 @@ use smithay::reexports::wayland_protocols_wlr::screencopy::v1::server::{
     zwlr_screencopy_frame_v1, zwlr_screencopy_manager_v1,
 };
 use smithay::reexports::wayland_server::protocol::wl_buffer::WlBuffer;
-use smithay::reexports::wayland_server::protocol::wl_shm;
+use smithay::reexports::wayland_server::protocol::wl_shm::Format;
 use smithay::reexports::wayland_server::{
     Client, DataInit, Dispatch, DisplayHandle, GlobalDispatch, New, Resource,
 };
-use smithay::utils::{Physical, Point, Rectangle, Size};
-use smithay::wayland::shm;
+use smithay::utils::{Physical, Point, Rectangle, Size, Transform};
+use smithay::wayland::{dmabuf, shm};
 
-// We do not support copy_with_damage() semantics yet.
-const VERSION: u32 = 1;
+const VERSION: u32 = 3;
 
-pub struct ScreencopyManagerState;
+pub struct ScreencopyQueue {
+    damage_tracker: OutputDamageTracker,
+    screencopies: Vec<Screencopy>,
+}
+
+impl Default for ScreencopyQueue {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl ScreencopyQueue {
+    pub fn new() -> Self {
+        Self {
+            damage_tracker: OutputDamageTracker::new((0, 0), 1.0, Transform::Normal),
+            screencopies: Vec::new(),
+        }
+    }
+
+    pub fn split(&mut self) -> (&mut OutputDamageTracker, Option<&Screencopy>) {
+        let ScreencopyQueue {
+            damage_tracker,
+            screencopies,
+        } = self;
+        (damage_tracker, screencopies.first())
+    }
+
+    pub fn push(&mut self, screencopy: Screencopy) {
+        self.screencopies.push(screencopy);
+    }
+
+    pub fn pop(&mut self) -> Screencopy {
+        self.screencopies.pop().unwrap()
+    }
+
+    pub fn remove_output(&mut self, output: &Output) {
+        self.screencopies
+            .retain(|screencopy| screencopy.output() != output);
+    }
+}
+
+#[derive(Default)]
+pub struct ScreencopyManagerState {
+    queues: HashMap<ZwlrScreencopyManagerV1, ScreencopyQueue>,
+}
 
 pub struct ScreencopyManagerGlobalData {
     filter: Box<dyn for<'c> Fn(&'c Client) -> bool + Send + Sync>,
@@ -42,7 +89,28 @@ impl ScreencopyManagerState {
         };
         display.create_global::<D, ZwlrScreencopyManagerV1, _>(VERSION, global_data);
 
-        Self
+        Self {
+            queues: HashMap::new(),
+        }
+    }
+
+    pub fn bind(&mut self, manager: &ZwlrScreencopyManagerV1) {
+        // Clean up all entries if its manager is dead and its queue is empty.
+        self.queues
+            .retain(|k, v| k.is_alive() || !v.screencopies.is_empty());
+
+        self.queues.insert(manager.clone(), ScreencopyQueue::new());
+    }
+
+    pub fn get_queue_mut(
+        &mut self,
+        manager: &ZwlrScreencopyManagerV1,
+    ) -> Option<&mut ScreencopyQueue> {
+        self.queues.get_mut(manager)
+    }
+
+    pub fn queues_mut(&mut self) -> impl Iterator<Item = &mut ScreencopyQueue> {
+        self.queues.values_mut()
     }
 }
 
@@ -56,14 +124,15 @@ where
     D: 'static,
 {
     fn bind(
-        _state: &mut D,
+        state: &mut D,
         _display: &DisplayHandle,
         _client: &Client,
         manager: New<ZwlrScreencopyManagerV1>,
         _manager_state: &ScreencopyManagerGlobalData,
         data_init: &mut DataInit<'_, D>,
     ) {
-        data_init.init(manager, ());
+        let manager = data_init.init(manager, ());
+        state.screencopy_state().bind(&manager);
     }
 
     fn can_view(client: Client, global_data: &ScreencopyManagerGlobalData) -> bool {
@@ -82,7 +151,7 @@ where
     fn request(
         _state: &mut D,
         _client: &Client,
-        _manager: &ZwlrScreencopyManagerV1,
+        manager: &ZwlrScreencopyManagerV1,
         request: zwlr_screencopy_manager_v1::Request,
         _data: &(),
         _display: &DisplayHandle,
@@ -174,6 +243,7 @@ where
         let frame = data_init.init(
             frame,
             ScreencopyFrameState::Pending {
+                manager: manager.clone(),
                 info,
                 copied: Arc::new(AtomicBool::new(false)),
             },
@@ -181,30 +251,31 @@ where
 
         // Send desired SHM buffer parameters.
         frame.buffer(
-            wl_shm::Format::Argb8888,
+            Format::Xrgb8888,
             buffer_size.w as u32,
             buffer_size.h as u32,
             buffer_size.w as u32 * 4,
         );
 
-        // if manager.version() >= 3 {
-        //     // Send desired DMA buffer parameters.
-        //     frame.linux_dmabuf(
-        //         Fourcc::Argb8888 as u32,
-        //         buffer_size.w as u32,
-        //         buffer_size.h as u32,
-        //     );
-        //
-        //     // Notify client that all supported buffers were enumerated.
-        //     frame.buffer_done();
-        // }
+        if frame.version() >= 3 {
+            // Send desired DMA buffer parameters.
+            frame.linux_dmabuf(
+                Fourcc::Xrgb8888 as u32,
+                buffer_size.w as u32,
+                buffer_size.h as u32,
+            );
+
+            // Notify client that all supported buffers were enumerated.
+            frame.buffer_done();
+        }
     }
 }
 
 /// Handler trait for wlr-screencopy.
 pub trait ScreencopyHandler {
     /// Handle new screencopy request.
-    fn frame(&mut self, frame: Screencopy);
+    fn frame(&mut self, manager: &ZwlrScreencopyManagerV1, screencopy: Screencopy);
+    fn screencopy_state(&mut self) -> &mut ScreencopyManagerState;
 }
 
 #[allow(missing_docs)]
@@ -236,6 +307,7 @@ pub struct ScreencopyFrameInfo {
 pub enum ScreencopyFrameState {
     Failed,
     Pending {
+        manager: ZwlrScreencopyManagerV1,
         info: ScreencopyFrameInfo,
         copied: Arc<AtomicBool>,
     },
@@ -260,9 +332,13 @@ where
             return;
         }
 
-        let (info, copied) = match data {
-            ScreencopyFrameState::Failed => return,
-            ScreencopyFrameState::Pending { info, copied } => (info, copied),
+        let ScreencopyFrameState::Pending {
+            manager,
+            info,
+            copied,
+        } = data
+        else {
+            return;
         };
 
         if copied.load(Ordering::SeqCst) {
@@ -275,44 +351,71 @@ where
 
         let (buffer, with_damage) = match request {
             zwlr_screencopy_frame_v1::Request::Copy { buffer } => (buffer, false),
-            // zwlr_screencopy_frame_v1::Request::CopyWithDamage { buffer } => (buffer, true),
+            zwlr_screencopy_frame_v1::Request::CopyWithDamage { buffer } => (buffer, true),
             _ => unreachable!(),
         };
 
-        if !shm::with_buffer_contents(&buffer, |_buf, shm_len, buffer_data| {
-            buffer_data.format == wl_shm::Format::Argb8888
-                && buffer_data.stride == info.buffer_size.w * 4
-                && buffer_data.height == info.buffer_size.h
-                && shm_len as i32 == buffer_data.stride * buffer_data.height
+        let size = info.buffer_size;
+
+        let buffer = if let Ok(dmabuf) = dmabuf::get_dmabuf(&buffer) {
+            if dmabuf.format().code == Fourcc::Xrgb8888
+                && dmabuf.width() == size.w as u32
+                && dmabuf.height() == size.h as u32
+            {
+                ScreencopyBuffer::Dmabuf(dmabuf.clone())
+            } else {
+                frame.post_error(
+                    zwlr_screencopy_frame_v1::Error::InvalidBuffer,
+                    "invalid dmabuf parameters",
+                );
+                return;
+            }
+        } else if shm::with_buffer_contents(&buffer, |_, shm_len, buffer_data| {
+            buffer_data.format == Format::Xrgb8888
+                && buffer_data.width == size.w
+                && buffer_data.height == size.h
+                && buffer_data.stride == size.w * 4
+                && shm_len == buffer_data.stride as usize * buffer_data.height as usize
         })
         .unwrap_or(false)
         {
+            ScreencopyBuffer::Shm(buffer)
+        } else {
             frame.post_error(
                 zwlr_screencopy_frame_v1::Error::InvalidBuffer,
                 "invalid buffer",
             );
             return;
-        }
+        };
 
         copied.store(true, Ordering::SeqCst);
 
-        state.frame(Screencopy {
-            with_damage,
-            buffer,
-            frame: frame.clone(),
-            info: info.clone(),
-            submitted: false,
-        });
+        state.frame(
+            manager,
+            Screencopy {
+                buffer,
+                frame: frame.clone(),
+                info: info.clone(),
+                with_damage,
+                submitted: false,
+            },
+        );
     }
 }
 
+/// Screencopy buffer.
+#[derive(Clone)]
+pub enum ScreencopyBuffer {
+    Dmabuf(Dmabuf),
+    Shm(WlBuffer),
+}
+
 /// Screencopy frame.
 pub struct Screencopy {
     info: ScreencopyFrameInfo,
     frame: ZwlrScreencopyFrameV1,
-    #[allow(unused)]
+    buffer: ScreencopyBuffer,
     with_damage: bool,
-    buffer: WlBuffer,
     submitted: bool,
 }
 
@@ -326,7 +429,7 @@ impl Drop for Screencopy {
 
 impl Screencopy {
     /// Get the target buffer to copy to.
-    pub fn buffer(&self) -> &WlBuffer {
+    pub fn buffer(&self) -> &ScreencopyBuffer {
         &self.buffer
     }
 
@@ -346,14 +449,16 @@ impl Screencopy {
         self.info.overlay_cursor
     }
 
-    // pub fn damage(&mut self, damage: &[Rectangle<i32, Physical>]) {
-    //     assert!(self.with_damage);
-    //
-    //     for Rectangle { loc, size } in damage {
-    //         self.frame
-    //             .damage(loc.x as u32, loc.y as u32, size.w as u32, size.h as u32);
-    //     }
-    // }
+    pub fn with_damage(&self) -> bool {
+        self.with_damage
+    }
+
+    pub fn damage(&self, damages: &[Rectangle<i32, Physical>]) {
+        for Rectangle { loc, size } in damages {
+            self.frame
+                .damage(loc.x as u32, loc.y as u32, size.w as u32, size.h as u32);
+        }
+    }
 
     /// Submit the copied content.
     pub fn submit(mut self, y_invert: bool) {
@@ -374,25 +479,4 @@ impl Screencopy {
         // Mark frame as submitted to ensure destructor isn't run.
         self.submitted = true;
     }
-
-    // pub fn submit_after_sync<T>(
-    //     self,
-    //     y_invert: bool,
-    //     sync_point: Option<OwnedFd>,
-    //     event_loop: &LoopHandle<'_, T>,
-    // ) {
-    //     match sync_point {
-    //         None => self.submit(y_invert),
-    //         Some(sync_fd) => {
-    //             let source = Generic::new(sync_fd, Interest::READ, Mode::OneShot);
-    //             let mut screencopy = Some(self);
-    //             event_loop
-    //                 .insert_source(source, move |_, _, _| {
-    //                     screencopy.take().unwrap().submit(y_invert);
-    //                     Ok(PostAction::Remove)
-    //                 })
-    //                 .unwrap();
-    //         }
-    //     }
-    // }
 }