Add find_root_shell_surface() that goes through popups

2 files changed, 24 insertions, 9 deletions

diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
index 9ffe5bab..8c1374d3 100644
--- a/src/handlers/mod.rs
+++ b/src/handlers/mod.rs
@@ -343,13 +343,12 @@ impl ClientDndGrabHandler for State {
         // example. On successful drop, additionally activate the target window.
         let mut activate_output = true;
         if let Some(target) = validated.then_some(target).flatten() {
-            if let Some(root) = self.niri.root_surface.get(&target) {
-                if let Some((mapped, _)) = self.niri.layout.find_window_and_output(root) {
-                    let window = mapped.window.clone();
-                    self.niri.layout.activate_window(&window);
-                    self.niri.layer_shell_on_demand_focus = None;
-                    activate_output = false;
-                }
+            let root = self.niri.find_root_shell_surface(&target);
+            if let Some((mapped, _)) = self.niri.layout.find_window_and_output(&root) {
+                let window = mapped.window.clone();
+                self.niri.layout.activate_window(&window);
+                self.niri.layer_shell_on_demand_focus = None;
+                activate_output = false;
             }
         }
 
diff --git a/src/niri.rs b/src/niri.rs
index 687f8539..dbbc9ff2 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -41,8 +41,8 @@ use smithay::desktop::utils::{
     under_from_surface_tree, update_surface_primary_scanout_output, OutputPresentationFeedback,
 };
 use smithay::desktop::{
-    layer_map_for_output, LayerMap, LayerSurface, PopupGrab, PopupManager, PopupUngrabStrategy,
-    Space, Window, WindowSurfaceType,
+    find_popup_root_surface, layer_map_for_output, LayerMap, LayerSurface, PopupGrab, PopupManager,
+    PopupUngrabStrategy, Space, Window, WindowSurfaceType,
 };
 use smithay::input::keyboard::Layout as KeyboardLayout;
 use smithay::input::pointer::{CursorIcon, CursorImageAttributes, CursorImageStatus, MotionEvent};
@@ -4907,6 +4907,22 @@ impl Niri {
         }
     }
 
+    /// Tries to find and return the root shell surface for a given surface.
+    ///
+    /// I.e. for popups, this function will try to find the parent toplevel or layer surface. For
+    /// regular subsurfaces, it will find the root surface.
+    pub fn find_root_shell_surface(&self, surface: &WlSurface) -> WlSurface {
+        let Some(root) = self.root_surface.get(surface) else {
+            return surface.clone();
+        };
+
+        if let Some(popup) = self.popups.find_popup(root) {
+            return find_popup_root_surface(&popup).unwrap_or_else(|_| root.clone());
+        }
+
+        root.clone()
+    }
+
     #[cfg(feature = "dbus")]
     pub fn on_ipc_outputs_changed(&self) {
         let _span = tracy_client::span!("Niri::on_ipc_outputs_changed");