Implement area selection screenshots

resources/default-config.kdl

@@ -213,7 +213,8 @@ binds {
     Mod+Minus { set-column-width "-10%"; }
     Mod+Equal { set-column-width "+10%"; }
 
-    Print { screenshot-screen; }
+    Print { screenshot; }
+    Ctrl+Print { screenshot-screen; }
     Alt+Print { screenshot-window; }
 
     Mod+Shift+E { quit; }

src/config.rs

@@ -236,6 +236,11 @@ pub enum Action {
     PowerOffMonitors,
     ToggleDebugTint,
     Spawn(#[knuffel(arguments)] Vec<String>),
+    #[knuffel(skip)]
+    ConfirmScreenshot,
+    #[knuffel(skip)]
+    CancelScreenshot,
+    Screenshot,
     ScreenshotScreen,
     ScreenshotWindow,
     CloseWindow,

src/input.rs

@@ -19,6 +19,7 @@ use smithay::wayland::tablet_manager::{TabletDescriptor, TabletSeatTrait};
 
 use crate::config::{Action, Binds, Modifiers};
 use crate::niri::State;
+use crate::screenshot_ui::ScreenshotUi;
 use crate::utils::{center, get_monotonic_time, spawn};
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -70,6 +71,7 @@ impl State {
                             raw,
                             pressed,
                             *mods,
+                            &this.niri.screenshot_ui,
                         )
                     },
                 ) else {
@@ -129,6 +131,32 @@ impl State {
                             }
                         }
                     }
