From 4de316cbce322c2e17adb7bb784b126e0d5ca677 Mon Sep 17 00:00:00 2001
From: Victoria Brekenfeld <victoria@system76.com>
Date: Mon, 11 Jul 2022 22:35:48 +0200
Subject: [PATCH] shell: Fix output_for_surfaces for layer surfaces

---
 src/shell/mod.rs | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index 74cb1562..c5a09db4 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -441,17 +441,21 @@ impl Shell {
         }
     }
 
-    pub fn outputs_for_surface(&self, surface: &WlSurface) -> impl Iterator<Item = Output> {
-        self.space_for_surface(surface)
-            .and_then(|w| {
+    pub fn outputs_for_surface<'a>(&'a self, surface: &'a WlSurface) -> impl Iterator<Item = Output> + 'a {
+        match self.outputs.iter().find(|o| {
+            let map = layer_map_for_output(o);
+            map.layer_for_surface(surface, WindowSurfaceType::ALL).is_some()
+        }) {
+            Some(output) => Box::new(std::iter::once(output.clone())) as Box<dyn Iterator<Item = Output>>,
+            None => Box::new(self.spaces.iter().filter_map(|w| {
                 if let Some(window) = w.space.window_for_surface(surface, WindowSurfaceType::ALL) {
                     Some(w.space.outputs_for_window(&window).into_iter())
                 } else {
                     None
                 }
             })
-            .into_iter()
-            .flatten()
+            .flatten()),
+        }
     }
 
     pub fn space_for_surface(&self, surface: &WlSurface) -> Option<&Workspace> {