Fix some cases of incomplete search for surface output

Most visibly, fixes screen not immediately redrawing upon layer-shell
popup commits.

There's still a number of places with questionable handling left, mostly
to do with subsurfaces (like, find_popup_root_surface() doesn't go up to
subsurfaces), and session-lock. I don't have good clients to test these.

4 files changed, 36 insertions, 23 deletions

diff --git a/src/handlers/compositor.rs b/src/handlers/compositor.rs
index 94ee8549..30963f47 100644
--- a/src/handlers/compositor.rs
+++ b/src/handlers/compositor.rs
@@ -1,7 +1,6 @@
 use std::collections::hash_map::Entry;
 
 use smithay::backend::renderer::utils::{on_commit_buffer_handler, with_renderer_surface_state};
-use smithay::desktop::find_popup_root_surface;
 use smithay::input::pointer::CursorImageStatus;
 use smithay::reexports::calloop::Interest;
 use smithay::reexports::wayland_server::protocol::wl_buffer;
@@ -30,7 +29,12 @@ impl CompositorHandler for State {
     }
 
     fn new_subsurface(&mut self, surface: &WlSurface, parent: &WlSurface) {
-        if let Some((_, output)) = self.niri.layout.find_window_and_output(parent) {
+        let mut root = parent.clone();
+        while let Some(parent) = get_parent(&root) {
+            root = parent;
+        }
+
+        if let Some(output) = self.niri.output_for_root(&root) {
             let scale = output.current_scale().integer_scale();
             let transform = output.current_transform();
             with_states(surface, |data| {
@@ -150,11 +154,8 @@ impl CompositorHandler for State {
         // This might be a popup.
         self.popups_handle_commit(surface);
         if let Some(popup) = self.niri.popups.find_popup(surface) {
-            if let Ok(root) = find_popup_root_surface(&popup) {
-                let root_window_output = self.niri.layout.find_window_and_output(&root);
-                if let Some((_window, output)) = root_window_output {
-                    self.niri.queue_redraw(output);
-                }
+            if let Some(output) = self.output_for_popup(&popup) {
+                self.niri.queue_redraw(output);
             }
         }
 
diff --git a/src/handlers/mod.rs b/src/handlers/mod.rs
index 65e7d101..ebf53c62 100644
--- a/src/handlers/mod.rs
+++ b/src/handlers/mod.rs
@@ -80,10 +80,8 @@ delegate_text_input_manager!(State);
 
 impl InputMethodHandler for State {
     fn new_popup(&mut self, surface: PopupSurface) {
-        if let Some((_, output)) = surface
-            .get_parent()
-            .and_then(|parent| self.niri.layout.find_window_and_output(&parent.surface))
-        {
+        let popup = PopupKind::from(surface.clone());
+        if let Some(output) = self.output_for_popup(&popup) {
             let scale = output.current_scale().integer_scale();
             let transform = output.current_transform();
             let wl_surface = surface.wl_surface();
@@ -91,7 +89,7 @@ impl InputMethodHandler for State {
                 send_surface_state(wl_surface, data, scale, transform);
             });
         }
-        if let Err(err) = self.niri.popups.track_popup(PopupKind::from(surface)) {
+        if let Err(err) = self.niri.popups.track_popup(popup) {
             warn!("error tracking ime popup {err:?}");
         }
     }
diff --git a/src/handlers/xdg_shell.rs b/src/handlers/xdg_shell.rs
index a99b344f..3eafcdbf 100644
--- a/src/handlers/xdg_shell.rs
+++ b/src/handlers/xdg_shell.rs
@@ -175,11 +175,8 @@ impl XdgShellHandler for State {
     }
 
     fn popup_destroyed(&mut self, surface: PopupSurface) {
-        if let Ok(root) = find_popup_root_surface(&surface.into()) {
-            let root_window_output = self.niri.layout.find_window_and_output(&root);
-            if let Some((_window, output)) = root_window_output {
-                self.niri.queue_redraw(output);
-            }
+        if let Some(output) = self.output_for_popup(&PopupKind::Xdg(surface)) {
+            self.niri.queue_redraw(output);
         }
     }
 }
@@ -271,12 +268,8 @@ impl State {
                             .initial_configure_sent
                     });
                     if !initial_configure_sent {
-                        if let Some(output) = popup.get_parent_surface().and_then(|parent| {
-                            self.niri
-                                .layout
-                                .find_window_and_output(&parent)
-                                .map(|(_, output)| output)
-                        }) {
+                        if let Some(output) = self.output_for_popup(&PopupKind::Xdg(popup.clone()))
+                        {
                             let scale = output.current_scale().integer_scale();
                             let transform = output.current_transform();
                             with_states(surface, |data| {
@@ -291,4 +284,9 @@ impl State {
             }
         }
     }
+
+    pub fn output_for_popup(&self, popup: &PopupKind) -> Option<Output> {
+        let root = find_popup_root_surface(popup).ok()?;
+        self.niri.output_for_root(&root)
+    }
 }
diff --git a/src/niri.rs b/src/niri.rs
index 8221ba6c..084d51e9 100644
--- a/src/niri.rs
+++ b/src/niri.rs
@@ -1211,6 +1211,22 @@ impl Niri {
             .or_else(|| self.global_space.outputs().next())
     }
 
+    pub fn output_for_root(&self, root: &WlSurface) -> Option<Output> {
+        // Check the main layout.
+        let win_out = self.layout.find_window_and_output(root);
+        let layout_output = win_out.map(|(_, output)| output);
+
+        // Check layer-shell.
+        let has_layer_surface = |o: &&Output| {
+            layer_map_for_output(o)
+                .layer_for_surface(root, WindowSurfaceType::TOPLEVEL)
+                .is_some()
+        };
+        let layer_shell_output = || self.layout.outputs().find(has_layer_surface).cloned();
+
+        layout_output.or_else(layer_shell_output)
+    }
+
     fn lock_surface_focus(&self) -> Option<WlSurface> {
         let output_under_cursor = self.output_under_cursor();
         let output = output_under_cursor