+                    Action::ConfirmScreenshot => {
+                        if let Some(renderer) = self.backend.renderer() {
+                            match self.niri.screenshot_ui.capture(renderer) {
+                                Ok((size, pixels)) => {
+                                    if let Err(err) = self.niri.save_screenshot(size, pixels) {
+                                        warn!("error saving screenshot: {err:?}");
+                                    }
+                                }
+                                Err(err) => {
+                                    warn!("error capturing screenshot: {err:?}");
+                                }
+                            }
+                        }
+
+                        self.niri.screenshot_ui.close();
+                        self.niri.queue_redraw_all();
+                    }
+                    Action::CancelScreenshot => {
+                        self.niri.screenshot_ui.close();
+                        self.niri.queue_redraw_all();
+                    }
+                    Action::Screenshot => {
+                        if let Some(renderer) = self.backend.renderer() {
+                            self.niri.open_screenshot_ui(renderer);
+                        }
+                    }
                     Action::ScreenshotWindow => {
                         let active = self.niri.layout.active_window();
                         if let Some((window, output)) = active {
@@ -335,6 +363,21 @@ impl State {
                     }
                 }
 
+                if let Some(output) = self.niri.screenshot_ui.selection_output() {
+                    let geom = self.niri.global_space.output_geometry(output).unwrap();
+                    let mut point = new_pos;
+                    point.x = point
+                        .x
+                        .clamp(geom.loc.x as f64, (geom.loc.x + geom.size.w - 1) as f64);
+                    point.y = point
+                        .y
+                        .clamp(geom.loc.y as f64, (geom.loc.y + geom.size.h - 1) as f64);
+                    let point = (point - geom.loc.to_f64())
+                        .to_physical(output.current_scale().fractional_scale())
+                        .to_i32_round();
+                    self.niri.screenshot_ui.pointer_motion(point);
+                }
+
                 let under = self.niri.surface_under_and_global_space(new_pos);
                 self.niri.pointer_focus = under.clone();
                 let under = under.map(|u| u.surface);
@@ -378,6 +421,21 @@ impl State {
 
                 let pointer = self.niri.seat.get_pointer().unwrap();
 
+                if let Some(output) = self.niri.screenshot_ui.selection_output() {
+                    let geom = self.niri.global_space.output_geometry(output).unwrap();
+                    let mut point = pos;
+                    point.x = point
+                        .x
+                        .clamp(geom.loc.x as f64, (geom.loc.x + geom.size.w - 1) as f64);
+                    point.y = point
+                        .y
+                        .clamp(geom.loc.y as f64, (geom.loc.y + geom.size.h - 1) as f64);
+                    let point = (point - geom.loc.to_f64())
+                        .to_physical(output.current_scale().fractional_scale())
+                        .to_i32_round();
+                    self.niri.screenshot_ui.pointer_motion(point);
+                }
+
                 let under = self.niri.surface_under_and_global_space(pos);
                 self.niri.pointer_focus = under.clone();
                 let under = under.map(|u| u.surface);
@@ -418,6 +476,34 @@ impl State {
 
                 self.update_pointer_focus();
 
+                if let Some(button) = event.button() {
+                    let pos = pointer.current_location();
+                    if let Some((output, _)) = self.niri.output_under(pos) {
+                        let output = output.clone();
+                        let geom = self.niri.global_space.output_geometry(&output).unwrap();
+                        let mut point = pos;
+                        // Re-clamp as pointer can be within 0.5 from the limit which will round up
+                        // to a wrong value.
+                        point.x = point
+                            .x
+                            .clamp(geom.loc.x as f64, (geom.loc.x + geom.size.w - 1) as f64);
+                        point.y = point
+                            .y
+                            .clamp(geom.loc.y as f64, (geom.loc.y + geom.size.h - 1) as f64);
+                        let point = (point - geom.loc.to_f64())
+                            .to_physical(output.current_scale().fractional_scale())
+                            .to_i32_round();
+                        if self.niri.screenshot_ui.pointer_button(
+                            output,
+                            point,
+                            button,
+                            button_state,
+                        ) {
+                            self.niri.queue_redraw_all();
+                        }
+                    }
+                }
+
                 pointer.button(
                     self,
                     &ButtonEvent {
@@ -842,6 +928,7 @@ fn should_intercept_key(
     raw: Option<Keysym>,
     pressed: bool,
     mods: ModifiersState,
+    screenshot_ui: &ScreenshotUi,
 ) -> FilterResult<Option<Action>> {
     // Actions are only triggered on presses, release of the key
     // shouldn't try to intercept anything unless we have marked
@@ -850,7 +937,20 @@ fn should_intercept_key(
         return FilterResult::Forward;
     }
 
-    match (action(bindings, comp_mod, modified, raw, mods), pressed) {
+    let mut final_action = action(bindings, comp_mod, modified, raw, mods);
+    if screenshot_ui.is_open()
+        // Allow only a subset of compositor actions while the screenshot UI is open,
+        // since the user cannot see the screen.
+        && !matches!(
+            final_action,
+            Some(Action::Quit | Action::ChangeVt(_) | Action::Suspend | Action::PowerOffMonitors)
+        )
+    {
+        // Otherwise, use the screenshot UI action.
+        final_action = screenshot_ui.action(raw, mods);
+    }
+
+    match (final_action, pressed) {
         (Some(action), true) => {
             suppressed_keys.insert(key_code);
             FilterResult::Intercept(Some(action))
@@ -965,6 +1065,8 @@ mod tests {
         let comp_mod = CompositorMod::Super;
         let mut suppressed_keys = HashSet::new();
 
+        let screenshot_ui = ScreenshotUi::new();
+
         // The key_code we pick is arbitrary, the only thing
         // that matters is that they are different between cases.
 
@@ -979,6 +1081,7 @@ mod tests {
                 Some(close_keysym),
                 pressed,
                 mods,
+                &screenshot_ui,
             )
         };
 
@@ -993,6 +1096,7 @@ mod tests {
                 Some(Keysym::l),
                 pressed,
                 mods,
+                &screenshot_ui,
             )
         };
 

src/main.rs

@@ -12,6 +12,7 @@ mod handlers;
 mod input;
 mod layout;
 mod niri;
+mod screenshot_ui;
 mod utils;
 mod watcher;
 

src/niri.rs

@@ -20,6 +20,7 @@ use smithay::backend::renderer::element::{
     RenderElementStates,
 };
 use smithay::backend::renderer::gles::{GlesMapping, GlesRenderer, GlesTexture};
+use smithay::backend::renderer::sync::SyncPoint;
 use smithay::backend::renderer::{Bind, ExportMem, Frame, ImportAll, Offscreen, Renderer};
 use smithay::desktop::utils::{
     bbox_from_surface_tree, output_update, send_dmabuf_feedback_surface_tree,
@@ -85,6 +86,7 @@ use crate::frame_clock::FrameClock;
 use crate::handlers::configure_lock_surface;
 use crate::layout::{output_size, Layout, MonitorRenderElement};
 use crate::pw_utils::{Cast, PipeWire};
+use crate::screenshot_ui::{ScreenshotUi, ScreenshotUiRenderElement};
 use crate::utils::{center, get_monotonic_time, make_screenshot_path, write_png_rgba8};
 
 const CLEAR_COLOR: [f32; 4] = [0.2, 0.2, 0.2, 1.];
@@ -150,6 +152,8 @@ pub struct Niri {
 
     pub lock_state: LockState,
 
+    pub screenshot_ui: ScreenshotUi,
+
     #[cfg(feature = "dbus")]
     pub dbus: Option<crate::dbus::DBusServers>,
     #[cfg(feature = "dbus")]
@@ -310,7 +314,7 @@ impl State {
         let pointer = &self.niri.seat.get_pointer().unwrap();
         let location = pointer.current_location();
 
-        if !self.niri.is_locked() {
+        if !self.niri.is_locked() && !self.niri.screenshot_ui.is_open() {
             // Don't refresh cursor focus during transitions.
             if let Some((output, _)) = self.niri.output_under(location) {
                 let monitor = self.niri.layout.monitor_for_output(output).unwrap();
@@ -366,6 +370,8 @@ impl State {
     pub fn update_focus(&mut self) {
         let focus = if self.niri.is_locked() {
             self.niri.lock_surface_focus()
+        } else if self.niri.screenshot_ui.is_open() {
+            None
         } else {
             self.niri.layer_surface_focus().or_else(|| {
                 self.niri
@@ -580,6 +586,8 @@ impl Niri {
         let cursor_manager =
             CursorManager::new(&config_.cursor.xcursor_theme, config_.cursor.xcursor_size);
 
+        let screenshot_ui = ScreenshotUi::new();
+
         let socket_source = ListeningSocketSource::new_auto().unwrap();
         let socket_name = socket_source.socket_name().to_os_string();
         event_loop
@@ -657,6 +665,8 @@ impl Niri {
 
             lock_state: LockState::Unlocked,
 
+            screenshot_ui,
+
             #[cfg(feature = "dbus")]
             dbus: None,
             #[cfg(feature = "dbus")]
@@ -832,6 +842,10 @@ impl Niri {
             }
             lock_state => self.lock_state = lock_state,
         }
+
+        if self.screenshot_ui.close() {
+            self.queue_redraw_all();
+        }
     }
 
     pub fn output_resized(&mut self, output: Output) {
@@ -852,6 +866,18 @@ impl Niri {
             }
         }
 
+        // If the output size changed with an open screenshot UI, close the screenshot UI.
+        if let Some(old_size) = self.screenshot_ui.output_size(&output) {
+            let output_transform = output.current_transform();
+            let output_mode = output.current_mode().unwrap();
+            let size = output_transform.transform_size(output_mode.size);
+            if old_size != size {
+                self.screenshot_ui.close();
+                self.queue_redraw_all();
+                return;
+            }
+        }
+
         self.queue_redraw(output);
     }
 
@@ -889,7 +915,7 @@ impl Niri {
     }
 
     pub fn window_under_cursor(&self) -> Option<&Window> {
-        if self.is_locked() {
+        if self.is_locked() || self.screenshot_ui.is_open() {
             return None;
         }
 
@@ -928,6 +954,10 @@ impl Niri {
             });
         }
 
+        if self.screenshot_ui.is_open() {
+            return None;
+        }
+
         let (window, win_pos_within_output) =
             self.layout.window_under(&output, pos_within_output)?;
 
@@ -1321,6 +1351,32 @@ impl Niri {
             return elements;
         }
 
+        // Prepare the background element.
+        let state = self.output_state.get(output).unwrap();
+        let background = SolidColorRenderElement::from_buffer(
+            &state.background_buffer,
+            (0, 0),
+            output_scale,
+            1.,
+            Kind::Unspecified,
+        )
+        .into();
+
+        // If the screenshot UI is open, draw it.
+        if self.screenshot_ui.is_open() {
+            elements.extend(
+                self.screenshot_ui
+                    .render_output(output)
+                    .into_iter()
+                    .map(OutputRenderElements::from),
+            );
+
+            // Add the background for outputs that were connected while the screenshot UI was open.
+            elements.push(background);
+
+            return elements;
+        }
+
         // Get monitor elements.
         let mon = self.layout.monitor_for_output(output).unwrap();
         let monitor_elements = mon.render_elements(renderer);
@@ -1363,17 +1419,7 @@ impl Niri {
         extend_from_layer(&mut elements, Layer::Background);
 
         // Then the background.
-        let state = self.output_state.get(output).unwrap();
-        elements.push(
-            SolidColorRenderElement::from_buffer(
-                &state.background_buffer,
-                (0, 0),
-                output_scale,
-                1.,
-                Kind::Unspecified,
-            )
-            .into(),
-        );
+        elements.push(background);
 
         elements
     }
@@ -1879,6 +1925,42 @@ impl Niri {
         }
     }
 
+    pub fn open_screenshot_ui(&mut self, renderer: &mut GlesRenderer) {
+        if self.is_locked() || self.screenshot_ui.is_open() {
+            return;
+        }
+
+        let Some(default_output) = self.output_under_cursor() else {
+            return;
+        };
+
+        let screenshots = self
+            .global_space
+            .outputs()
+            .cloned()
+            .filter_map(|output| {
+                let size = output.current_mode().unwrap().size;
+                let scale = Scale::from(output.current_scale().fractional_scale());
+                let elements = self.render(renderer, &output, true);
+
+                let res = render_to_texture(renderer, size, scale, Fourcc::Abgr8888, &elements);
+                let screenshot = match res {
+                    Ok((texture, _)) => texture,
+                    Err(err) => {
+                        warn!("error rendering output {}: {err:?}", output.name());
+                        return None;
+                    }
+                };
+
+                Some((output, screenshot))
+            })
+            .collect();
+
+        self.screenshot_ui
+            .open(renderer, screenshots, default_output);
+        self.queue_redraw_all();
+    }
+
     pub fn screenshot(&self, renderer: &mut GlesRenderer, output: &Output) -> anyhow::Result<()> {
         let _span = tracy_client::span!("Niri::screenshot");
 
@@ -1916,7 +1998,11 @@ impl Niri {
             .context("error saving screenshot")
     }
 
-    fn save_screenshot(&self, size: Size<i32, Physical>, pixels: Vec<u8>) -> anyhow::Result<()> {
+    pub fn save_screenshot(
+        &self,
+        size: Size<i32, Physical>,
+        pixels: Vec<u8>,
+    ) -> anyhow::Result<()> {
         let path = make_screenshot_path().context("error making screenshot path")?;
         debug!("saving screenshot to {path:?}");
 
@@ -2029,6 +2115,8 @@ impl Niri {
     pub fn lock(&mut self, confirmation: SessionLocker) {
         info!("locking session");
 
+        self.screenshot_ui.close();
+
         self.lock_state = LockState::Locking(confirmation);
         self.queue_redraw_all();
     }
@@ -2065,6 +2153,7 @@ render_elements! {
     Wayland = WaylandSurfaceRenderElement<R>,
     NamedPointer = TextureRenderElement<<R as Renderer>::TextureId>,
     SolidColor = SolidColorRenderElement,
+    ScreenshotUi = ScreenshotUiRenderElement<R>,
 }
 
 #[derive(Default)]

src/screenshot_ui.rs

@@ -0,0 +1,448 @@
+use std::cmp::{max, min};
+use std::collections::HashMap;
+use std::iter::zip;
+use std::mem;
+
+use anyhow::Context;
+use arrayvec::ArrayVec;
+use smithay::backend::allocator::Fourcc;
+use smithay::backend::input::{ButtonState, MouseButton};
+use smithay::backend::renderer::element::solid::{SolidColorBuffer, SolidColorRenderElement};
+use smithay::backend::renderer::element::texture::{TextureBuffer, TextureRenderElement};
+use smithay::backend::renderer::element::Kind;
+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::render_elements;
+use smithay::utils::{Physical, Point, Rectangle, Size, Transform};
+
+use crate::config::Action;
+
+const BORDER: i32 = 2;
+
+// Ideally the screenshot UI should support cross-output selections. However, that poses some
+// technical challenges when the outputs have different scales and such. So, this implementation
+// allows only single-output selections for now.
+//
+// As a consequence of this, selection coordinates are in output-local coordinate space.
+pub enum ScreenshotUi {
+    Closed {
+        last_selection: Option<(WeakOutput, Rectangle<i32, Physical>)>,
+    },
+    Open {
+        selection: (Output, Point<i32, Physical>, Point<i32, Physical>),
+        output_data: HashMap<Output, OutputData>,
+        mouse_down: bool,
+    },
+}
+
+pub struct OutputData {
+    size: Size<i32, Physical>,
+    texture: GlesTexture,
+    texture_buffer: TextureBuffer<GlesTexture>,
+    buffers: [SolidColorBuffer; 8],
+    locations: [Point<i32, Physical>; 8],
+}
+
+render_elements! {
+    #[derive(Debug)]
+    pub ScreenshotUiRenderElement<R>;
+    Screenshot = TextureRenderElement<R::TextureId>,
+    SolidColor = SolidColorRenderElement,
+}
+
+impl ScreenshotUi {
+    pub fn new() -> Self {
+        Self::Closed {
+            last_selection: None,
+        }
+    }
+
+    pub fn open(
+        &mut self,
+        renderer: &GlesRenderer,
+        screenshots: HashMap<Output, GlesTexture>,
+        default_output: Output,
+    ) -> bool {
+        if screenshots.is_empty() {
+            return false;
+        }
+
+        let Self::Closed { last_selection } = self else {
+            return false;
+        };
+
+        let last_selection = last_selection
+            .take()
+            .and_then(|(weak, sel)| weak.upgrade().map(|output| (output, sel)));
+        let selection = match last_selection {
+            Some(selection) if screenshots.contains_key(&selection.0) => selection,
+            _ => {
+                let output = default_output;
+                let output_transform = output.current_transform();
+                let output_mode = output.current_mode().unwrap();
+                let size = output_transform.transform_size(output_mode.size);
+                (
+                    output,
+                    Rectangle::from_loc_and_size(
+                        (size.w / 4, size.h / 4),
+                        (size.w / 2, size.h / 2),
+                    ),
+                )
+            }
+        };
+        let scale = selection.0.current_scale().integer_scale();
+        let selection = (
+            selection.0,
+            selection.1.loc,
+            selection.1.loc + selection.1.size - Size::from((scale, scale)),
+        );
+
+        let output_data = screenshots
+            .into_iter()
+            .map(|(output, texture)| {
+                let output_transform = output.current_transform();
+                let output_mode = output.current_mode().unwrap();
+                let size = output_transform.transform_size(output_mode.size);
+                let texture_buffer = TextureBuffer::from_texture(
+                    renderer,
+                    texture.clone(),
+                    output.current_scale().integer_scale(),
+                    Transform::Normal,
+                    None,
+                );
+                let buffers = [
+                    SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
+                    SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
+                    SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
+                    SolidColorBuffer::new((0, 0), [1., 1., 1., 1.]),
+                    SolidColorBuffer::new((0, 0), [0., 0., 0., 0.5]),
+                    SolidColorBuffer::new((0, 0), [0., 0., 0., 0.5]),
+                    SolidColorBuffer::new((0, 0), [0., 0., 0., 0.5]),
+                    SolidColorBuffer::new((0, 0), [0., 0., 0., 0.5]),
+                ];
+                let locations = [Default::default(); 8];
+                let data = OutputData {
+                    size,
+                    texture,
+                    texture_buffer,
+                    buffers,
+                    locations,
+                };
+                (output, data)
+            })
+            .collect();
+
+        *self = Self::Open {
+            selection,
+            output_data,
+            mouse_down: false,
+        };
+
+        self.update_buffers();
+
+        true
+    }
+
+    pub fn close(&mut self) -> bool {
+        let selection = match mem::take(self) {
+            Self::Open { selection, .. } => selection,
+            closed @ Self::Closed { .. } => {
+                // Put it back.
+                *self = closed;
+                return false;
+            }
+        };
+
+        let scale = selection.0.current_scale().integer_scale();
+        let last_selection = Some((
+            selection.0.downgrade(),
+            rect_from_corner_points(selection.1, selection.2, scale),
+        ));
+
+        *self = Self::Closed { last_selection };
+
+        true
+    }
+
+    pub fn is_open(&self) -> bool {
+        matches!(self, ScreenshotUi::Open { .. })
+    }
+
+    fn update_buffers(&mut self) {
+        let Self::Open {
+            selection,
+            output_data,
+            ..
+        } = self
+        else {
+            panic!("screenshot UI must be open to update buffers");
+        };
+
+        let (selection_output, a, b) = selection;
+        let scale = selection_output.current_scale().integer_scale();
+        let mut rect = rect_from_corner_points(*a, *b, scale);
+
+        for (output, data) in output_data {
+            let buffers = &mut data.buffers;
+            let locations = &mut data.locations;
+            let size = data.size;
+
+            if output == selection_output {
+                let scale = output.current_scale().integer_scale();
+
+                // Check if the selection is still valid. If not, reset it back to default.
+                if !Rectangle::from_loc_and_size((0, 0), size).contains_rect(rect) {
+                    rect = Rectangle::from_loc_and_size(
+                        (size.w / 4, size.h / 4),
+                        (size.w / 2, size.h / 2),
+                    );
+                    *a = rect.loc;
+                    *b = rect.loc + rect.size - Size::from((scale, scale));
+                }
+
+                let border = BORDER * scale;
+
+                buffers[0].resize((rect.size.w + border * 2, border));
+                buffers[1].resize((rect.size.w + border * 2, border));
+                buffers[2].resize((border, rect.size.h));
+                buffers[3].resize((border, rect.size.h));
+
+                buffers[4].resize((size.w, rect.loc.y));
+                buffers[5].resize((size.w, size.h - rect.loc.y - rect.size.h));
+                buffers[6].resize((rect.loc.x, rect.size.h));
+                buffers[7].resize((size.w - rect.loc.x - rect.size.w, rect.size.h));
+
+                locations[0] = Point::from((rect.loc.x - border, rect.loc.y - border));
+                locations[1] = Point::from((rect.loc.x - border, rect.loc.y + rect.size.h));
+                locations[2] = Point::from((rect.loc.x - border, rect.loc.y));
+                locations[3] = Point::from((rect.loc.x + rect.size.w, rect.loc.y));
+
+                locations[5] = Point::from((0, rect.loc.y + rect.size.h));
+                locations[6] = Point::from((0, rect.loc.y));
+                locations[7] = Point::from((rect.loc.x + rect.size.w, rect.loc.y));
+            } else {
+                buffers[0].resize((0, 0));
+                buffers[1].resize((0, 0));
+                buffers[2].resize((0, 0));
+                buffers[3].resize((0, 0));
+
+                buffers[4].resize(size.to_logical(1));
+                buffers[5].resize((0, 0));
+                buffers[6].resize((0, 0));
+                buffers[7].resize((0, 0));
+            }
+        }
+    }
+
+    pub fn render_output(
+        &self,
+        output: &Output,
+    ) -> ArrayVec<ScreenshotUiRenderElement<GlesRenderer>, 9> {
+        let _span = tracy_client::span!("ScreenshotUi::render_output");
+
+        let Self::Open { output_data, .. } = self else {
+            panic!("screenshot UI must be open to render it");
+        };
+
+        let mut elements = ArrayVec::new();
+
+        let Some(output_data) = output_data.get(output) else {
+            return elements;
+        };
+
+        let buf_loc = zip(&output_data.buffers, &output_data.locations);
+        elements.extend(buf_loc.map(|(buffer, loc)| {
+            SolidColorRenderElement::from_buffer(
+                buffer,
+                *loc,
+                1., // We treat these as physical coordinates.
+                1.,
+                Kind::Unspecified,
+            )
+            .into()
+        }));
+
+        // The screenshot itself goes last.
+        elements.push(
+            TextureRenderElement::from_texture_buffer(
+                (0., 0.),
+                &output_data.texture_buffer,
+                None,
+                None,
+                None,
+                Kind::Unspecified,
+            )
+            .into(),
+        );
+
+        elements
+    }
+
+    pub fn capture(
+        &self,
+        renderer: &mut GlesRenderer,
+    ) -> anyhow::Result<(Size<i32, Physical>, Vec<u8>)> {
+        let _span = tracy_client::span!("ScreenshotUi::capture");
+
+        let Self::Open {
+            selection,
+            output_data,
+            ..
+        } = self
+        else {
+            panic!("screenshot UI must be open to capture");
+        };
+
+        let data = &output_data[&selection.0];
+        let scale = selection.0.current_scale().integer_scale();
+        let rect = rect_from_corner_points(selection.1, selection.2, scale);
+        let buf_rect = rect
+            .to_logical(1)
+            .to_buffer(1, Transform::Normal, &data.size.to_logical(1));
+
+        let mapping = renderer
+            .copy_texture(&data.texture, buf_rect, Fourcc::Abgr8888)
+            .context("error copying texture")?;
+        let copy = renderer
+            .map_texture(&mapping)
+            .context("error mapping texture")?;
+
+        Ok((rect.size, copy.to_vec()))
+    }
+
+    pub fn action(&self, raw: Option<Keysym>, mods: ModifiersState) -> Option<Action> {
+        if !matches!(self, Self::Open { .. }) {
+            return None;
+        }
+
+        action(raw?, mods)
+    }
+
+    pub fn selection_output(&self) -> Option<&Output> {
+        if let Self::Open {
+            selection: (output, _, _),
+            ..
+        } = self
+        {
+            Some(output)
+        } else {
+            None
+        }
+    }
+
+    pub fn output_size(&self, output: &Output) -> Option<Size<i32, Physical>> {
+        if let Self::Open { output_data, .. } = self {
+            Some(output_data.get(output)?.size)
+        } else {
+            None
+        }
+    }
+
+    /// The pointer has moved to `point` relative to the current selection output.
+    pub fn pointer_motion(&mut self, point: Point<i32, Physical>) {
+        let Self::Open {
+            selection,
+            mouse_down: true,
+            ..
+        } = self
+        else {
+            return;
+        };
+
+        selection.2 = point;
+        self.update_buffers();
+    }
+
+    pub fn pointer_button(
+        &mut self,
+        output: Output,
+        point: Point<i32, Physical>,
+        button: MouseButton,
+        state: ButtonState,
+    ) -> bool {
+        let Self::Open {
+            selection,
+            output_data,
+            mouse_down,
+        } = self
+        else {
+            return false;
+        };
+
+        if button != MouseButton::Left {
+            return false;
+        }
+
+        let down = state == ButtonState::Pressed;
+        if *mouse_down == down {
+            return false;
+        }
+
+        if !output_data.contains_key(&output) {
+            return false;
+        }
+
+        *mouse_down = down;
+
+        if down {
+            *selection = (output, point, point);
+        } else {
+            // Check if the resulting selection is zero-sized, and try to come up with a small
+            // default rectangle.
+            let (output, a, b) = selection;
+            let scale = output.current_scale().integer_scale();
+            let mut rect = rect_from_corner_points(*a, *b, scale);
+            if rect.size.is_empty() || rect.size == Size::from((scale, scale)) {
+                let data = &output_data[output];
+                rect = Rectangle::from_loc_and_size((rect.loc.x - 16, rect.loc.y - 16), (32, 32))
+                    .intersection(Rectangle::from_loc_and_size((0, 0), data.size))
+                    .unwrap_or_default();
+                let scale = output.current_scale().integer_scale();
+                *a = rect.loc;
+                *b = rect.loc + rect.size - Size::from((scale, scale));
+            }
+        }
+
+        self.update_buffers();
+
+        true
+    }
+}
+
+impl Default for ScreenshotUi {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+fn action(raw: Keysym, mods: ModifiersState) -> Option<Action> {
+    if raw == Keysym::Escape {
+        return Some(Action::CancelScreenshot);
+    }
+
+    if mods.alt || mods.shift {
+        return None;
+    }
+
+    if (mods.ctrl && raw == Keysym::c)
+        || (!mods.ctrl && (raw == Keysym::space || raw == Keysym::Return))
+    {
+        return Some(Action::ConfirmScreenshot);
+    }
+
+    None
+}
+
+pub fn rect_from_corner_points(
+    a: Point<i32, Physical>,
+    b: Point<i32, Physical>,
+    scale: i32,
+) -> Rectangle<i32, Physical> {
+    let x1 = min(a.x, b.x);
+    let y1 = min(a.y, b.y);
+    let x2 = max(a.x, b.x);
+    let y2 = max(a.y, b.y);
+    Rectangle::from_extemities((x1, y1), (x2 + scale, y2 + scale))
+}