Add with_toplevel_role() util function

src/ipc/server.rs

@@ -19,13 +19,11 @@ use niri_ipc::{Event, KeyboardLayouts, OutputConfigChanged, Reply, Request, Resp
 use smithay::reexports::calloop::generic::Generic;
 use smithay::reexports::calloop::{Interest, LoopHandle, Mode, PostAction};
 use smithay::reexports::rustix::fs::unlink;
-use smithay::wayland::compositor::with_states;
-use smithay::wayland::shell::xdg::XdgToplevelSurfaceData;
 
 use crate::backend::IpcOutputMap;
 use crate::layout::workspace::WorkspaceId;
 use crate::niri::State;
-use crate::utils::version;
+use crate::utils::{version, with_toplevel_role};
 use crate::window::Mapped;
 
 // If an event stream client fails to read events fast enough that we accumulate more than this
@@ -361,22 +359,12 @@ async fn handle_event_stream_client(client: EventStreamClient) -> anyhow::Result
 }
 
 fn make_ipc_window(mapped: &Mapped, workspace_id: Option<WorkspaceId>) -> niri_ipc::Window {
-    let wl_surface = mapped.toplevel().wl_surface();
-    with_states(wl_surface, |states| {
-        let role = states
-            .data_map
-            .get::<XdgToplevelSurfaceData>()
-            .unwrap()
-            .lock()
-            .unwrap();
-
-        niri_ipc::Window {
+    with_toplevel_role(mapped.toplevel(), |role| niri_ipc::Window {
         id: mapped.id().get(),
         title: role.title.clone(),
         app_id: role.app_id.clone(),
         workspace_id: workspace_id.map(|id| id.get()),
         is_focused: mapped.is_focused(),
-        }
     })
 }
 
@@ -559,15 +547,7 @@ impl State {
             let workspace_id = ws_id.map(|id| id.get());
             let mut changed = ipc_win.workspace_id != workspace_id;
 
-            let wl_surface = mapped.toplevel().wl_surface();
-            changed |= with_states(wl_surface, |states| {
-                let role = states
-                    .data_map
-                    .get::<XdgToplevelSurfaceData>()
-                    .unwrap()
-                    .lock()
-                    .unwrap();
-
+            changed |= with_toplevel_role(mapped.toplevel(), |role| {
                 ipc_win.title != role.title || ipc_win.app_id != role.app_id
             });
 

src/niri.rs

@@ -1541,7 +1541,7 @@ impl State {
         to_introspect: &async_channel::Sender<NiriToIntrospect>,
         msg: IntrospectToNiri,
     ) {
-        use smithay::wayland::shell::xdg::XdgToplevelSurfaceData;
+        use crate::utils::with_toplevel_role;
 
         let IntrospectToNiri::GetWindows = msg;
         let _span = tracy_client::span!("GetWindows");
@@ -1549,21 +1549,8 @@ impl State {
         let mut windows = HashMap::new();
 
         self.niri.layout.with_windows(|mapped, _, _| {
-            let wl_surface = mapped
-                .window
-                .toplevel()
-                .expect("no X11 support")
-                .wl_surface();
-
             let id = mapped.id().get();
-            let props = with_states(wl_surface, |states| {
-                let role = states
-                    .data_map
-                    .get::<XdgToplevelSurfaceData>()
-                    .unwrap()
-                    .lock()
-                    .unwrap();
-
+            let props = with_toplevel_role(mapped.toplevel(), |role| {
                 gnome_shell_introspect::WindowProperties {
                     title: role.title.clone().unwrap_or_default(),
                     app_id: role.app_id.clone().unwrap_or_default(),

src/protocols/foreign_toplevel.rs

@@ -11,10 +11,7 @@ use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::{
     Client, DataInit, Dispatch, DisplayHandle, GlobalDispatch, New, Resource,
 };
-use smithay::wayland::compositor::with_states;
-use smithay::wayland::shell::xdg::{
-    ToplevelStateSet, XdgToplevelSurfaceData, XdgToplevelSurfaceRoleAttributes,
-};
+use smithay::wayland::shell::xdg::{ToplevelStateSet, XdgToplevelSurfaceRoleAttributes};
 use wayland_protocols_wlr::foreign_toplevel::v1::server::{
     zwlr_foreign_toplevel_handle_v1, zwlr_foreign_toplevel_manager_v1,
 };
@@ -22,6 +19,7 @@ use zwlr_foreign_toplevel_handle_v1::ZwlrForeignToplevelHandleV1;
 use zwlr_foreign_toplevel_manager_v1::ZwlrForeignToplevelManagerV1;
 
 use crate::niri::State;
+use crate::utils::with_toplevel_role;
 
 const VERSION: u32 = 3;
 
@@ -96,37 +94,23 @@ pub fn refresh(state: &mut State) {
     // the previous window and only then activate the newly focused window.
     let mut focused = None;
     state.niri.layout.with_windows(|mapped, output, _| {
-        let wl_surface = mapped.toplevel().wl_surface();
-
-        with_states(wl_surface, |states| {
-            let role = states
-                .data_map
-                .get::<XdgToplevelSurfaceData>()
-                .unwrap()
-                .lock()
-                .unwrap();
-
+        let toplevel = mapped.toplevel();
+        let wl_surface = toplevel.wl_surface();
+        with_toplevel_role(toplevel, |role| {
             if state.niri.keyboard_focus.surface() == Some(wl_surface) {
                 focused = Some((mapped.window.clone(), output.cloned()));
             } else {
-                refresh_toplevel(protocol_state, wl_surface, &role, output, false);
+                refresh_toplevel(protocol_state, wl_surface, role, output, false);
             }
         });
     });
 
     // Finally, refresh the focused window.
     if let Some((window, output)) = focused {
-        let wl_surface = window.toplevel().expect("no x11 support").wl_surface();
-
-        with_states(wl_surface, |states| {
-            let role = states
-                .data_map
-                .get::<XdgToplevelSurfaceData>()
-                .unwrap()
-                .lock()
-                .unwrap();
-
-            refresh_toplevel(protocol_state, wl_surface, &role, output.as_ref(), true);
+        let toplevel = window.toplevel().expect("no X11 support");
+        let wl_surface = toplevel.wl_surface();
+        with_toplevel_role(toplevel, |role| {
+            refresh_toplevel(protocol_state, wl_surface, role, output.as_ref(), true);
         });
     }
 }

src/utils/mod.rs

@@ -17,8 +17,11 @@ use smithay::reexports::rustix::time::{clock_gettime, ClockId};
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::utils::{Coordinate, Logical, Point, Rectangle, Size, Transform};
-use smithay::wayland::compositor::{send_surface_state, SurfaceData};
+use smithay::wayland::compositor::{send_surface_state, with_states, SurfaceData};
 use smithay::wayland::fractional_scale::with_fractional_scale;
+use smithay::wayland::shell::xdg::{
+    ToplevelSurface, XdgToplevelSurfaceData, XdgToplevelSurfaceRoleAttributes,
+};
 
 pub mod id;
 pub mod scale;
@@ -221,6 +224,22 @@ pub fn output_matches_name(output: &Output, target: &str) -> bool {
     name.matches(target)
 }
 
+pub fn with_toplevel_role<T>(
+    toplevel: &ToplevelSurface,
+    f: impl FnOnce(&mut XdgToplevelSurfaceRoleAttributes) -> T,
+) -> T {
+    with_states(toplevel.wl_surface(), |states| {
+        let mut role = states
+            .data_map
+            .get::<XdgToplevelSurfaceData>()
+            .unwrap()
+            .lock()
+            .unwrap();
+
+        f(&mut role)
+    })
+}
+
 #[cfg(feature = "dbus")]
 pub fn show_screenshot_notification(image_path: Option<PathBuf>) {
     let mut notification = notify_rust::Notification::new();

src/window/mapped.rs

@@ -15,7 +15,7 @@ use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::reexports::wayland_server::Resource as _;
 use smithay::utils::{Logical, Point, Rectangle, Scale, Serial, Size, Transform};
 use smithay::wayland::compositor::{remove_pre_commit_hook, with_states, HookId};
-use smithay::wayland::shell::xdg::{SurfaceCachedState, ToplevelSurface, XdgToplevelSurfaceData};
+use smithay::wayland::shell::xdg::{SurfaceCachedState, ToplevelSurface};
 
 use super::{ResolvedWindowRules, WindowRef};
 use crate::handlers::KdeDecorationsModeState;
@@ -33,7 +33,7 @@ use crate::render_helpers::surface::render_snapshot_from_surface_tree;
 use crate::render_helpers::{BakedBuffer, RenderTarget, SplitElements};
 use crate::utils::id::IdCounter;
 use crate::utils::transaction::Transaction;
-use crate::utils::{send_scale_transform, ResizeEdge};
+use crate::utils::{send_scale_transform, with_toplevel_role, ResizeEdge};
 
 #[derive(Debug)]
 pub struct Mapped {
@@ -631,14 +631,7 @@ impl LayoutElement for Mapped {
         let _span =
             trace_span!("configure_intent", surface = ?self.toplevel().wl_surface().id()).entered();
 
-        with_states(self.toplevel().wl_surface(), |states| {
-            let attributes = states
-                .data_map
-                .get::<XdgToplevelSurfaceData>()
-                .unwrap()
-                .lock()
-                .unwrap();
-
+        with_toplevel_role(self.toplevel(), |attributes| {
             if let Some(server_pending) = &attributes.server_pending {
                 let current_server = attributes.current_server_state();
                 if server_pending != current_server {

src/window/mod.rs

@@ -1,11 +1,9 @@
 use niri_config::{BlockOutFrom, BorderRule, CornerRadius, Match, WindowRule};
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
-use smithay::wayland::compositor::with_states;
-use smithay::wayland::shell::xdg::{
-    ToplevelSurface, XdgToplevelSurfaceData, XdgToplevelSurfaceRoleAttributes,
-};
+use smithay::wayland::shell::xdg::{ToplevelSurface, XdgToplevelSurfaceRoleAttributes};
 
 use crate::layout::workspace::ColumnWidth;
+use crate::utils::with_toplevel_role;
 
 pub mod mapped;
 pub use mapped::Mapped;
@@ -144,15 +142,7 @@ impl ResolvedWindowRules {
 
         let mut resolved = ResolvedWindowRules::empty();
 
-        let toplevel = window.toplevel();
-        with_states(toplevel.wl_surface(), |states| {
-            let mut role = states
-                .data_map
-                .get::<XdgToplevelSurfaceData>()
-                .unwrap()
-                .lock()
-                .unwrap();
-
+        with_toplevel_role(window.toplevel(), |role| {
             // Ensure server_pending like in Smithay's with_pending_state().
             if role.server_pending.is_none() {
                 role.server_pending = Some(role.current_server_state().clone());
@@ -169,7 +159,7 @@ impl ResolvedWindowRules {
                         }
                     }
 
-                    window_matches(window, &role, m)
+                    window_matches(window, role, m)
                 };
 
                 if !(rule.matches.is_empty() || rule.matches.iter().any(matches)) {