Clean up stuff

src/api.rs

@@ -725,7 +725,7 @@ impl tag_service_server::TagService for TagService {
             let Some(tag) = tag_id.tag(state) else { return };
             let Some(output) = tag.output(state) else { return };
 
-            output.with_state(|state| {
+            output.with_state_mut(|state| {
                 for op_tag in state.tags.iter_mut() {
                     op_tag.set_active(false);
                 }
@@ -792,7 +792,7 @@ impl tag_service_server::TagService for TagService {
                 .outputs()
                 .find(|output| output.name() == output_name.0)
             {
-                output.with_state(|state| {
+                output.with_state_mut(|state| {
                     state.tags.extend(new_tags.clone());
                     tracing::debug!("tags added, are now {:?}", state.tags);
                 });
@@ -800,7 +800,7 @@ impl tag_service_server::TagService for TagService {
 
             for tag in new_tags {
                 for window in state.windows.iter() {
-                    window.with_state(|state| {
+                    window.with_state_mut(|state| {
                         for win_tag in state.tags.iter_mut() {
                             if win_tag.id() == tag.id() {
                                 *win_tag = tag.clone();
@@ -826,7 +826,7 @@ impl tag_service_server::TagService for TagService {
 
             for output in state.space.outputs().cloned().collect::<Vec<_>>() {
                 // TODO: seriously, convert state.tags into a hashset
-                output.with_state(|state| {
+                output.with_state_mut(|state| {
                     for tag_to_remove in tags_to_remove.iter() {
                         state.tags.retain(|tag| tag != tag_to_remove);
                     }
@@ -1061,8 +1061,7 @@ impl output_service_server::OutputService for OutputService {
             let y = output.as_ref().map(|output| output.current_location().y);
 
             let focused = state
-                .output_focus_stack
+                .focused_output()
-                .current_focus()
                 .and_then(|foc_op| output.as_ref().map(|op| op == foc_op));
 
             let tag_ids = output
@@ -1176,7 +1175,7 @@ impl window_service_server::WindowService for WindowService {
 
             let rect = Rectangle::from_loc_and_size(window_loc, window_size);
             // window.change_geometry(rect);
-            window.with_state(|state| {
+            window.with_state_mut(|state| {
                 use crate::window::window_state::FloatingOrTiled;
                 state.floating_or_tiled = match state.floating_or_tiled {
                     FloatingOrTiled::Floating(_) => FloatingOrTiled::Floating(rect),
@@ -1374,7 +1373,7 @@ impl window_service_server::WindowService for WindowService {
             match set_or_toggle {
                 SetOrToggle::Set => {
                     window.set_activate(true);
-                    output.with_state(|state| state.focus_stack.set_focus(window.clone()));
+                    output.with_state_mut(|state| state.focus_stack.set_focus(window.clone()));
                     state.output_focus_stack.set_focus(output.clone());
                     if let Some(keyboard) = state.seat.get_keyboard() {
                         keyboard.set_focus(
@@ -1385,24 +1384,22 @@ impl window_service_server::WindowService for WindowService {
                     }
                 }
                 SetOrToggle::Unset => {
-                    if output.with_state(|state| state.focus_stack.current_focus() == Some(&window))
+                    if state.focused_window(&output) == Some(window) {
-                    {
+                        output.with_state_mut(|state| state.focus_stack.unset_focus());
-                        output.with_state(|state| state.focus_stack.unset_focus());
                         if let Some(keyboard) = state.seat.get_keyboard() {
                             keyboard.set_focus(state, None, SERIAL_COUNTER.next_serial());
                         }
                     }
                 }
                 SetOrToggle::Toggle => {
-                    if output.with_state(|state| state.focus_stack.current_focus() == Some(&window))
+                    if state.focused_window(&output).as_ref() == Some(&window) {
-                    {
+                        output.with_state_mut(|state| state.focus_stack.unset_focus());
-                        output.with_state(|state| state.focus_stack.unset_focus());
                         if let Some(keyboard) = state.seat.get_keyboard() {
                             keyboard.set_focus(state, None, SERIAL_COUNTER.next_serial());
                         }
                     } else {
                         window.set_activate(true);
-                        output.with_state(|state| state.focus_stack.set_focus(window.clone()));
+                        output.with_state_mut(|state| state.focus_stack.set_focus(window.clone()));
                         state.output_focus_stack.set_focus(output.clone());
                         if let Some(keyboard) = state.seat.get_keyboard() {
                             keyboard.set_focus(
@@ -1449,7 +1446,7 @@ impl window_service_server::WindowService for WindowService {
         run_unary_no_response(&self.sender, move |state| {
             let Some(window) = window_id.window(state) else { return };
             let Some(tag) = tag_id.tag(state) else { return };
-            window.with_state(|state| {
+            window.with_state_mut(|state| {
                 state.tags = vec![tag.clone()];
             });
             let Some(output) = tag.output(state) else { return };
@@ -1486,14 +1483,14 @@ impl window_service_server::WindowService for WindowService {
 
             // TODO: turn state.tags into a hashset
             match set_or_toggle {
-                SetOrToggle::Set => window.with_state(|state| {
+                SetOrToggle::Set => window.with_state_mut(|state| {
                     state.tags.retain(|tg| tg != &tag);
                     state.tags.push(tag.clone());
                 }),
-                SetOrToggle::Unset => window.with_state(|state| {
+                SetOrToggle::Unset => window.with_state_mut(|state| {
                     state.tags.retain(|tg| tg != &tag);
                 }),
-                SetOrToggle::Toggle => window.with_state(|state| {
+                SetOrToggle::Toggle => window.with_state_mut(|state| {
                     if !state.tags.contains(&tag) {
                         state.tags.push(tag.clone());
                     } else {
@@ -1684,8 +1681,7 @@ impl window_service_server::WindowService for WindowService {
 
             let focused = window.as_ref().and_then(|win| {
                 state
-                    .output_focus_stack
+                    .focused_output()
-                    .current_focus()
                     .and_then(|output| state.focused_window(output))
                     .map(|foc_win| win == &foc_win)
             });

src/backend/udev.rs

@@ -972,7 +972,7 @@ impl State {
             output.change_current_state(None, None, None, Some(*loc));
             self.space.map_output(&output, *loc);
 
-            output.with_state(|state| state.tags = tags.clone());
+            output.with_state_mut(|state| state.tags = tags.clone());
         } else {
             self.signal_state.output_connect.signal(|buffer| {
                 buffer.push_back(OutputConnectResponse {

src/backend/winit.rs

@@ -318,7 +318,7 @@ impl State {
 
                 // Send frames to the cursor surface so it updates correctly
                 if let CursorImageStatus::Surface(surf) = &self.cursor_status {
-                    if let Some(op) = self.output_focus_stack.current_focus() {
+                    if let Some(op) = self.focused_output() {
                         send_frames_surface_tree(surf, op, time, Some(Duration::ZERO), |_, _| None);
                     }
                 }

src/config.rs

@@ -299,7 +299,7 @@ impl State {
 
         debug!("Clearing tags");
         for output in self.space.outputs() {
-            output.with_state(|state| state.tags.clear());
+            output.with_state_mut(|state| state.tags.clear());
         }
 
         TagId::reset();

src/focus.rs

@@ -11,33 +11,23 @@ pub mod keyboard;
 pub mod pointer;
 
 impl State {
-    /// Get the currently focused window on `output`
+    /// Get the currently focused window on `output`.
-    /// that isn't an override redirect window, if any.
+    ///
+    /// This returns the topmost window on the keyboard focus stack that is on an active tag.
     pub fn focused_window(&self, output: &Output) -> Option<WindowElement> {
         // TODO: see if the below is necessary
         // output.with_state(|state| state.focus_stack.stack.retain(|win| win.alive()));
 
-        let windows = output.with_state(|state| {
+        output.with_state(|state| {
             state
                 .focus_stack
                 .stack
                 .iter()
                 .rev()
-                .filter(|win| {
+                .filter(|win| win.is_on_active_tag())
-                    let win_tags = win.with_state(|state| state.tags.clone());
-                    let output_tags = state.focused_tags().cloned().collect::<Vec<_>>();
-
-                    win_tags
-                        .iter()
-                        .any(|win_tag| output_tags.iter().any(|op_tag| win_tag == op_tag))
-                })
-                .cloned()
-                .collect::<Vec<_>>()
-        });
-
-        windows
-            .into_iter()
                 .find(|win| !win.is_x11_override_redirect())
+                .cloned()
+        })
     }
 
     /// Update the keyboard focus.
@@ -64,41 +54,51 @@ impl State {
             self.space.raise_element(win, false);
         }
     }
+
+    /// Get the currently focused output, or the first mapped output if there is none, or None.
+    pub fn focused_output(&self) -> Option<&Output> {
+        self.output_focus_stack
+            .stack
+            .last()
+            .or_else(|| self.space.outputs().next())
+    }
 }
 
-/// A vector of windows, with the last one being the one in focus and the first
+#[derive(Debug, Clone, Default)]
-/// being the one at the bottom of the focus stack.
+pub struct OutputFocusStack {
-#[derive(Debug)]
+    stack: Vec<Output>,
-pub struct FocusStack<T> {
+}
-    pub stack: Vec<T>,
+
+impl OutputFocusStack {
+    // Set the new focused output.
+    pub fn set_focus(&mut self, output: Output) {
+        self.stack.retain(|op| op != &output);
+        self.stack.push(output);
+    }
+}
+
+/// A stack of windows, with the top one being the one in focus.
+#[derive(Debug, Default)]
+pub struct WindowKeyboardFocusStack {
+    pub stack: Vec<WindowElement>,
     focused: bool,
 }
 
-impl<T> Default for FocusStack<T> {
+impl WindowKeyboardFocusStack {
-    fn default() -> Self {
+    /// Set `window` to be focused.
-        Self {
-            stack: Default::default(),
-            focused: Default::default(),
-        }
-    }
-}
-
-impl<T: PartialEq> FocusStack<T> {
-    /// Set `focus` to be focused.
     ///
     /// If it's already in the stack, it will be removed then pushed.
     /// If it isn't, it will just be pushed.
-    pub fn set_focus(&mut self, focus: T) {
+    pub fn set_focus(&mut self, window: WindowElement) {
-        self.stack.retain(|foc| foc != &focus);
+        self.stack.retain(|win| win != &window);
-        self.stack.push(focus);
+        self.stack.push(window);
         self.focused = true;
     }
 
+    /// Unset the focus by marking this stack as unfocused.
+    ///
+    /// This will cause [`Self::current_focus`] to return `None`.
     pub fn unset_focus(&mut self) {
         self.focused = false;
     }
-
-    pub fn current_focus(&self) -> Option<&T> {
-        self.focused.then(|| self.stack.last())?
-    }
 }

src/grab/move_grab.rs

@@ -131,7 +131,7 @@ impl PointerGrab<State> for MoveSurfaceGrab {
                 .expect("window wasn't mapped")
                 .size;
 
-            self.window.with_state(|state| {
+            self.window.with_state_mut(|state| {
                 if state.floating_or_tiled.is_floating() {
                     state.floating_or_tiled =
                         FloatingOrTiled::Floating(Rectangle::from_loc_and_size(new_loc, size));

src/grab/resize_grab.rs

@@ -64,7 +64,7 @@ impl ResizeSurfaceGrab {
         initial_window_rect: Rectangle<i32, Logical>,
         button_used: u32,
     ) -> Option<Self> {
-        window.wl_surface()?.with_state(|state| {
+        window.wl_surface()?.with_state_mut(|state| {
             state.resize_state = ResizeSurfaceState::Resizing {
                 edges,
                 initial_window_rect,
@@ -215,7 +215,7 @@ impl PointerGrab<State> for ResizeSurfaceGrab {
 
                     toplevel.send_pending_configure();
 
-                    toplevel.wl_surface().with_state(|state| {
+                    toplevel.wl_surface().with_state_mut(|state| {
                         // TODO: validate resize state
                         state.resize_state = ResizeSurfaceState::WaitingForLastCommit {
                             edges: self.edges,
@@ -228,7 +228,7 @@ impl PointerGrab<State> for ResizeSurfaceGrab {
                         return;
                     }
                     let Some(surface) = surface.wl_surface() else { return };
-                    surface.with_state(|state| {
+                    surface.with_state_mut(|state| {
                         state.resize_state = ResizeSurfaceState::WaitingForLastCommit {
                             edges: self.edges,
                             initial_window_rect: self.initial_window_rect,
@@ -363,7 +363,7 @@ pub fn handle_commit(state: &mut State, surface: &WlSurface) -> Option<()> {
     let mut window_loc = state.space.element_location(&window)?;
     let geometry = window.geometry();
 
-    let new_loc: Point<Option<i32>, Logical> = surface.with_state(|state| {
+    let new_loc: Point<Option<i32>, Logical> = surface.with_state_mut(|state| {
         state
             .resize_state
             .commit()
@@ -406,7 +406,7 @@ pub fn handle_commit(state: &mut State, surface: &WlSurface) -> Option<()> {
         .expect("called element_geometry on unmapped window")
         .size;
 
-    window.with_state(|state| {
+    window.with_state_mut(|state| {
         if state.floating_or_tiled.is_floating() {
             state.floating_or_tiled =
                 FloatingOrTiled::Floating(Rectangle::from_loc_and_size(window_loc, size));

src/handlers.rs

@@ -120,7 +120,7 @@ impl CompositorHandler for State {
         if !compositor::is_sync_subsurface(surface) {
             if let Some(window) = self.window_for_surface(&root) {
                 window.on_commit();
-                if let Some(loc) = window.with_state(|state| state.target_loc.take()) {
+                if let Some(loc) = window.with_state_mut(|state| state.target_loc.take()) {
                     self.space.map_element(window.clone(), loc, false);
                 }
             }
@@ -145,21 +145,21 @@ impl CompositorHandler for State {
                 self.windows.push(new_window.clone());
                 self.z_index_stack.set_focus(new_window.clone());
 
-                if let (Some(output), _) | (None, Some(output)) = (
+                if let Some(output) = self.focused_output() {
-                    self.output_focus_stack.current_focus(),
-                    self.space.outputs().next(),
-                ) {
                     tracing::debug!("Placing toplevel");
                     new_window.place_on_output(output);
-                    output.with_state(|state| state.focus_stack.set_focus(new_window.clone()));
+                    output.with_state_mut(|state| state.focus_stack.set_focus(new_window.clone()));
                 }
 
+                // FIXME: I'm mapping way offscreen here then sending a frame to prevent a window from
+                // |      mapping with its default geometry then immediately resizing
+                // |      because I don't set a target geometry before the initial configure.
                 self.space
                     .map_element(new_window.clone(), (1000000, 0), true);
 
                 self.apply_window_rules(&new_window);
 
-                if let Some(focused_output) = self.output_focus_stack.current_focus().cloned() {
+                if let Some(focused_output) = self.focused_output().cloned() {
                     self.update_windows(&focused_output);
                     new_window.send_frame(
                         &focused_output,

src/handlers/xdg_shell.rs

@@ -63,7 +63,7 @@ impl XdgShellHandler for State {
         });
 
         for output in self.space.outputs() {
-            output.with_state(|state| {
+            output.with_state_mut(|state| {
                 state.focus_stack.stack.retain(|window| {
                     window
                         .wl_surface()
@@ -103,9 +103,7 @@ impl XdgShellHandler for State {
     fn new_popup(&mut self, surface: PopupSurface, mut positioner: PositionerState) {
         tracing::debug!(?positioner.constraint_adjustment, ?positioner.gravity);
         let output_rect = self
-            .output_focus_stack
+            .focused_output()
-            .current_focus()
-            .or_else(|| self.space.outputs().next())
             .and_then(|op| self.space.output_geometry(op));
 
         /// Horizontal direction

src/handlers/xwayland.rs

@@ -51,15 +51,13 @@ impl XwmHandler for State {
             .expect("called element_bbox on an unmapped window");
 
         let output_size = self
-            .output_focus_stack
+            .focused_output()
-            .current_focus()
             .and_then(|op| self.space.output_geometry(op))
             .map(|geo| geo.size)
             .unwrap_or((2, 2).into());
 
         let output_loc = self
-            .output_focus_stack
+            .focused_output()
-            .current_focus()
             .map(|op| op.current_location())
             .unwrap_or((0, 0).into());
 
@@ -87,15 +85,12 @@ impl XwmHandler for State {
             .expect("failed to configure x11 window");
         // TODO: ssd
 
-        if let (Some(output), _) | (None, Some(output)) = (
+        if let Some(output) = self.focused_output() {
-            self.output_focus_stack.current_focus(),
-            self.space.outputs().next(),
-        ) {
             window.place_on_output(output);
         }
 
         if should_float(surface) {
-            window.with_state(|state| {
+            window.with_state_mut(|state| {
                 state.floating_or_tiled = FloatingOrTiled::Floating(bbox);
             });
         }
@@ -107,7 +102,7 @@ impl XwmHandler for State {
         self.apply_window_rules(&window);
 
         if let Some(output) = window.output(self) {
-            output.with_state(|state| state.focus_stack.set_focus(window.clone()));
+            output.with_state_mut(|state| state.focus_stack.set_focus(window.clone()));
             self.update_windows(&output);
         }
 
@@ -136,14 +131,11 @@ impl XwmHandler for State {
         self.windows.push(window.clone());
         self.z_index_stack.set_focus(window.clone());
 
-        if let (Some(output), _) | (None, Some(output)) = (
+        if let Some(output) = self.focused_output() {
-            self.output_focus_stack.current_focus(),
-            self.space.outputs().next(),
-        ) {
             window.place_on_output(output);
             // FIXME: setting focus here may possibly muck things up
             // |      or maybe they won't idk
-            output.with_state(|state| state.focus_stack.set_focus(window.clone()))
+            output.with_state_mut(|state| state.focus_stack.set_focus(window.clone()))
         }
 
         self.space.map_element(window, loc, true);
@@ -151,7 +143,7 @@ impl XwmHandler for State {
 
     fn unmapped_window(&mut self, _xwm: XwmId, surface: X11Surface) {
         for output in self.space.outputs() {
-            output.with_state(|state| {
+            output.with_state_mut(|state| {
                 state.focus_stack.stack.retain(|win| {
                     win.wl_surface()
                         .is_some_and(|surf| Some(surf) != surface.wl_surface())
@@ -206,7 +198,7 @@ impl XwmHandler for State {
 
     fn destroyed_window(&mut self, _xwm: XwmId, surface: X11Surface) {
         for output in self.space.outputs() {
-            output.with_state(|state| {
+            output.with_state_mut(|state| {
                 state.focus_stack.stack.retain(|win| {
                     win.wl_surface()
                         .is_some_and(|surf| Some(surf) != surface.wl_surface())

src/input.rs

@@ -218,7 +218,7 @@ impl State {
             .space
             .elements()
             .rev()
-            .filter(|win| win.is_on_active_tag(self.space.outputs()))
+            .filter(|win| win.is_on_active_tag())
             .find_map(|win| {
                 let loc = self
                     .space
@@ -378,7 +378,7 @@ impl State {
                     self.space.raise_element(&window, true);
                     self.z_index_stack.set_focus(window.clone());
                     if let Some(output) = window.output(self) {
-                        output.with_state(|state| state.focus_stack.set_focus(window.clone()));
+                        output.with_state_mut(|state| state.focus_stack.set_focus(window.clone()));
                     }
                 }
 
@@ -395,8 +395,8 @@ impl State {
                     }
                 }
             } else {
-                if let Some(focused_op) = self.output_focus_stack.current_focus() {
+                if let Some(focused_op) = self.focused_output() {
-                    focused_op.with_state(|state| {
+                    focused_op.with_state_mut(|state| {
                         state.focus_stack.unset_focus();
                         for window in state.focus_stack.stack.iter() {
                             window.set_activate(false);
@@ -572,7 +572,7 @@ impl State {
 
             pointer.frame(self);
 
-            if let Some(output) = self.output_focus_stack.current_focus().cloned() {
+            if let Some(output) = self.focused_output().cloned() {
                 self.schedule_render(&output);
             }
         }

src/layout.rs

@@ -130,14 +130,14 @@ impl State {
                         if pending {
                             pending_wins.push((win.clone(), toplevel.send_configure()))
                         } else {
-                            let loc = win.with_state(|state| state.target_loc.take());
+                            let loc = win.with_state_mut(|state| state.target_loc.take());
                             if let Some(loc) = loc {
                                 non_pending_wins.push((loc, win.clone()));
                             }
                         }
                     }
                     WindowSurface::X11(_) => {
-                        let loc = win.with_state(|state| state.target_loc.take());
+                        let loc = win.with_state_mut(|state| state.target_loc.take());
                         if let Some(loc) = loc {
                             self.space.map_element(win.clone(), loc, false);
                         }

src/output.rs

@@ -5,10 +5,9 @@ use std::cell::RefCell;
 use smithay::output::Output;
 
 use crate::{
-    focus::FocusStack,
+    focus::WindowKeyboardFocusStack,
     state::{State, WithState},
     tag::Tag,
-    window::WindowElement,
 };
 
 /// A unique identifier for an output.
@@ -33,13 +32,24 @@ impl OutputName {
 #[derive(Default, Debug)]
 pub struct OutputState {
     pub tags: Vec<Tag>,
-    pub focus_stack: FocusStack<WindowElement>,
+    pub focus_stack: WindowKeyboardFocusStack,
 }
 
 impl WithState for Output {
     type State = OutputState;
 
     fn with_state<F, T>(&self, func: F) -> T
+    where
+        F: FnOnce(&Self::State) -> T,
+    {
+        let state = self
+            .user_data()
+            .get_or_insert(RefCell::<Self::State>::default);
+
+        func(&state.borrow())
+    }
+
+    fn with_state_mut<F, T>(&self, func: F) -> T
     where
         F: FnOnce(&mut Self::State) -> T,
     {

src/render.rs

@@ -147,9 +147,7 @@ where
     let elements = windows
         .iter()
         .rev() // rev because I treat the focus stack backwards vs how the renderer orders it
-        .filter(|win| {
+        .filter(|win| win.is_on_active_tag())
-            win.is_on_active_tag(space.outputs())
-        })
         .map(|win| {
             // subtract win.geometry().loc to align decorations correctly
             let loc = (space.element_location(win).unwrap_or((0, 0).into()) - win.geometry().loc - output.current_location())
@@ -162,7 +160,9 @@ where
 
             (win.render_elements::<WaylandSurfaceRenderElement<R>>(renderer, loc, scale, 1.0), elem_geo)
         }).flat_map(|(elems, rect)| {
-            // elems.into_iter().map(OutputRenderElements::from).collect::<Vec<_>>()
+            // We're cropping everything down to its expected size to stop any sussy windows that
+            // don't want to cooperate. Unfortunately this also truncates shadows and other
+            // external decorations.
             match rect {
                 Some(rect) => {
                     elems.into_iter().filter_map(|elem| {

src/state.rs

@@ -1,14 +1,18 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
 use crate::{
-    api::signal::SignalState, backend::Backend, config::Config, cursor::Cursor, focus::FocusStack,
+    api::signal::SignalState,
-    grab::resize_grab::ResizeSurfaceState, window::WindowElement,
+    backend::Backend,
+    config::Config,
+    cursor::Cursor,
+    focus::{OutputFocusStack, WindowKeyboardFocusStack},
+    grab::resize_grab::ResizeSurfaceState,
+    window::WindowElement,
 };
 use anyhow::Context;
 use smithay::{
     desktop::{PopupManager, Space},
     input::{keyboard::XkbConfig, pointer::CursorImageStatus, Seat, SeatState},
-    output::Output,
     reexports::{
         calloop::{generic::Generic, Interest, LoopHandle, LoopSignal, Mode, PostAction},
         wayland_server::{
@@ -72,8 +76,8 @@ pub struct State {
     /// The state of key and mousebinds along with libinput settings
     pub input_state: InputState,
 
-    pub output_focus_stack: FocusStack<Output>,
+    pub output_focus_stack: OutputFocusStack,
-    pub z_index_stack: FocusStack<WindowElement>,
+    pub z_index_stack: WindowKeyboardFocusStack,
 
     pub popup_manager: PopupManager,
 
@@ -250,8 +254,8 @@ impl State {
 
             input_state: InputState::new(),
 
-            output_focus_stack: FocusStack::default(),
+            output_focus_stack: OutputFocusStack::default(),
-            z_index_stack: FocusStack::default(),
+            z_index_stack: WindowKeyboardFocusStack::default(),
 
             config: Config::new(no_config, config_dir),
 
@@ -331,8 +335,16 @@ pub trait WithState {
 
     /// Access data map state.
     ///
-    /// RefCell Safety: This function will panic if called within itself.
+    /// RefCell Safety: This function will panic if called within [`with_state_mut`][Self::with_state_mut].
     fn with_state<F, T>(&self, func: F) -> T
+    where
+        F: FnOnce(&Self::State) -> T;
+
+    /// Access data map state mutably.
+    ///
+    /// RefCell Safety: This function will panic if called within itself or
+    /// [`with_state`][Self::with_state].
+    fn with_state_mut<F, T>(&self, func: F) -> T
     where
         F: FnOnce(&mut Self::State) -> T;
 }
@@ -346,6 +358,19 @@ impl WithState for WlSurface {
     type State = WlSurfaceState;
 
     fn with_state<F, T>(&self, func: F) -> T
+    where
+        F: FnOnce(&Self::State) -> T,
+    {
+        compositor::with_states(self, |states| {
+            let state = states
+                .data_map
+                .get_or_insert(RefCell::<Self::State>::default);
+
+            func(&state.borrow())
+        })
+    }
+
+    fn with_state_mut<F, T>(&self, func: F) -> T
     where
         F: FnOnce(&mut Self::State) -> T,
     {

src/window.rs

@@ -49,7 +49,6 @@ impl WindowElement {
                 });
             }
             WindowSurface::X11(surface) => {
-                // TODO: maybe move this check elsewhere idk
                 if !surface.is_override_redirect() {
                     surface
                         .configure(new_geo)
@@ -57,11 +56,12 @@ impl WindowElement {
                 }
             }
         }
-        self.with_state(|state| {
+        self.with_state_mut(|state| {
             state.target_loc = Some(new_geo.loc);
         });
     }
 
+    /// Get this window's class (app id in Wayland but hey old habits die hard).
     pub fn class(&self) -> Option<String> {
         match self.0.underlying_surface() {
             WindowSurface::Wayland(toplevel) => {
@@ -80,6 +80,7 @@ impl WindowElement {
         }
     }
 
+    /// Get this window's title.
     pub fn title(&self) -> Option<String> {
         match self.0.underlying_surface() {
             WindowSurface::Wayland(toplevel) => {
@@ -109,26 +110,16 @@ impl WindowElement {
 
     /// Returns whether or not this window has an active tag.
     ///
-    /// RefCell Safety: This uses RefCells on both `self` and everything in `outputs`.
+    /// RefCell Safety: This calls `with_state` on `self`.
-    pub fn is_on_active_tag<'a>(&self, outputs: impl IntoIterator<Item = &'a Output>) -> bool {
+    pub fn is_on_active_tag(&self) -> bool {
-        let tags = outputs
+        self.with_state(|state| state.tags.iter().any(|tag| tag.active()))
-            .into_iter()
-            .flat_map(|op| op.with_state(|state| state.focused_tags().cloned().collect::<Vec<_>>()))
-            .collect::<Vec<_>>();
-
-        self.with_state(|state| {
-            state
-                .tags
-                .iter()
-                .any(|tag| tags.iter().any(|tag2| tag == tag2))
-        })
     }
 
     /// Place this window on the given output, giving it the output's focused tags.
     ///
-    /// RefCell Safety: Uses refcells on both the window and the output.
+    /// RefCell Safety: Uses `with_state_mut` on the window and `with_state` on the output
     pub fn place_on_output(&self, output: &Output) {
-        self.with_state(|state| {
+        self.with_state_mut(|state| {
             state.tags = output.with_state(|state| {
                 let output_tags = state.focused_tags().cloned().collect::<Vec<_>>();
                 if !output_tags.is_empty() {
@@ -197,6 +188,17 @@ impl WithState for WindowElement {
     type State = WindowElementState;
 
     fn with_state<F, T>(&self, func: F) -> T
+    where
+        F: FnOnce(&Self::State) -> T,
+    {
+        let state = self
+            .user_data()
+            .get_or_insert(|| RefCell::new(WindowElementState::new()));
+
+        func(&state.borrow())
+    }
+
+    fn with_state_mut<F, T>(&self, func: F) -> T
     where
         F: FnOnce(&mut Self::State) -> T,
     {
@@ -222,6 +224,9 @@ impl State {
             .cloned()
     }
 
+    /// `window_for_surface` but for windows that haven't commited a buffer yet.
+    ///
+    /// Currently only used in `ensure_initial_configure` in [`handlers`][crate::handlers].
     pub fn new_window_for_surface(&self, surface: &WlSurface) -> Option<WindowElement> {
         self.new_windows
             .iter()

src/window/rules.rs

@@ -189,7 +189,7 @@ impl State {
                         let tags = output
                             .with_state(|state| state.focused_tags().cloned().collect::<Vec<_>>());
 
-                        window.with_state(|state| state.tags = tags.clone());
+                        window.with_state_mut(|state| state.tags = tags.clone());
                     }
                 }
 
@@ -199,7 +199,7 @@ impl State {
                         .filter_map(|tag_id| tag_id.tag(self))
                         .collect::<Vec<_>>();
 
-                    window.with_state(|state| state.tags = tags.clone());
+                    window.with_state_mut(|state| state.tags = tags.clone());
                 }
 
                 if let Some(floating_or_tiled) = floating_or_tiled {
@@ -218,7 +218,7 @@ impl State {
                 }
 
                 if let Some(fs_or_max) = fullscreen_or_maximized {
-                    window.with_state(|state| state.fullscreen_or_maximized = *fs_or_max);
+                    window.with_state_mut(|state| state.fullscreen_or_maximized = *fs_or_max);
                 }
 
                 if let Some((w, h)) = size {
@@ -229,7 +229,7 @@ impl State {
                     match window.with_state(|state| state.floating_or_tiled) {
                         window_state::FloatingOrTiled::Floating(mut rect) => {
                             rect.size = (u32::from(*w) as i32, u32::from(*h) as i32).into();
-                            window.with_state(|state| {
+                            window.with_state_mut(|state| {
                                 state.floating_or_tiled =
                                     window_state::FloatingOrTiled::Floating(rect)
                             });
@@ -238,7 +238,7 @@ impl State {
                             if let Some(rect) = rect.as_mut() {
                                 rect.size = (u32::from(*w) as i32, u32::from(*h) as i32).into();
                             }
-                            window.with_state(|state| {
+                            window.with_state_mut(|state| {
                                 state.floating_or_tiled = window_state::FloatingOrTiled::Tiled(rect)
                             });
                         }
@@ -249,7 +249,7 @@ impl State {
                     match window.with_state(|state| state.floating_or_tiled) {
                         window_state::FloatingOrTiled::Floating(mut rect) => {
                             rect.loc = (*loc).into();
-                            window.with_state(|state| {
+                            window.with_state_mut(|state| {
                                 state.floating_or_tiled =
                                     window_state::FloatingOrTiled::Floating(rect)
                             });
@@ -263,7 +263,7 @@ impl State {
                                 Rectangle::from_loc_and_size(Point::from(*loc), size)
                             });
 
-                            window.with_state(|state| {
+                            window.with_state_mut(|state| {
                                 state.floating_or_tiled =
                                     window_state::FloatingOrTiled::Tiled(Some(rect))
                             });

src/window/window_state.rs

@@ -54,7 +54,7 @@ impl WindowElement {
     pub fn toggle_floating(&self) {
         match self.with_state(|state| state.floating_or_tiled) {
             FloatingOrTiled::Floating(current_rect) => {
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.floating_or_tiled = FloatingOrTiled::Tiled(Some(current_rect))
                 });
                 self.set_tiled_states();
@@ -62,7 +62,7 @@ impl WindowElement {
             FloatingOrTiled::Tiled(prev_rect) => {
                 let prev_rect = prev_rect.unwrap_or_else(|| self.geometry());
 
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.floating_or_tiled = FloatingOrTiled::Floating(prev_rect);
                 });
 
@@ -77,7 +77,7 @@ impl WindowElement {
     pub fn toggle_fullscreen(&self) {
         match self.with_state(|state| state.fullscreen_or_maximized) {
             FullscreenOrMaximized::Neither | FullscreenOrMaximized::Maximized => {
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.fullscreen_or_maximized = FullscreenOrMaximized::Fullscreen;
                 });
 
@@ -105,7 +105,7 @@ impl WindowElement {
                 }
             }
             FullscreenOrMaximized::Fullscreen => {
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.fullscreen_or_maximized = FullscreenOrMaximized::Neither;
                 });
 
@@ -124,7 +124,7 @@ impl WindowElement {
     pub fn toggle_maximized(&self) {
         match self.with_state(|state| state.fullscreen_or_maximized) {
             FullscreenOrMaximized::Neither | FullscreenOrMaximized::Fullscreen => {
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.fullscreen_or_maximized = FullscreenOrMaximized::Maximized;
                 });
 
@@ -152,7 +152,7 @@ impl WindowElement {
                 }
             }
             FullscreenOrMaximized::Maximized => {
-                self.with_state(|state| {
+                self.with_state_mut(|state| {
                     state.fullscreen_or_maximized = FullscreenOrMaximized::Neither;
                 });