Add touch selection support to the screenshot UI

3 files changed, 82 insertions, 13 deletions

diff --git a/src/input/mod.rs b/src/input/mod.rs
index 9cc39f74..2fd5ab55 100644
--- a/src/input/mod.rs
+++ b/src/input/mod.rs
@@ -2112,7 +2112,7 @@ impl State {
             point.x = point.x.clamp(0, size.w - 1);
             point.y = point.y.clamp(0, size.h - 1);
 
-            self.niri.screenshot_ui.pointer_motion(point);
+            self.niri.screenshot_ui.pointer_motion(point, None);
         }
 
         let under = self.niri.contents_under(new_pos);
@@ -2245,7 +2245,7 @@ impl State {
             point.x = point.x.clamp(0, size.w - 1);
             point.y = point.y.clamp(0, size.h - 1);
 
-            self.niri.screenshot_ui.pointer_motion(point);
+            self.niri.screenshot_ui.pointer_motion(point, None);
         }
 
         let under = self.niri.contents_under(pos);
@@ -2579,11 +2579,11 @@ impl State {
                     point.x = min(size.w - 1, point.x);
                     point.y = min(size.h - 1, point.y);
 
-                    if self.niri.screenshot_ui.pointer_down(output, point) {
+                    if self.niri.screenshot_ui.pointer_down(output, point, None) {
                         self.niri.queue_redraw_all();
                     }
                 }
-            } else if self.niri.screenshot_ui.pointer_up() {
+            } else if self.niri.screenshot_ui.pointer_up(None) {
                 self.niri.queue_redraw_all();
             }
         }
@@ -3034,7 +3034,7 @@ impl State {
             point.x = point.x.clamp(0, size.w - 1);
             point.y = point.y.clamp(0, size.h - 1);
 
-            self.niri.screenshot_ui.pointer_motion(point);
+            self.niri.screenshot_ui.pointer_motion(point, None);
         }
 
         let under = self.niri.contents_under(pos);
