Add initial popup unconstraining implementation

Loosely based on the one in cosmic-comp:
https://github.com/pop-os/cosmic-comp/blob/6eb8b82440e3185462d217a3474a968d50122ceb/src/wayland/handlers/xdg_shell/popup.rs

Doesn't update reactive popups yet.

src/handlers/layer_shell.rs

@@ -3,13 +3,16 @@ use smithay::desktop::{layer_map_for_output, LayerSurface, WindowSurfaceType};
 use smithay::output::Output;
 use smithay::reexports::wayland_server::protocol::wl_output::WlOutput;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
+use smithay::utils::Rectangle;
 use smithay::wayland::compositor::{send_surface_state, with_states};
 use smithay::wayland::shell::wlr_layer::{
     Layer, LayerSurface as WlrLayerSurface, LayerSurfaceData, WlrLayerShellHandler,
     WlrLayerShellState,
 };
+use smithay::wayland::shell::xdg::PopupSurface;
 
 use crate::niri::State;
+use crate::popups::unconstrain_popup;
 
 impl WlrLayerShellHandler for State {
     fn shell_state(&mut self) -> &mut WlrLayerShellState {
@@ -52,6 +55,10 @@ impl WlrLayerShellHandler for State {
             self.niri.output_resized(output);
         }
     }
+
+    fn new_popup(&mut self, parent: WlrLayerSurface, popup: PopupSurface) {
+        self.unconstrain_layer_shell_popup(&parent, &popup);
+    }
 }
 delegate_layer_shell!(State);
 
@@ -104,4 +111,28 @@ impl State {
 
         self.niri.output_resized(output);
     }
+
+    fn unconstrain_layer_shell_popup(
+        &self,
+        parent: &WlrLayerSurface,
+        popup: &PopupSurface,
+    ) -> Option<()> {
+        let (output, layer_surface) = self.niri.layout.outputs().find_map(|o| {
+            let map = layer_map_for_output(o);
+            let layer_surface =
+                map.layer_for_surface(parent.wl_surface(), WindowSurfaceType::TOPLEVEL)?;
+            Some((o, layer_surface.clone()))
+        })?;
+
+        let output_geo = self.niri.global_space.output_geometry(output).unwrap();
+        let map = layer_map_for_output(output);
+        let layer_geo = map.layer_geometry(&layer_surface)?;
+
+        // The target geometry for the positioner should be relative to its parent's geometry, so
+        // we will compute that here.
+        let mut target = Rectangle::from_loc_and_size((0, 0), output_geo.size);
+        target.loc -= layer_geo.loc;
+
+        unconstrain_popup(target, popup)
+    }
 }

src/handlers/xdg_shell.rs

@@ -1,21 +1,25 @@
+use std::sync::Mutex;
+
 use smithay::desktop::{find_popup_root_surface, PopupKind, Window};
 use smithay::output::Output;
 use smithay::reexports::wayland_protocols::xdg::decoration::zv1::server::zxdg_toplevel_decoration_v1;
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel::{self, ResizeEdge};
 use smithay::reexports::wayland_server::protocol::wl_output;
 use smithay::reexports::wayland_server::protocol::wl_seat::WlSeat;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
-use smithay::utils::Serial;
-use smithay::wayland::compositor::{send_surface_state, with_states};
+use smithay::utils::{Logical, Point, Rectangle, Serial};
+use smithay::wayland::compositor::{get_role, send_surface_state, with_states};
 use smithay::wayland::shell::kde::decoration::{KdeDecorationHandler, KdeDecorationState};
 use smithay::wayland::shell::xdg::decoration::XdgDecorationHandler;
 use smithay::wayland::shell::xdg::{
-    PopupSurface, PositionerState, ToplevelSurface, XdgPopupSurfaceData, XdgShellHandler,
-    XdgShellState, XdgToplevelSurfaceData,
+    PopupSurface, PositionerState, SurfaceCachedState, ToplevelSurface, XdgPopupSurfaceData,
+    XdgPopupSurfaceRoleAttributes, XdgShellHandler, XdgShellState, XdgToplevelSurfaceData,
+    XDG_POPUP_ROLE,
 };
 use smithay::{delegate_kde_decoration, delegate_xdg_decoration, delegate_xdg_shell};
 
 use crate::niri::State;
