Pass Un/Mapped to window rule resolution

4 files changed, 42 insertions, 15 deletions

diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 7bd47717..1c200ee5 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -28,7 +28,7 @@ use smithay::{
 
 use crate::layout::workspace::ColumnWidth;
 use crate::niri::{PopupGrabState, State};
-use crate::window::{InitialConfigureState, ResolvedWindowRules, Unmapped};
+use crate::window::{InitialConfigureState, ResolvedWindowRules, Unmapped, WindowRef};
 
 impl XdgShellHandler for State {
     fn xdg_shell_state(&mut self) -> &mut XdgShellState {
@@ -489,6 +489,10 @@ impl State {
             return;
         };
 
+        let config = self.niri.config.borrow();
+        let rules =
+            ResolvedWindowRules::compute(&config.window_rules, WindowRef::Unmapped(unmapped));
+
         let Unmapped { window, state } = unmapped;
 
         let InitialConfigureState::NotConfigured { wants_fullscreen } = state else {
@@ -496,9 +500,6 @@ impl State {
             return;
         };
 
-        let config = self.niri.config.borrow();
-        let rules = ResolvedWindowRules::compute(&config.window_rules, toplevel);
-
         // Pick the target monitor. First, check if we had an output set in the window rules.
         let mon = rules
             .open_on_output
@@ -730,19 +731,23 @@ impl State {
     }
 
     pub fn update_window_rules(&mut self, toplevel: &ToplevelSurface) {
-        let resolve =
-            || ResolvedWindowRules::compute(&self.niri.config.borrow().window_rules, toplevel);
+        let config = self.niri.config.borrow();
+        let window_rules = &config.window_rules;
 
         if let Some(unmapped) = self.niri.unmapped_windows.get_mut(toplevel.wl_surface()) {
+            let new_rules =
+                ResolvedWindowRules::compute(window_rules, WindowRef::Unmapped(unmapped));
             if let InitialConfigureState::Configured { rules, .. } = &mut unmapped.state {
-                *rules = resolve();
+                *rules = new_rules;
             }
         } else if let Some((mapped, output)) = self
             .niri
             .layout
             .find_window_and_output_mut(toplevel.wl_surface())
         {
-            let new_rules = resolve();
+            let new_rules = ResolvedWindowRules::compute(window_rules, WindowRef::Mapped(mapped));
+            drop(config);
+
             if mapped.rules != new_rules {
                 mapped.rules = new_rules;
 
diff --git a/src/niri.rs b/src/niri.rs
index 49592cf6..54c99453 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -115,7 +115,7 @@ use crate::utils::{
     center, center_f64, get_monotonic_time, make_screenshot_path, output_size, write_png_rgba8,
 };
 use crate::wheel_tracker::WheelTracker;
-use crate::window::{InitialConfigureState, Mapped, ResolvedWindowRules, Unmapped};
+use crate::window::{InitialConfigureState, Mapped, ResolvedWindowRules, Unmapped, WindowRef};
 use crate::{animation, niri_render_elements};
 
 const CLEAR_COLOR: [f32; 4] = [0.2, 0.2, 0.2, 1.];
@@ -880,17 +880,17 @@ impl State {
             let window_rules = &self.niri.config.borrow().window_rules;
 
             for unmapped in self.niri.unmapped_windows.values_mut() {
+                let new_rules =
+                    ResolvedWindowRules::compute(window_rules, WindowRef::Unmapped(unmapped));
                 if let InitialConfigureState::Configured { rules, .. } = &mut unmapped.state {
-                    *rules = ResolvedWindowRules::compute(
-                        window_rules,
-                        unmapped.window.toplevel().expect("no X11 support"),
-                    );
+                    *rules = new_rules;
                 }
             }
 
             let mut windows = vec![];
             self.niri.layout.with_windows_mut(|mapped, _| {
-                mapped.rules = ResolvedWindowRules::compute(window_rules, mapped.toplevel());
+                mapped.rules =
+                    ResolvedWindowRules::compute(window_rules, WindowRef::Mapped(mapped));
                 windows.push(mapped.window.clone());
             });
             for win in windows {
diff --git a/src/window/mod.rs b/src/window/mod.rs
index bb6f2151..4d225223 100644
--- a/src/window/mod.rs
+++ b/src/window/mod.rs
@@ -12,6 +12,13 @@ pub use mapped::Mapped;
 pub mod unmapped;
 pub use unmapped::{InitialConfigureState, Unmapped};
 
+/// Reference to a mapped or unmapped window.
+#[derive(Debug, Clone, Copy)]
+pub enum WindowRef<'a> {
+    Unmapped(&'a Unmapped),
+    Mapped(&'a Mapped),
+}
+
 /// Rules fully resolved for a window.
 #[derive(Debug, PartialEq)]
 pub struct ResolvedWindowRules {
@@ -46,6 +53,15 @@ pub struct ResolvedWindowRules {
     pub draw_border_with_background: Option<bool>,
 }
 
+impl<'a> WindowRef<'a> {
+    pub fn toplevel(self) -> &'a ToplevelSurface {
+        match self {
+            WindowRef::Unmapped(unmapped) => unmapped.toplevel(),
+            WindowRef::Mapped(mapped) => mapped.toplevel(),
+        }
+    }
+}
+
 impl ResolvedWindowRules {
     pub const fn empty() -> Self {
         Self {
@@ -61,11 +77,12 @@ impl ResolvedWindowRules {
         }
     }
 
-    pub fn compute(rules: &[WindowRule], toplevel: &ToplevelSurface) -> Self {
+    pub fn compute(rules: &[WindowRule], window: WindowRef) -> Self {
         let _span = tracy_client::span!("ResolvedWindowRules::compute");
 
         let mut resolved = ResolvedWindowRules::empty();
 
+        let toplevel = window.toplevel();
         with_states(toplevel.wl_surface(), |states| {
             let role = states
                 .data_map
diff --git a/src/window/unmapped.rs b/src/window/unmapped.rs
index 51e37e07..9cde4ca7 100644
--- a/src/window/unmapped.rs
+++ b/src/window/unmapped.rs
@@ -1,5 +1,6 @@
 use smithay::desktop::Window;
 use smithay::output::Output;
+use smithay::wayland::shell::xdg::ToplevelSurface;
 
 use super::ResolvedWindowRules;
 use crate::layout::workspace::ColumnWidth;
@@ -58,4 +59,8 @@ impl Unmapped {
     pub fn needs_initial_configure(&self) -> bool {
         matches!(self.state, InitialConfigureState::NotConfigured { .. })
     }
+
+    pub fn toplevel(&self) -> &ToplevelSurface {
+        self.window.toplevel().expect("no X11 support")
+    }
 }