Add freeze functionality

libwayshot/src/dispatch.rs

@@ -1,20 +1,26 @@
 use std::{
+    collections::HashSet,
     process::exit,
     sync::atomic::{AtomicBool, Ordering},
 };
 use wayland_client::{
     delegate_noop,
     globals::GlobalListContents,
     protocol::{
-        wl_buffer::WlBuffer, wl_output, wl_output::WlOutput, wl_registry, wl_registry::WlRegistry,
-        wl_shm::WlShm, wl_shm_pool::WlShmPool,
+        wl_buffer::WlBuffer, wl_compositor::WlCompositor, wl_output, wl_output::WlOutput,
+        wl_registry, wl_registry::WlRegistry, wl_shm::WlShm, wl_shm_pool::WlShmPool,
+        wl_surface::WlSurface,
     },
     Connection, Dispatch, QueueHandle, WEnum,
     WEnum::Value,
 };
 use wayland_protocols::xdg::xdg_output::zv1::client::{
     zxdg_output_manager_v1::ZxdgOutputManagerV1, zxdg_output_v1, zxdg_output_v1::ZxdgOutputV1,
 };
+use wayland_protocols_wlr::layer_shell::v1::client::{
+    zwlr_layer_shell_v1::ZwlrLayerShellV1,
+    zwlr_layer_surface_v1::{self, ZwlrLayerSurfaceV1},
+};
 use wayland_protocols_wlr::screencopy::v1::client::{
     zwlr_screencopy_frame_v1, zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1,
     zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
@@ -58,16 +64,9 @@ impl Dispatch<WlRegistry, ()> for OutputCaptureState {
                         name: "".to_string(),
                         description: String::new(),
                         transform: wl_output::Transform::Normal,
-                        dimensions: OutputPositioning {
-                            x: 0,
-                            y: 0,
-                            width: 0,
-                            height: 0,
-                        },
-                        mode: WlOutputMode {
-                            width: 0,
-                            height: 0,
-                        },
+                        scale: 1,
+                        dimensions: OutputPositioning::default(),
+                        mode: WlOutputMode::default(),
                     });
                 } else {
                     tracing::error!("Ignoring a wl_output with version < 4.");
@@ -109,7 +108,11 @@ impl Dispatch<WlOutput, ()> for OutputCaptureState {
             } => {
                 output.transform = transform;
             }
-            _ => (),
+            wl_output::Event::Scale { factor } => {
+                output.scale = factor;
+            }
+            wl_output::Event::Done => {}
+            _ => {}
         }
     }
 }
@@ -227,3 +230,43 @@ impl wayland_client::Dispatch<wl_registry::WlRegistry, GlobalListContents> for W
     ) {
     }
 }
+
+pub struct LayerShellState {
+    pub configured_outputs: HashSet<WlOutput>,
+}
+
+delegate_noop!(LayerShellState: ignore WlCompositor);
+delegate_noop!(LayerShellState: ignore WlShm);
+delegate_noop!(LayerShellState: ignore WlShmPool);
+delegate_noop!(LayerShellState: ignore WlBuffer);
+delegate_noop!(LayerShellState: ignore ZwlrLayerShellV1);
+delegate_noop!(LayerShellState: ignore WlSurface);
+
+impl wayland_client::Dispatch<ZwlrLayerSurfaceV1, WlOutput> for LayerShellState {
+    // No need to instrument here, span from lib.rs is automatically used.
+    fn event(
+        state: &mut Self,
+        proxy: &ZwlrLayerSurfaceV1,
+        event: <ZwlrLayerSurfaceV1 as wayland_client::Proxy>::Event,
+        data: &WlOutput,
+        _conn: &Connection,
+        _qhandle: &QueueHandle<Self>,
+    ) {
+        match event {
+            zwlr_layer_surface_v1::Event::Configure {
+                serial,
+                width: _,
+                height: _,
+            } => {
+                tracing::debug!("Acking configure");
+                state.configured_outputs.insert(data.clone());
+                proxy.ack_configure(serial);
+                tracing::trace!("Acked configure");
+            }
+            zwlr_layer_surface_v1::Event::Closed => {
+                tracing::debug!("Closed")
+            }
+            _ => {}
+        }
+    }
+}

libwayshot/src/error.rs

