Keep monitor aspect ratio and clamp to monitor for tablets

Before, the full tablet area was used, even if the aspect ratio didn't
match the monitor. Also, the coordinates weren't clamped.

2 files changed, 86 insertions, 20 deletions

diff --git a/src/input.rs b/src/input.rs
index 063d321c..a86045bd 100644
--- a/src/input.rs
+++ b/src/input.rs
@@ -1,3 +1,4 @@
+use std::any::Any;
 use std::collections::HashSet;
 
 use smithay::backend::input::{
@@ -14,7 +15,8 @@ use smithay::input::pointer::{
     GesturePinchBeginEvent, GesturePinchEndEvent, GesturePinchUpdateEvent, GestureSwipeBeginEvent,
     GestureSwipeEndEvent, GestureSwipeUpdateEvent, MotionEvent, RelativeMotionEvent,
 };
-use smithay::utils::SERIAL_COUNTER;
+use smithay::reexports::input;
+use smithay::utils::{Logical, Point, SERIAL_COUNTER};
 use smithay::wayland::tablet_manager::{TabletDescriptor, TabletSeatTrait};
 
 use crate::config::{Action, Binds, LayoutAction, Modifiers};
@@ -28,8 +30,16 @@ pub enum CompositorMod {
     Alt,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct TabletData {
+    pub aspect_ratio: f64,
+}
+
 impl State {
-    pub fn process_input_event<I: InputBackend>(&mut self, event: InputEvent<I>) {
+    pub fn process_input_event<I: InputBackend>(&mut self, event: InputEvent<I>)
+    where
+        I::Device: 'static, // Needed for downcasting.
+    {
         let _span = tracy_client::span!("process_input_event");
 
         // A bit of a hack, but animation end runs some logic (i.e. workspace clean-up) and it
@@ -582,14 +592,10 @@ impl State {
                 pointer.frame(self);
             }
             InputEvent::TabletToolAxis { event, .. } => {
-                let Some(output) = self.niri.output_for_tablet() else {
+                let Some(pos) = self.compute_tablet_position(&event) else {
                     return;
                 };
 
-                let output_geo = self.niri.global_space.output_geometry(output).unwrap();
-
-                let pos = event.position_transformed(output_geo.size) + output_geo.loc.to_f64();
-
                 let serial = SERIAL_COUNTER.next_serial();
 
                 let pointer = self.niri.seat.get_pointer().unwrap();
@@ -671,14 +677,10 @@ impl State {
                 }
             }
             InputEvent::TabletToolProximity { event, .. } => {
-                let Some(output) = self.niri.output_for_tablet() else {
+                let Some(pos) = self.compute_tablet_position(&event) else {
                     return;
                 };
 
-                let output_geo = self.niri.global_space.output_geometry(output).unwrap();
-
-                let pos = event.position_transformed(output_geo.size) + output_geo.loc.to_f64();
-
                 let serial = SERIAL_COUNTER.next_serial();
 
                 let pointer = self.niri.seat.get_pointer().unwrap();
@@ -927,17 +929,75 @@ impl State {
     }
 
     pub fn process_libinput_event(&mut self, event: &mut InputEvent<LibinputInputBackend>) {
-        if let InputEvent::DeviceAdded { device } = event {
-            // According to Mutter code, this setting is specific to touchpads.
-            let is_touchpad = device.config_tap_finger_count() > 0;
-            if is_touchpad {
-                let c = &self.niri.config.borrow().input.touchpad;
-                let _ = device.config_tap_set_enabled(c.tap);
-                let _ = device.config_scroll_set_natural_scroll_enabled(c.natural_scroll);
-                let _ = device.config_accel_set_speed(c.accel_speed);
+        match event {
+            InputEvent::DeviceAdded { device } => {
+                // According to Mutter code, this setting is specific to touchpads.
+                let is_touchpad = device.config_tap_finger_count() > 0;
+                if is_touchpad {
+                    let c = &self.niri.config.borrow().input.touchpad;
+                    let _ = device.config_tap_set_enabled(c.tap);
+                    let _ = device.config_scroll_set_natural_scroll_enabled(c.natural_scroll);
+                    let _ = device.config_accel_set_speed(c.accel_speed);
+                }
+
+                if device.has_capability(smithay::reexports::input::DeviceCapability::TabletTool) {
+                    match device.size() {
+                        Some((w, h)) => {
+                            self.niri.tablets.insert(
+                                device.clone(),
+                                TabletData {
+                                    aspect_ratio: w / h,
+                                },
+                            );
+                        }
+                        None => {
+                            warn!("tablet tool device has no size");
+                        }
+                    }
+                }
+            }
+            InputEvent::DeviceRemoved { device } => {
+                self.niri.tablets.remove(device);
             }
+            _ => (),
         }
     }
+
+    fn compute_tablet_position<I: InputBackend>(
+        &mut self,
+        event: &(impl Event<I> + TabletToolEvent<I>),
+    ) -> Option<Point<f64, Logical>>
+    where
+        I::Device: 'static,
+    {
+        let output = self.niri.output_for_tablet()?;
+        let output_geo = self.niri.global_space.output_geometry(output).unwrap();
+
+        let mut pos = event.position_transformed(output_geo.size);
+        pos.x /= output_geo.size.w as f64;
+        pos.y /= output_geo.size.h as f64;
+
+        let device = event.device();
+        if let Some(device) = (&device as &dyn Any).downcast_ref::<input::Device>() {
+            if let Some(data) = self.niri.tablets.get(device) {
+                // This code does the same thing as mutter with "keep aspect ratio" enabled.
+                let output_aspect_ratio = output_geo.size.w as f64 / output_geo.size.h as f64;
+                let ratio = data.aspect_ratio / output_aspect_ratio;
+
+                if ratio > 1. {
+                    pos.x *= ratio;
+                } else {
+                    pos.y /= ratio;
+                }
+            }
+        };
+
+        pos.x *= output_geo.size.w as f64;
+        pos.y *= output_geo.size.h as f64;
+        pos.x = pos.x.clamp(0.0, output_geo.size.w as f64 - 1.);
+        pos.y = pos.y.clamp(0.0, output_geo.size.h as f64 - 1.);
+        Some(pos + output_geo.loc.to_f64())
+    }
 }
 
 /// Check whether the key should be intercepted and mark intercepted
diff --git a/src/niri.rs b/src/niri.rs
index c9f7487e..66863f15 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -41,6 +41,7 @@ use smithay::reexports::calloop::timer::{TimeoutAction, Timer};
 use smithay::reexports::calloop::{
     self, Idle, Interest, LoopHandle, LoopSignal, Mode, PostAction, RegistrationToken,
 };
+use smithay::reexports::input;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel::WmCapabilities;
 use smithay::reexports::wayland_protocols_misc::server_decoration as _server_decoration;
 use smithay::reexports::wayland_server::backend::{
@@ -86,6 +87,7 @@ use crate::dbus::gnome_shell_screenshot::{NiriToScreenshot, ScreenshotToNiri};
 use crate::dbus::mutter_screen_cast::{self, ScreenCastToNiri};
 use crate::frame_clock::FrameClock;
 use crate::handlers::configure_lock_surface;
+use crate::input::TabletData;
 use crate::layout::{output_size, Layout, MonitorRenderElement};
 use crate::pw_utils::{Cast, PipeWire};
 use crate::screenshot_ui::{ScreenshotUi, ScreenshotUiRenderElement};
@@ -121,6 +123,8 @@ pub struct Niri {
     // When false, we're idling with monitors powered off.
     pub monitors_active: bool,
 
+    pub tablets: HashMap<input::Device, TabletData>,
+
     // Smithay state.
     pub compositor_state: CompositorState,
     pub xdg_shell_state: XdgShellState,
@@ -694,6 +698,8 @@ impl Niri {
             unmapped_windows: HashMap::new(),
             monitors_active: true,
 
+            tablets: HashMap::new(),
+
             compositor_state,
             xdg_shell_state,
             xdg_decoration_state,