@@ -3110,7 +3110,7 @@ impl State {
                             point.x = min(size.w - 1, point.x);
                             point.y = min(size.h - 1, point.y);
 
-                            if self.niri.screenshot_ui.pointer_down(output, point) {
+                            if self.niri.screenshot_ui.pointer_down(output, point, None) {
                                 self.niri.queue_redraw_all();
                             }
                         }
@@ -3151,7 +3151,7 @@ impl State {
                 }
             }
             TabletToolTipState::Up => {
-                if self.niri.screenshot_ui.pointer_up() {
+                if self.niri.screenshot_ui.pointer_up(None) {
                     self.niri.queue_redraw_all();
                 }
 
@@ -3550,7 +3550,28 @@ impl State {
 
         let mod_key = self.backend.mod_key(&self.niri.config.borrow());
 
-        if !handle.is_grabbed() {
+        if self.niri.screenshot_ui.is_open() {
+            if let Some(output) = under.output.clone() {
+                let geom = self.niri.global_space.output_geometry(&output).unwrap();
+                let mut point = (pos - geom.loc.to_f64())
+                    .to_physical(output.current_scale().fractional_scale())
+                    .to_i32_round();
+
+                let size = output.current_mode().unwrap().size;
+                let transform = output.current_transform();
+                let size = transform.transform_size(size);
+                point.x = min(size.w - 1, point.x);
+                point.y = min(size.h - 1, point.y);
+
+                if self
+                    .niri
+                    .screenshot_ui
+                    .pointer_down(output, point, Some(slot))
+                {
+                    self.niri.queue_redraw_all();
+                }
+            }
+        } else if !handle.is_grabbed() {
             let mods = self.niri.seat.get_keyboard().unwrap().modifier_state();
             let mods = modifiers_from_state(mods);
             let mod_down = mods.contains(mod_key.to_modifiers());
@@ -3638,11 +3659,17 @@ impl State {
         let Some(handle) = self.niri.seat.get_touch() else {
             return;
         };
+        let slot = evt.slot();
+
+        if self.niri.screenshot_ui.pointer_up(Some(slot)) {
+            self.niri.queue_redraw_all();
+        }
+
         let serial = SERIAL_COUNTER.next_serial();
         handle.up(
             self,
             &UpEvent {
-                slot: evt.slot(),
+                slot,
                 serial,
                 time: evt.time_msec(),
             },
@@ -3655,12 +3682,30 @@ impl State {
         let Some(pos) = self.compute_touch_location(&evt) else {
             return;
         };
+        let slot = evt.slot();
+
+        if let Some(output) = self.niri.screenshot_ui.selection_output().cloned() {
+            let geom = self.niri.global_space.output_geometry(&output).unwrap();
+            let mut point = (pos - geom.loc.to_f64())
+                .to_physical(output.current_scale().fractional_scale())
+                .to_i32_round::<i32>();
+
+            let size = output.current_mode().unwrap().size;
+            let transform = output.current_transform();
+            let size = transform.transform_size(size);
+            point.x = point.x.clamp(0, size.w - 1);
+            point.y = point.y.clamp(0, size.h - 1);
+
+            self.niri.screenshot_ui.pointer_motion(point, Some(slot));
+            self.niri.queue_redraw(&output);
+        }
+
         let under = self.niri.contents_under(pos);
         handle.motion(
             self,
             under.surface,
             &TouchMotionEvent {
-                slot: evt.slot(),
+                slot,
                 location: pos,
                 time: evt.time_msec(),
             },
diff --git a/src/niri.rs b/src/niri.rs
index 43a582b6..54a3dd58 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1712,6 +1712,9 @@ impl State {
             SERIAL_COUNTER.next_serial(),
             get_monotonic_time().as_millis() as u32,
         );
+        if let Some(touch) = self.niri.seat.get_touch() {
+            touch.unset_grab(self);
+        }
 
         self.backend.with_primary_renderer(|renderer| {
             self.niri
diff --git a/src/ui/screenshot_ui.rs b/src/ui/screenshot_ui.rs
index 2de3ddbc..14e1e84a 100644
--- a/src/ui/screenshot_ui.rs
+++ b/src/ui/screenshot_ui.rs
@@ -11,6 +11,7 @@ use niri_ipc::SizeChange;
 use pango::{Alignment, FontDescription};
 use pangocairo::cairo::{self, ImageSurface};
 use smithay::backend::allocator::Fourcc;
+use smithay::backend::input::TouchSlot;
 use smithay::backend::renderer::element::utils::{Relocate, RelocateRenderElement};
 use smithay::backend::renderer::element::Kind;
 use smithay::backend::renderer::gles::{GlesRenderer, GlesTexture};
@@ -56,6 +57,7 @@ pub enum ScreenshotUi {
         selection: (Output, Point<i32, Physical>, Point<i32, Physical>),
         output_data: HashMap<Output, OutputData>,
         mouse_down: bool,
+        touch_slot: Option<TouchSlot>,
         show_pointer: bool,
         open_anim: Animation,
         clock: Clock,
@@ -193,6 +195,7 @@ impl ScreenshotUi {
             selection,
             output_data,
             mouse_down: false,
+            touch_slot: None,
             show_pointer,
             open_anim,
             clock: clock.clone(),
@@ -659,25 +662,36 @@ impl ScreenshotUi {
     }
 
     /// The pointer has moved to `point` relative to the current selection output.
-    pub fn pointer_motion(&mut self, point: Point<i32, Physical>) {
+    pub fn pointer_motion(&mut self, point: Point<i32, Physical>, slot: Option<TouchSlot>) {
         let Self::Open {
             selection,
             mouse_down: true,
+            touch_slot,
             ..
         } = self
         else {
             return;
         };
 
+        if *touch_slot != slot {
+            return;
+        }
+
         selection.2 = point;
         self.update_buffers();
     }
 
-    pub fn pointer_down(&mut self, output: Output, point: Point<i32, Physical>) -> bool {
+    pub fn pointer_down(
+        &mut self,
+        output: Output,
+        point: Point<i32, Physical>,
+        slot: Option<TouchSlot>,
+    ) -> bool {
         let Self::Open {
             selection,
             output_data,
             mouse_down,
+            touch_slot,
             ..
         } = self
         else {
@@ -694,17 +708,19 @@ impl ScreenshotUi {
 
         *mouse_down = true;
         *selection = (output, point, point);
+        *touch_slot = slot;
 
         self.update_buffers();
 
         true
     }
 
-    pub fn pointer_up(&mut self) -> bool {
+    pub fn pointer_up(&mut self, slot: Option<TouchSlot>) -> bool {
         let Self::Open {
             selection,
             output_data,
             mouse_down,
+            touch_slot,
             ..
         } = self
         else {
@@ -715,7 +731,12 @@ impl ScreenshotUi {
             return false;
         }
 
+        if *touch_slot != slot {
+            return false;
+        }
+
         *mouse_down = false;
+        *touch_slot = None;
 
         // Check if the resulting selection is zero-sized, and try to come up with a small
         // default rectangle.