@@ -27,4 +27,6 @@ pub enum Error {
     NoSupportedBufferFormat,
     #[error("Cannot find required wayland protocol")]
     ProtocolNotFound(String),
+    #[error("error occurred in freeze callback")]
+    FreezeCallbackError,
 }

libwayshot/src/lib.rs

@@ -12,13 +12,15 @@ pub mod region;
 mod screencopy;
 
 use std::{
+    collections::HashSet,
     fs::File,
     os::fd::AsFd,
     process::exit,
     sync::atomic::{AtomicBool, Ordering},
     thread,
 };
 
+use dispatch::LayerShellState;
 use image::{imageops::replace, DynamicImage};
 use memmap2::MmapMut;
 use region::{EmbeddedRegion, RegionCapturer};
@@ -27,6 +29,7 @@ use tracing::debug;
 use wayland_client::{
     globals::{registry_queue_init, GlobalList},
     protocol::{
+        wl_compositor::WlCompositor,
         wl_output::WlOutput,
         wl_shm::{self, WlShm},
     },
@@ -35,9 +38,15 @@ use wayland_client::{
 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::ZwlrScreencopyFrameV1,
-    zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
+use wayland_protocols_wlr::{
+    layer_shell::v1::client::{
+        zwlr_layer_shell_v1::{Layer, ZwlrLayerShellV1},
+        zwlr_layer_surface_v1::Anchor,
+    },
+    screencopy::v1::client::{
+        zwlr_screencopy_frame_v1::ZwlrScreencopyFrameV1,
+        zwlr_screencopy_manager_v1::ZwlrScreencopyManagerV1,
+    },
 };
 
 use crate::{
@@ -406,6 +415,87 @@ impl WayshotConnection {
         Ok(frame_copies)
     }
 
+    fn overlay_frames(&self, frames: &Vec<(FrameCopy, FrameGuard, OutputInfo)>) -> Result<()> {
+        let mut state = LayerShellState {
+            configured_outputs: HashSet::new(),
+        };
+        let mut event_queue: EventQueue<LayerShellState> =
+            self.conn.new_event_queue::<LayerShellState>();
+        let qh = event_queue.handle();
+
+        let compositor = match self.globals.bind::<WlCompositor, _, _>(&qh, 3..=3, ()) {
+            Ok(x) => x,
+            Err(e) => {
+                tracing::error!(
+                    "Failed to create compositor Does your compositor implement WlCompositor?"
+                );
+                tracing::error!("err: {e}");
+                return Err(Error::ProtocolNotFound(
+                    "WlCompositor not found".to_string(),
+                ));
+            }
+        };
+        let layer_shell = match self.globals.bind::<ZwlrLayerShellV1, _, _>(&qh, 1..=1, ()) {
+            Ok(x) => x,
+            Err(e) => {
+                tracing::error!(
+                    "Failed to create layer shell. Does your compositor implement WlrLayerShellV1?"
+                );
+                tracing::error!("err: {e}");
+                return Err(Error::ProtocolNotFound(
+                    "WlrLayerShellV1 not found".to_string(),
+                ));
+            }
+        };
+
+        for (frame_copy, frame_guard, output_info) in frames {
+            tracing::span!(
+                tracing::Level::DEBUG,
+                "overlay_frames::surface",
+                output = output_info.name.as_str()
+            )
+            .in_scope(|| -> Result<()> {
+                let surface = compositor.create_surface(&qh, ());
+
+                let layer_surface = layer_shell.get_layer_surface(
+                    &surface,
+                    Some(&output_info.wl_output),
+                    Layer::Top,
+                    "wayshot".to_string(),
+                    &qh,
+                    output_info.wl_output.clone(),
+                );
+
+                layer_surface.set_exclusive_zone(-1);
+                layer_surface.set_anchor(Anchor::Top | Anchor::Left);
+                layer_surface.set_size(
+                    frame_copy.frame_format.width,
+                    frame_copy.frame_format.height,
+                );
+
+                debug!("Committing surface creation changes.");
+                surface.commit();
+
+                debug!("Waiting for layer surface to be configured.");
+                while !state.configured_outputs.contains(&output_info.wl_output) {
+                    event_queue.blocking_dispatch(&mut state)?;
+                }
+
+                surface.set_buffer_transform(output_info.transform);
+                surface.set_buffer_scale(output_info.scale);
+                surface.attach(Some(&frame_guard.buffer), 0, 0);
+
+                debug!("Committing surface with attached buffer.");
+                surface.commit();
+
+                event_queue.blocking_dispatch(&mut state)?;
+
+                Ok(())
+            })?;
+        }
+        Ok(())
+    }
+
     /// Take a screenshot from the specified region.
     fn screenshot_region_capturer(
         &self,
@@ -445,13 +535,21 @@ impl WayshotConnection {
                         })
                     })
                     .collect(),
+                RegionCapturer::Freeze(_) => self
+                    .get_all_outputs()
+                    .into_iter()
+                    .map(|output_info| (output_info.clone(), None))
+                    .collect(),
             };
 
         let frames = self.capture_frame_copies(outputs_capture_regions, cursor_overlay)?;
 
         let capture_region: LogicalRegion = match region_capturer {
             RegionCapturer::Outputs(ref outputs) => outputs.try_into()?,
             RegionCapturer::Region(region) => region,
+            RegionCapturer::Freeze(callback) => {
+                self.overlay_frames(&frames).and_then(|_| callback())?
+            }
         };
 
         thread::scope(|scope| {
@@ -529,6 +627,15 @@ impl WayshotConnection {
         self.screenshot_region_capturer(RegionCapturer::Region(capture_region), cursor_overlay)
     }
 
+    /// Take a screenshot, overlay the screenshot, run the callback, and then
+    /// unfreeze the screenshot and return the selected region.
+    pub fn screenshot_freeze(
+        &self,
+        callback: Box<dyn Fn() -> Result<LogicalRegion>>,
+        cursor_overlay: bool,
+    ) -> Result<DynamicImage> {
+        self.screenshot_region_capturer(RegionCapturer::Freeze(callback), cursor_overlay)
+    }
     /// shot one ouput
     pub fn screenshot_single_output(
         &self,

libwayshot/src/output.rs

@@ -3,23 +3,24 @@ use wayland_client::protocol::{wl_output, wl_output::WlOutput};
 /// Represents an accessible wayland output.
 ///
 /// Do not instantiate, instead use [`crate::WayshotConnection::get_all_outputs`].
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct OutputInfo {
     pub wl_output: WlOutput,
     pub name: String,
     pub description: String,
     pub transform: wl_output::Transform,
+    pub scale: i32,
     pub dimensions: OutputPositioning,
     pub mode: WlOutputMode,
 }
 
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
+#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
 pub struct WlOutputMode {
     pub width: i32,
     pub height: i32,
 }
 
-#[derive(Default, Debug, Clone, PartialEq, Eq)]
+#[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
 pub struct OutputPositioning {
     pub x: i32,
     pub y: i32,

libwayshot/src/region.rs

@@ -2,7 +2,7 @@ use std::cmp;
 
 use wayland_client::protocol::wl_output::Transform;
 
-use crate::error::Error;
+use crate::error::{Error, Result};
 use crate::output::OutputInfo;
 use crate::screencopy::FrameCopy;
 
@@ -12,6 +12,10 @@ pub enum RegionCapturer {
     Outputs(Vec<OutputInfo>),
     /// Capture an already known `LogicalRegion`.
     Region(LogicalRegion),
+    /// The outputs will be "frozen" to the user at which point the given
+    /// callback is called to get the region to capture. This callback is often
+    /// a user interaction to let the user select a region.
+    Freeze(Box<dyn Fn() -> Result<LogicalRegion>>),
 }
 
 /// `Region` where the coordinate system is the logical coordinate system used

wayshot/Cargo.toml

@@ -29,6 +29,7 @@ image = { version = "0.24", default-features = false, features = [
 ] }
 
 dialoguer = { version = "0.11.0", features = ["fuzzy-select"] }
+eyre = "0.6.8"
 
 [[bin]]
 name = "wayshot"

wayshot/src/clap.rs

@@ -12,10 +12,10 @@ pub fn set_flags() -> Command {
                 .help("Enable debug mode"),
         )
         .arg(
-            arg!(-s --slurp <GEOMETRY>)
+            arg!(-s --slurp <SLURP_ARGS>)
                 .required(false)
                 .action(ArgAction::Set)
-                .help("Choose a portion of your display to screenshot using slurp"),
+                .help("Arguments to call slurp with for selecting a region"),
         )
         .arg(
             arg!(-f - -file <FILE_PATH>)