Implement toggling pointer for the screenshot UI

3 files changed, 179 insertions, 34 deletions

diff --git a/src/input/mod.rs b/src/input/mod.rs
index e74ff100..5251f94f 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -508,6 +508,10 @@ impl State {
                     .set_cursor_image(CursorImageStatus::default_named());
                 self.niri.queue_redraw_all();
             }
+            Action::ScreenshotTogglePointer => {
+                self.niri.screenshot_ui.toggle_pointer();
+                self.niri.queue_redraw_all();
+            }
             Action::Screenshot => {
                 self.backend.with_primary_renderer(|renderer| {
                     self.niri.open_screenshot_ui(renderer);
diff --git a/src/niri.rs b/src/niri.rs
index 124c3558..852dc5ee 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -124,13 +124,14 @@ use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
 use crate::render_helpers::renderer::NiriRenderer;
 use crate::render_helpers::texture::TextureBuffer;
 use crate::render_helpers::{
-    render_to_shm, render_to_texture, render_to_vec, shaders, RenderTarget,
+    render_to_encompassing_texture, render_to_shm, render_to_texture, render_to_vec, shaders,
+    RenderTarget,
 };
 use crate::ui::config_error_notification::ConfigErrorNotification;
 use crate::ui::exit_confirm_dialog::ExitConfirmDialog;
 use crate::ui::hotkey_overlay::HotkeyOverlay;
 use crate::ui::screen_transition::{self, ScreenTransition};
-use crate::ui::screenshot_ui::{ScreenshotUi, ScreenshotUiRenderElement};
+use crate::ui::screenshot_ui::{OutputScreenshot, ScreenshotUi, ScreenshotUiRenderElement};
 use crate::utils::scale::{closest_representable_scale, guess_monitor_scale};
 use crate::utils::spawning::CHILD_ENV;
 use crate::utils::{
@@ -3790,8 +3791,8 @@ impl Niri {
                     RenderTarget::Screencast,
                     RenderTarget::ScreenCapture,
                 ];
-                let textures = targets.map(|target| {
-                    let elements = self.render::<GlesRenderer>(renderer, &output, true, target);
+                let screenshot = targets.map(|target| {
+                    let elements = self.render::<GlesRenderer>(renderer, &output, false, target);
                     let elements = elements.iter().rev();
 
                     let res = render_to_texture(
@@ -3802,25 +3803,48 @@ impl Niri {
                         Fourcc::Abgr8888,
                         elements,
                     );
-
                     if let Err(err) = &res {
                         warn!("error rendering output {}: {err:?}", output.name());
                     }
+                    let res_output = res.ok();
 
-                    res
+                    let pointer = self.pointer_element(renderer, &output);
+                    let res_pointer = if pointer.is_empty() {
+                        None
+                    } else {
+                        let res = render_to_encompassing_texture(
+                            renderer,
+                            scale,
+                            Transform::Normal,
+                            Fourcc::Abgr8888,
+                            &pointer,
+                        );
+                        if let Err(err) = &res {
+                            warn!("error rendering pointer for {}: {err:?}", output.name());
+                        }
+                        res.ok()
+                    };
+
+                    res_output.map(|(texture, _)| {
+                        OutputScreenshot::from_textures(
+                            renderer,
+                            scale,
+                            texture,
+                            res_pointer.map(|(texture, _, geo)| (texture, geo)),
+                        )
+                    })
                 });
 
-                if textures.iter().any(|res| res.is_err()) {
+                if screenshot.iter().any(|res| res.is_none()) {
                     return None;
                 }
 
-                let textures = textures.map(|res| res.unwrap().0);
-                Some((output, textures))
+                let screenshot = screenshot.map(|res| res.unwrap());
+                Some((output, screenshot))
             })
             .collect();
 
-        self.screenshot_ui
-            .open(renderer, screenshots, default_output);
+        self.screenshot_ui.open(screenshots, default_output);
         self.cursor_manager
             .set_cursor_image(CursorImageStatus::Named(CursorIcon::Crosshair));
         self.queue_redraw_all();
diff --git a/src/ui/screenshot_ui.rs b/src/ui/screenshot_ui.rs
index 6d073220..2f321637 100644
--- a/src/ui/screenshot_ui.rs
+++ b/src/ui/screenshot_ui.rs
@@ -13,13 +13,13 @@ use smithay::backend::renderer::gles::{GlesRenderer, GlesTexture};
 use smithay::backend::renderer::ExportMem;
 use smithay::input::keyboard::{Keysym, ModifiersState};
 use smithay::output::{Output, WeakOutput};
-use smithay::utils::{Physical, Point, Rectangle, Size, Transform};
+use smithay::utils::{Logical, Physical, Point, Rectangle, Scale, Size, Transform};
 
 use crate::niri_render_elements;
 use crate::render_helpers::primary_gpu_texture::PrimaryGpuTextureRenderElement;
 use crate::render_helpers::solid_color::{SolidColorBuffer, SolidColorRenderElement};
 use crate::render_helpers::texture::{TextureBuffer, TextureRenderElement};
-use crate::render_helpers::RenderTarget;
+use crate::render_helpers::{render_to_texture, RenderTarget};
 use crate::utils::to_physical_precise_round;
 
 const BORDER: i32 = 2;
@@ -37,6 +37,7 @@ pub enum ScreenshotUi {
         selection: (Output, Point<i32, Physical>, Point<i32, Physical>),
         output_data: HashMap<Output, OutputData>,
         mouse_down: bool,
+        show_pointer: bool,
     },
 }
 
@@ -45,12 +46,17 @@ pub struct OutputData {
     scale: f64,
     transform: Transform,
     // Output, screencast, screen capture.
-    texture: [GlesTexture; 3],
-    texture_buffer: [TextureBuffer<GlesTexture>; 3],
+    screenshot: [OutputScreenshot; 3],
     buffers: [SolidColorBuffer; 8],
     locations: [Point<i32, Physical>; 8],
 }
 
+pub struct OutputScreenshot {
+    texture: GlesTexture,
+    buffer: TextureBuffer<GlesTexture>,
+    pointer: Option<(TextureBuffer<GlesTexture>, Rectangle<f64, Logical>)>,
+}
+
 niri_render_elements! {
     ScreenshotUiRenderElement => {
         Screenshot = PrimaryGpuTextureRenderElement,
@@ -67,9 +73,8 @@ impl ScreenshotUi {
 
     pub fn open(
         &mut self,
-        renderer: &GlesRenderer,
         // Output, screencast, screen capture.
-        screenshots: HashMap<Output, [GlesTexture; 3]>,
+        screenshots: HashMap<Output, [OutputScreenshot; 3]>,
         default_output: Output,
     ) -> bool {
         if screenshots.is_empty() {
@@ -108,20 +113,11 @@ impl ScreenshotUi {
 
         let output_data = screenshots
             .into_iter()
-            .map(|(output, texture)| {
+            .map(|(output, screenshot)| {
                 let transform = output.current_transform();
                 let output_mode = output.current_mode().unwrap();
                 let size = transform.transform_size(output_mode.size);
                 let scale = output.current_scale().fractional_scale();
-                let texture_buffer = texture.clone().map(|texture| {
-                    TextureBuffer::from_texture(
-                        renderer,
-                        texture,
-                        scale,
-                        Transform::Normal,
-                        Vec::new(),
-                    )
-                });
                 let buffers = [
                     SolidColorBuffer::new((0., 0.), [1., 1., 1., 1.]),
                     SolidColorBuffer::new((0., 0.), [1., 1., 1., 1.]),
@@ -137,8 +133,7 @@ impl ScreenshotUi {
                     size,
                     scale,
                     transform,
-                    texture,
-                    texture_buffer,
+                    screenshot,
                     buffers,
                     locations,
                 };
@@ -150,6 +145,7 @@ impl ScreenshotUi {
             selection,
             output_data,
             mouse_down: false,
+            show_pointer: true,
         };
 
         self.update_buffers();
@@ -177,6 +173,12 @@ impl ScreenshotUi {
         true
     }
 
+    pub fn toggle_pointer(&mut self) {
+        if let Self::Open { show_pointer, .. } = self {
+            *show_pointer = !*show_pointer;
+        }
+    }
+
     pub fn is_open(&self) -> bool {
         matches!(self, ScreenshotUi::Open { .. })
     }
@@ -258,10 +260,15 @@ impl ScreenshotUi {
         &self,
         output: &Output,
         target: RenderTarget,
-    ) -> ArrayVec<ScreenshotUiRenderElement, 9> {
+    ) -> ArrayVec<ScreenshotUiRenderElement, 10> {
         let _span = tracy_client::span!("ScreenshotUi::render_output");
 
-        let Self::Open { output_data, .. } = self else {
+        let Self::Open {
+            output_data,
+            show_pointer,
+            ..
+        } = self
+        else {
             panic!("screenshot UI must be open to render it");
         };
 
@@ -288,9 +295,26 @@ impl ScreenshotUi {
             RenderTarget::Screencast => 1,
             RenderTarget::ScreenCapture => 2,
         };
+
+        if *show_pointer {
+            if let Some((buffer, geo)) = output_data.screenshot[index].pointer.clone() {
+                elements.push(
+                    PrimaryGpuTextureRenderElement(TextureRenderElement::from_texture_buffer(
+                        buffer,
+                        geo.loc,
+                        1.,
+                        None,
+                        None,
+                        Kind::Unspecified,
+                    ))
+                    .into(),
+                );
+            }
+        }
+
         elements.push(
             PrimaryGpuTextureRenderElement(TextureRenderElement::from_texture_buffer(
-                output_data.texture_buffer[index].clone(),
+                output_data.screenshot[index].buffer.clone(),
                 (0., 0.),
                 1.,
                 None,
@@ -312,6 +336,7 @@ impl ScreenshotUi {
         let Self::Open {
             selection,
             output_data,
+            show_pointer,
             ..
         } = self
         else {
@@ -320,12 +345,67 @@ impl ScreenshotUi {
 
         let data = &output_data[&selection.0];
         let rect = rect_from_corner_points(selection.1, selection.2);
+
+        let screenshot = &data.screenshot[0];
+
+        // Composite the pointer on top if needed.
+        let mut tex_rect = None;
+        if *show_pointer {
+            if let Some((buffer, geo)) = screenshot.pointer.clone() {
+                let scale = buffer.texture_scale();
+                let offset = rect.loc.to_f64().to_logical(scale);
+                let offset = offset.upscale(-1.);
+
+                let mut elements = ArrayVec::<_, 2>::new();
+                elements.push(PrimaryGpuTextureRenderElement(
+                    TextureRenderElement::from_texture_buffer(
+                        buffer,
+                        geo.loc + offset,
+                        1.,
+                        None,
+                        None,
+                        Kind::Unspecified,
+                    ),
+                ));
+                elements.push(PrimaryGpuTextureRenderElement(
+                    TextureRenderElement::from_texture_buffer(
+                        screenshot.buffer.clone(),
+                        offset,
+                        1.,
+                        None,
+                        None,
+                        Kind::Unspecified,
+                    ),
+                ));
+                let elements = elements.iter().rev();
+
+                let res = render_to_texture(
+                    renderer,
+                    rect.size,
+                    scale,
+                    Transform::Normal,
+                    Fourcc::Abgr8888,
+                    elements,
+                );
+                match res {
+                    Ok((texture, _)) => {
+                        tex_rect = Some((texture, Rectangle::from_loc_and_size((0, 0), rect.size)));
+                    }
+                    Err(err) => {
+                        warn!("error compositing pointer onto screenshot: {err:?}");
+                    }
+                }
+            }
+        }
+
+        let (texture, rect) = tex_rect.unwrap_or_else(|| (screenshot.texture.clone(), rect));
+        // The size doesn't actually matter because we're not transforming anything.
         let buf_rect = rect
             .to_logical(1)
-            .to_buffer(1, Transform::Normal, &data.size.to_logical(1));
+            .to_buffer(1, Transform::Normal, &Size::from((1, 1)));
 
         let mapping = renderer
-            .copy_texture(&data.texture[0], buf_rect, Fourcc::Abgr8888)
+            .copy_texture(&texture, buf_rect, Fourcc::Abgr8888)
             .context("error copying texture")?;
         let copy = renderer
             .map_texture(&mapping)
@@ -389,6 +469,7 @@ impl ScreenshotUi {
             selection,
             output_data,
             mouse_down,
+            ..
         } = self
         else {
             return false;
@@ -438,6 +519,38 @@ impl Default for ScreenshotUi {
     }
 }
 
+impl OutputScreenshot {
+    pub fn from_textures(
+        renderer: &mut GlesRenderer,
+        scale: Scale<f64>,
+        texture: GlesTexture,
+        pointer: Option<(GlesTexture, Rectangle<i32, Physical>)>,
+    ) -> Self {
+        Self {
+            texture: texture.clone(),
+            buffer: TextureBuffer::from_texture(
+                renderer,
+                texture,
+                scale,
+                Transform::Normal,
+                Vec::new(),
+            ),
+            pointer: pointer.map(|(texture, geo)| {
+                (
+                    TextureBuffer::from_texture(
+                        renderer,
+                        texture,
+                        scale,
+                        Transform::Normal,
+                        Vec::new(),
+                    ),
+                    geo.to_f64().to_logical(scale),
+                )
+            }),
+        }
+    }
+}
+
 fn action(raw: Keysym, mods: ModifiersState) -> Option<Action> {
     if raw == Keysym::Escape {
         return Some(Action::CancelScreenshot);
@@ -453,6 +566,10 @@ fn action(raw: Keysym, mods: ModifiersState) -> Option<Action> {
         return Some(Action::ConfirmScreenshot);
     }
 
+    if !mods.ctrl && raw == Keysym::p {
+        return Some(Action::ScreenshotTogglePointer);
+    }
+
     None
 }