+use crate::popups::unconstrain_popup;
 
 impl XdgShellHandler for State {
     fn xdg_shell_state(&mut self) -> &mut XdgShellState {
@@ -48,9 +52,16 @@ impl XdgShellHandler for State {
         assert!(existing.is_none());
     }
 
-    fn new_popup(&mut self, surface: PopupSurface, _positioner: PositionerState) {
-        // FIXME: adjust the geometry so the popup doesn't overflow at least off the top and bottom
-        // screen edges, and ideally off the view size.
+    fn new_popup(&mut self, surface: PopupSurface, positioner: PositionerState) {
+        surface.with_pending_state(|state| {
+            let geometry = positioner.get_geometry();
+            state.geometry = geometry;
+            state.positioner = positioner;
+        });
+
+        // FIXME: for reactive popups, update them upon window resizes, monitor changes, etc.
+        self.unconstrain_xdg_popup(&surface);
+
         if let Err(err) = self.niri.popups.track_popup(PopupKind::Xdg(surface)) {
             warn!("error tracking popup: {err:?}");
         }
@@ -76,13 +87,14 @@ impl XdgShellHandler for State {
         positioner: PositionerState,
         token: u32,
     ) {
-        // FIXME: adjust the geometry so the popup doesn't overflow at least off the top and bottom
-        // screen edges, and ideally off the view size.
         surface.with_pending_state(|state| {
             let geometry = positioner.get_geometry();
             state.geometry = geometry;
             state.positioner = positioner;
         });
+
+        self.unconstrain_xdg_popup(&surface);
+
         surface.send_repositioned(token);
     }
 
@@ -291,4 +303,91 @@ impl State {
             }
         }
     }
+
+    fn unconstrain_xdg_popup(&self, popup: &PopupSurface) -> Option<()> {
+        // Popups with a NULL parent will get repositioned in their respective protocol handlers
+        // (i.e. layer-shell).
+        let toplevel = get_popup_toplevel(popup)?;
+
+        let (window, output) = self.niri.layout.find_window_and_output(&toplevel)?;
+
+        let window_geo = window.geometry();
+        let output_geo = self.niri.global_space.output_geometry(&output).unwrap();
+
+        // The target geometry for the positioner should be relative to its parent's geometry, so
+        // we will compute that here.
+        //
+        // We try to keep regular window popups within the window itself horizontally (since the
+        // window can be scrolled to both edges of the screen), but within the whole monitor's
+        // height.
+        let mut target =
+            Rectangle::from_loc_and_size((0, 0), (window_geo.size.w, output_geo.size.h));
+        target.loc -= get_popup_toplevel_coords(popup);
+        target.loc.y -= self.niri.layout.window_y(&window).unwrap();
+
+        unconstrain_popup(target, popup)
+    }
+}
+
+fn get_popup_toplevel(popup: &PopupSurface) -> Option<WlSurface> {
+    let mut parent = popup.get_parent_surface()?;
+    while get_role(&parent) == Some(XDG_POPUP_ROLE) {
+        parent = with_states(&parent, |states| {
+            states
+                .data_map
+                .get::<Mutex<XdgPopupSurfaceRoleAttributes>>()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .parent
+                .as_ref()
+                .cloned()
+                .unwrap()
+        });
+    }
+    Some(parent)
+}
+
+fn get_popup_toplevel_coords(popup: &PopupSurface) -> Point<i32, Logical> {
+    let mut parent = match popup.get_parent_surface() {
+        Some(parent) => parent,
+        None => return (0, 0).into(),
+    };
+
+    let mut offset = (0, 0).into();
+    while get_role(&parent) == Some(XDG_POPUP_ROLE) {
+        offset += with_states(&parent, |states| {
+            states
+                .data_map
+                .get::<Mutex<XdgPopupSurfaceRoleAttributes>>()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .current
+                .geometry
+                .loc
+        });
+        parent = with_states(&parent, |states| {
+            states
+                .data_map
+                .get::<Mutex<XdgPopupSurfaceRoleAttributes>>()
+                .unwrap()
+                .lock()
+                .unwrap()
+                .parent
+                .as_ref()
+                .cloned()
+                .unwrap()
+        });
+    }
+    offset += with_states(&parent, |states| {
+        states
+            .cached_state
+            .current::<SurfaceCachedState>()
+            .geometry
+            .map(|x| x.loc)
+            .unwrap_or_else(|| (0, 0).into())
+    });
+
+    offset
 }

src/layout.rs

@@ -795,6 +795,33 @@ impl<W: LayoutElement> Layout<W> {
         None
     }
 
+    pub fn window_y(&self, window: &W) -> Option<i32> {
+        match &self.monitor_set {
+            MonitorSet::Normal { monitors, .. } => {
+                for mon in monitors {
+                    for ws in &mon.workspaces {
+                        for col in &ws.columns {
+                            if let Some(idx) = col.windows.iter().position(|w| w == window) {
+                                return Some(col.window_y(idx));
+                            }
+                        }
+                    }
+                }
+            }
+            MonitorSet::NoOutputs { workspaces, .. } => {
+                for ws in workspaces {
+                    for col in &ws.columns {
+                        if let Some(idx) = col.windows.iter().position(|w| w == window) {
+                            return Some(col.window_y(idx));
+                        }
+                    }
+                }
+            }
+        }
+
+        None
+    }
+
     pub fn update_output_size(&mut self, output: &Output) {
         let MonitorSet::Normal { monitors, .. } = &mut self.monitor_set else {
             panic!()

src/main.rs

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