From e0d207fbe12d1602cc0c00d67b544e43d7ad8135 Mon Sep 17 00:00:00 2001
From: Victoria Brekenfeld <github@drakulix.de>
Date: Wed, 20 Dec 2023 20:04:51 +0000
Subject: [PATCH] shell: Make element_for_(wl_)surface consider sticky windows

---
 src/shell/mod.rs | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/shell/mod.rs b/src/shell/mod.rs
index 13cc25cf..8a5f692e 100644
--- a/src/shell/mod.rs
+++ b/src/shell/mod.rs
@@ -1264,15 +1264,32 @@ impl Shell {
     }
 
     pub fn element_for_surface(&self, surface: &CosmicSurface) -> Option<&CosmicMapped> {
-        self.workspaces
-            .spaces()
+        self.workspaces.sets.values().find_map(|set| {
+            set.sticky_layer
+                .mapped()
+                .find(|w| w.windows().any(|(s, _)| &s == surface))
+                .or_else(|| {
+                    set.workspaces
+                        .iter()
             .find_map(|w| w.element_for_surface(surface))
+                })
+        })
     }
 
     pub fn element_for_wl_surface(&self, surface: &WlSurface) -> Option<&CosmicMapped> {
-        self.workspaces
-            .spaces()
+        self.workspaces.sets.values().find_map(|set| {
+            set.sticky_layer
+                .mapped()
+                .find(|w| {
+                    w.windows()
+                        .any(|(s, _)| s.wl_surface().as_ref() == Some(surface))
+                })
+                .or_else(|| {
+                    set.workspaces
+                        .iter()
             .find_map(|w| w.element_for_wl_surface(surface))
+                })
+        })
     }
 
     pub fn space_for(&self, mapped: &CosmicMapped) -> Option<&Workspace> {