Implement cancelling DnD with Escape

The last Smithay upgrade includes a fix that lets us do this.

src/handlers/mod.rs

@@ -361,7 +361,7 @@ impl DndGrabHandler for State {
         trace!("dnd dropped, target: {target:?}, validated: {validated}");
 
         // End DnD before activating a specific window below so that it takes precedence.
-        self.niri.layout.dnd_end();
+        self.niri.on_maybe_dnd_ended();
 
         // Activate the target output, since that's how Firefox drag-tab-into-new-window works for
         // example. On successful drop, additionally activate the target window.
@@ -383,10 +383,15 @@ impl DndGrabHandler for State {
                 self.niri.layout.focus_output(&output);
             }
         }
+    }
+}
 
-        self.niri.dnd_icon = None;
+impl crate::niri::Niri {
+    pub fn on_maybe_dnd_ended(&mut self) {
+        self.layout.dnd_end();
+        self.dnd_icon = None;
         // FIXME: more granular
-        self.niri.queue_redraw_all();
+        self.queue_redraw_all();
     }
 }
 

src/input/mod.rs

@@ -41,6 +41,8 @@ use smithay::wayland::tablet_manager::{TabletDescriptor, TabletSeatTrait};
 use touch_overview_grab::TouchOverviewGrab;
 
 use self::move_grab::MoveGrab;
+use self::pick_color_grab::PickColorGrab;
+use self::pick_window_grab::PickWindowGrab;
 use self::resize_grab::ResizeGrab;
 use self::spatial_movement_grab::SpatialMovementGrab;
 #[cfg(feature = "dbus")]
@@ -488,19 +490,22 @@ impl State {
                     }
                 }
 
-                if pressed
-                    && raw == Some(Keysym::Escape)
-                    && (this.niri.pick_window.is_some() || this.niri.pick_color.is_some())
-                {
-                    // We window picking state so the pick window grab must be active.
-                    // Unsetting it cancels window picking.
-                    this.niri
-                        .seat
-                        .get_pointer()
-                        .unwrap()
-                        .unset_grab(this, serial, time);
-                    this.niri.suppressed_keys.insert(key_code);
-                    return FilterResult::Intercept(None);
+                if pressed && raw == Some(Keysym::Escape) {
+                    // Cancel certain grabs on Escape.
+                    let pointer = this.niri.seat.get_pointer().unwrap();
+                    if pointer
+                        .with_grab(|_, grab| Self::grab_can_be_cancelled_with_esc(grab))
+                        .unwrap_or(false)
+                    {
+                        pointer.unset_grab(this, serial, time);
+
+                        // If this was a DnD, we won't get DndGrabHandler::dropped(), so we need to
+                        // call the cleanup.
+                        this.niri.on_maybe_dnd_ended();
+
+                        this.niri.suppressed_keys.insert(key_code);
+                        return FilterResult::Intercept(None);
+                    }
                 }
 
                 if let Some(Keysym::space) = raw {
@@ -4296,6 +4301,12 @@ impl State {
             // Null-source DnD: weston-dnd --self-only
             || grab.is::<DnDGrab<Self, WlSurface, WlSurface>>()
     }
+
+    fn grab_can_be_cancelled_with_esc(grab: &(dyn PointerGrab<State> + 'static)) -> bool {
+        let grab = grab.as_any();
+
+        grab.is::<PickWindowGrab>() || grab.is::<PickColorGrab>() || Self::is_dnd_grab(grab)
+    }
 }
 
 /// Check whether the key should be intercepted and mark intercepted