Move unmapped check to a pre-commit hook

3 files changed, 57 insertions, 37 deletions

diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index 9473de10..785ad2fe 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -16,7 +16,7 @@ use smithay::wayland::dmabuf::get_dmabuf;
 use smithay::wayland::shm::{ShmHandler, ShmState};
 use smithay::{delegate_compositor, delegate_shm};
 
-use crate::layout::LayoutElement;
+use super::xdg_shell::add_mapped_toplevel_pre_commit_hook;
 use crate::niri::{ClientState, State};
 use crate::window::{InitialConfigureState, Mapped, ResolvedWindowRules, Unmapped};
 
@@ -81,6 +81,13 @@ impl CompositorHandler for State {
     fn commit(&mut self, surface: &WlSurface) {
         let _span = tracy_client::span!("CompositorHandler::commit");
 
+        on_commit_buffer_handler::<Self>(surface);
+        self.backend.early_import(surface);
+
+        if is_sync_subsurface(surface) {
+            return;
+        }
+
         let mut root_surface = surface.clone();
         while let Some(parent) = get_parent(&root_surface) {
             root_surface = parent;
@@ -91,30 +98,6 @@ impl CompositorHandler for State {
             .root_surface
             .insert(surface.clone(), root_surface.clone());
 
-        // Check if this root surface got unmapped to snapshot it before on_commit_buffer_handler()
-        // overwrites the buffer.
-        if surface == &root_surface {
-            let got_unmapped = with_states(surface, |states| {
-                let attrs = states.cached_state.current::<SurfaceAttributes>();
-                matches!(attrs.buffer, Some(BufferAssignment::Removed))
-            });
-
-            if got_unmapped {
-                if let Some((mapped, _)) = self.niri.layout.find_window_and_output(surface) {
-                    self.backend.with_primary_renderer(|renderer| {
-                        mapped.render_and_store_snapshot(renderer);
-                    });
-                }
-            }
-        }
-
-        on_commit_buffer_handler::<Self>(surface);
-        self.backend.early_import(surface);
-
-        if is_sync_subsurface(surface) {
-            return;
-        }
-
         if surface == &root_surface {
             // This is a root surface commit. It might have mapped a previously-unmapped toplevel.
             if let Entry::Occupied(entry) = self.niri.unmapped_windows.entry(surface.clone()) {
@@ -131,6 +114,8 @@ impl CompositorHandler for State {
 
                     window.on_commit();
 
+                    let toplevel = window.toplevel().expect("no X11 support");
+
                     let (rules, width, is_full_width, output) =
                         if let InitialConfigureState::Configured {
                             rules,
@@ -149,9 +134,7 @@ impl CompositorHandler for State {
                             (ResolvedWindowRules::empty(), None, false, None)
                         };
 
-                    let parent = window
-                        .toplevel()
-                        .expect("no x11 support")
+                    let parent = toplevel
                         .parent()
                         .and_then(|parent| self.niri.layout.find_window_and_output(&parent))
                         // Only consider the parent if we configured the window for the same
@@ -165,7 +148,8 @@ impl CompositorHandler for State {
                         })
                         .map(|(mapped, _)| mapped.window.clone());
 
-                    let mapped = Mapped::new(window, rules);
+                    let hook = add_mapped_toplevel_pre_commit_hook(toplevel);
+                    let mapped = Mapped::new(window, rules, hook);
                     let window = mapped.window.clone();
 
                     let output = if let Some(p) = parent {
@@ -218,11 +202,8 @@ impl CompositorHandler for State {
                             false
                         });
 
-                if is_mapped {
-                    // The surface remains mapped; clear any cached render snapshot.
-                    let _ = mapped.take_last_render();
-                } else {
-                    // Must start the close animation before window.on_commit().
+                // Must start the close animation before window.on_commit().
+                if !is_mapped {
                     self.backend.with_primary_renderer(|renderer| {
                         self.niri
                             .layout
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index 53001459..85ce5915 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -12,7 +12,10 @@ 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::{Logical, Rectangle, Serial};
-use smithay::wayland::compositor::{send_surface_state, with_states};
+use smithay::wayland::compositor::{
+    add_pre_commit_hook, send_surface_state, with_states, BufferAssignment, HookId,
+    SurfaceAttributes,
+};
 use smithay::wayland::input_method::InputMethodSeat;
 use smithay::wayland::shell::kde::decoration::{KdeDecorationHandler, KdeDecorationState};
 use smithay::wayland::shell::wlr_layer::Layer;
@@ -27,6 +30,7 @@ use smithay::{
 };
 
 use crate::layout::workspace::ColumnWidth;
+use crate::layout::LayoutElement as _;
 use crate::niri::{PopupGrabState, State};
 use crate::window::{InitialConfigureState, ResolvedWindowRules, Unmapped, WindowRef};
 
@@ -808,3 +812,26 @@ fn unconstrain_with_padding(
     // Could not unconstrain into the padded target, so resort to the regular one.
     positioner.get_unconstrained_geometry(target)
 }
+
+pub fn add_mapped_toplevel_pre_commit_hook(toplevel: &ToplevelSurface) -> HookId {
+    add_pre_commit_hook::<State, _>(toplevel.wl_surface(), move |state, _dh, surface| {
+        let Some((mapped, _)) = state.niri.layout.find_window_and_output_mut(surface) else {
+            error!("pre-commit hook for mapped surfaces must be removed upon unmapping");
+            return;
+        };
+
+        let got_unmapped = with_states(surface, |states| {
+            let attrs = states.cached_state.current::<SurfaceAttributes>();
+            matches!(attrs.buffer, Some(BufferAssignment::Removed))
+        });
+
+        if got_unmapped {
+            state.backend.with_primary_renderer(|renderer| {
+                mapped.render_and_store_snapshot(renderer);
+            });
+        } else {
+            // The toplevel remains mapped; clear any cached render snapshot.
+            let _ = mapped.take_last_render();
+        }
+    })
+}
diff --git a/src/window/mapped.rs b/src/window/mapped.rs
index 58f733ba..c2bcf87f 100644
--- a/src/window/mapped.rs
+++ b/src/window/mapped.rs
@@ -12,7 +12,9 @@ use smithay::reexports::wayland_protocols::xdg::decoration::zv1::server::zxdg_to
 use smithay::reexports::wayland_protocols::xdg::shell::server::xdg_toplevel;
 use smithay::reexports::wayland_server::protocol::wl_surface::WlSurface;
 use smithay::utils::{Logical, Point, Rectangle, Scale, Size, Transform};
-use smithay::wayland::compositor::{send_surface_state, with_states};
+use smithay::wayland::compositor::{
+    remove_pre_commit_hook, send_surface_state, with_states, HookId,
+};
 use smithay::wayland::shell::xdg::{SurfaceCachedState, ToplevelSurface};
 
 use super::{ResolvedWindowRules, WindowRef};
@@ -26,6 +28,9 @@ use crate::render_helpers::{BakedBuffer, RenderSnapshot, RenderTarget};
 pub struct Mapped {
     pub window: Window,
 
+    /// Pre-commit hook that we have on all mapped toplevel surfaces.
+    pre_commit_hook: HookId,
+
     /// Up-to-date rules.
     rules: ResolvedWindowRules,
 
@@ -46,9 +51,10 @@ pub struct Mapped {
 }
 
 impl Mapped {
-    pub fn new(window: Window, rules: ResolvedWindowRules) -> Self {
+    pub fn new(window: Window, rules: ResolvedWindowRules, hook: HookId) -> Self {
         Self {
             window,
+            pre_commit_hook: hook,
             rules,
             need_to_recompute_rules: false,
             is_focused: false,
@@ -132,6 +138,12 @@ impl Mapped {
     }
 }
 
+impl Drop for Mapped {
+    fn drop(&mut self) {
+        remove_pre_commit_hook(self.toplevel().wl_surface(), self.pre_commit_hook);
+    }
+}
+
 impl LayoutElement for Mapped {
     type Id = Window;