element: Introduce CosmicMappedKey for safely hashing windows across threads

src/backend/render/mod.rs

@@ -4,7 +4,7 @@ use std::{
     borrow::{Borrow, BorrowMut},
     cell::RefCell,
     collections::HashMap,
-    sync::Weak,
+    sync::{Arc, RwLock, Weak},
     time::Instant,
 };
 
@@ -14,6 +14,7 @@ use crate::{
     backend::render::element::DamageElement,
     config::Config,
     shell::{
+        element::CosmicMappedKey,
         focus::target::WindowGroup,
         grabs::{SeatMenuGrabState, SeatMoveGrabState},
         layout::tiling::ANIMATION_DURATION,
@@ -112,7 +113,7 @@ pub enum Usage {
 pub enum Key {
     Static(Id),
     Group(Weak<()>),
-    Window(Usage, CosmicMapped),
+    Window(Usage, CosmicMappedKey),
 }
 impl std::hash::Hash for Key {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {

src/shell/element/mod.rs

@@ -4,7 +4,7 @@ use crate::{
         GlMultiError, GlMultiFrame, GlMultiRenderer,
     },
     state::State,
-    utils::prelude::*,
+    utils::{iced::IcedElementInternal, prelude::*},
 };
 use calloop::LoopHandle;
 use id_tree::NodeId;
@@ -40,13 +40,15 @@ use smithay::{
     },
     xwayland::{xwm::X11Relatable, X11Surface},
 };
+use stack::CosmicStackInternal;
+use window::CosmicWindowInternal;
 
 use std::{
     borrow::Cow,
     collections::HashMap,
     fmt,
     hash::Hash,
-    sync::{atomic::AtomicBool, Arc, Mutex},
+    sync::{atomic::AtomicBool, Arc, Mutex, Weak},
 };
 
 pub mod surface;
@@ -126,6 +128,46 @@ impl fmt::Debug for CosmicMapped {
     }
 }
 
+#[derive(Clone)]
+pub struct CosmicMappedKey(CosmicMappedKeyInner);
+#[derive(Clone)]
+enum CosmicMappedKeyInner {
+    Window(Weak<Mutex<IcedElementInternal<CosmicWindowInternal>>>),
+    Stack(Weak<Mutex<IcedElementInternal<CosmicStackInternal>>>),
+}
+
+impl Hash for CosmicMappedKey {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        match &self.0 {
+            CosmicMappedKeyInner::Window(weak) => weak.as_ptr().hash(state),
+            CosmicMappedKeyInner::Stack(weak) => weak.as_ptr().hash(state),
+        }
+    }
+}
+
+impl IsAlive for CosmicMappedKey {
+    fn alive(&self) -> bool {
+        match &self.0 {
+            CosmicMappedKeyInner::Window(weak) => weak.strong_count() > 0,
+            CosmicMappedKeyInner::Stack(weak) => weak.strong_count() > 0,
+        }
+    }
+}
+
+impl PartialEq for CosmicMappedKey {
+    fn eq(&self, other: &Self) -> bool {
+        match (&self.0, &other.0) {
+            (CosmicMappedKeyInner::Window(weak1), CosmicMappedKeyInner::Window(weak2)) => {
+                Weak::ptr_eq(weak1, weak2)
+            }
+            (CosmicMappedKeyInner::Stack(weak1), CosmicMappedKeyInner::Stack(weak2)) => {
+                Weak::ptr_eq(weak1, weak2)
+            }
+            _ => false,
+        }
+    }
+}
+
 impl PartialEq for CosmicMapped {
     fn eq(&self, other: &Self) -> bool {
         self.element == other.element
@@ -835,6 +877,18 @@ impl CosmicMapped {
             CosmicMappedInternal::_GenericCatcher(_) => {}
         }
     }
+
+    pub fn key(&self) -> CosmicMappedKey {
+        CosmicMappedKey(match &self.element {
+            CosmicMappedInternal::Stack(stack) => {
+                CosmicMappedKeyInner::Stack(Arc::downgrade(&stack.0 .0))
+            }
+            CosmicMappedInternal::Window(window) => {
+                CosmicMappedKeyInner::Window(Arc::downgrade(&window.0 .0))
+            }
+            _ => unreachable!(),
+        })
+    }
 }
 
 impl IsAlive for CosmicMapped {

src/shell/element/stack.rs

@@ -77,7 +77,7 @@ use self::{
 static SCROLLABLE_ID: Lazy<Id> = Lazy::new(|| Id::new("scrollable"));
 
 #[derive(Clone, PartialEq, Eq, Hash)]
-pub struct CosmicStack(IcedElement<CosmicStackInternal>);
+pub struct CosmicStack(pub(super) IcedElement<CosmicStackInternal>);
 
 impl fmt::Debug for CosmicStack {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

src/shell/element/window.rs

@@ -62,7 +62,7 @@ use super::{
 };
 
 #[derive(Clone, PartialEq, Eq, Hash)]
-pub struct CosmicWindow(IcedElement<CosmicWindowInternal>);
+pub struct CosmicWindow(pub(super) IcedElement<CosmicWindowInternal>);
 
 impl fmt::Debug for CosmicWindow {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {

src/shell/grabs/moving.rs

@@ -111,7 +111,7 @@ impl MoveGrabState {
             Some(
                 CosmicMappedRenderElement::from(IndicatorShader::focus_element(
                     renderer,
-                    Key::Window(Usage::MoveGrabIndicator, self.window.clone()),
+                    Key::Window(Usage::MoveGrabIndicator, self.window.key()),
                     Rectangle::from_loc_and_size(
                         render_location,
                         self.window
@@ -151,7 +151,7 @@ impl MoveGrabState {
                 vec![
                     CosmicMappedRenderElement::from(IndicatorShader::element(
                         renderer,
-                        Key::Window(Usage::SnappingIndicator, self.window.clone()),
+                        Key::Window(Usage::SnappingIndicator, self.window.key()),
                         overlay_geometry,
                         3,
                         theme.radius_s()[0] as u8, // TODO: Fix once shaders support 4 corner radii customization
@@ -166,7 +166,7 @@ impl MoveGrabState {
                     .into(),
                     CosmicMappedRenderElement::from(BackdropShader::element(
                         renderer,
-                        Key::Window(Usage::SnappingIndicator, self.window.clone()),
+                        Key::Window(Usage::SnappingIndicator, self.window.key()),
                         t.overlay_geometry(non_exclusive_geometry, gaps),
                         theme.radius_s()[0], // TODO: Fix once shaders support 4 corner radii customization
                         0.4,

src/shell/layout/floating/mod.rs

@@ -1375,7 +1375,7 @@ impl FloatingLayout {
                 if indicator_thickness > 0 {
                     let element = IndicatorShader::focus_element(
                         renderer,
-                        Key::Window(Usage::FocusIndicator, elem.clone()),
+                        Key::Window(Usage::FocusIndicator, elem.key()),
                         geometry,
                         indicator_thickness,
                         output_scale,

src/shell/layout/tiling/mod.rs

@@ -4549,7 +4549,7 @@ where
                             elements.push(
                                 IndicatorShader::element(
                                     *renderer,
-                                    Key::Window(Usage::PotentialGroupIndicator, mapped.clone()),
+                                    Key::Window(Usage::PotentialGroupIndicator, mapped.key()),
                                     geo,
                                     4,
                                     8,
@@ -4589,7 +4589,7 @@ where
                             elements.push(
                                 BackdropShader::element(
                                     *renderer,
-                                    Key::Window(Usage::OverviewBackdrop, mapped.clone()),
+                                    Key::Window(Usage::OverviewBackdrop, mapped.key()),
                                     geo,
                                     8.,
                                     alpha
@@ -4800,7 +4800,7 @@ where
                     window_elements.push(CosmicMappedRenderElement::FocusIndicator(
                         IndicatorShader::focus_element(
                             renderer,
-                            Key::Window(Usage::FocusIndicator, mapped.clone().into()),
+                            Key::Window(Usage::FocusIndicator, mapped.clone().key()),
                             geo,
                             indicator_thickness,
                             output_scale,
@@ -5133,7 +5133,7 @@ where
                             renderer,
                             match data {
                                 Data::Mapped { mapped, .. } => {
-                                    Key::Window(Usage::FocusIndicator, mapped.clone().into())
+                                    Key::Window(Usage::FocusIndicator, mapped.clone().key())
                                 }
                                 Data::Group { alive, .. } => Key::Group(Arc::downgrade(alive)),
                                 _ => unreachable!(),
@@ -5250,7 +5250,7 @@ where
                                     .stack_ref()
                                     .map(|stack| &stack.active() == stack_window)
                                     .unwrap_or(false),
-                                _ => unreachable!(),
+                                _ => unreachable!(), // TODO: We could swap with a group
                             })
                             .unwrap_or(false)
                     })
@@ -5262,7 +5262,7 @@ where
                         0,
                         CosmicMappedRenderElement::Overlay(BackdropShader::element(
                             renderer,
-                            Key::Window(Usage::Overlay, mapped.clone()),
+                            Key::Window(Usage::Overlay, mapped.key()),
                             geo,
                             0.0,
                             0.3,

src/utils/iced.rs

@@ -67,7 +67,7 @@ use smithay::{
     },
 };
 
-pub struct IcedElement<P: Program + Send + 'static>(Arc<Mutex<IcedElementInternal<P>>>);
+pub struct IcedElement<P: Program + Send + 'static>(pub(crate) Arc<Mutex<IcedElementInternal<P>>>);
 
 impl<P: Program + Send + 'static> fmt::Debug for IcedElement<P> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -75,8 +75,8 @@ impl<P: Program + Send + 'static> fmt::Debug for IcedElement<P> {
     }
 }
 
-// SAFETY: We cannot really be sure about `iced_native::program::State` sadly,
+// SAFETY: It is not, we need to make sure we never move this into another thread and drop it there,
-// but the rest should be fine.
+//  as on drop it could trigger RefCells in calloop.
 unsafe impl<P: Program + Send + 'static> Send for IcedElementInternal<P> {}
 
 impl<P: Program + Send + 'static> Clone for IcedElement<P> {
@@ -139,7 +139,7 @@ impl<P: Program> IcedProgram for ProgramWrapper<P> {
     }
 }
 
-struct IcedElementInternal<P: Program + Send + 'static> {
+pub(crate) struct IcedElementInternal<P: Program + Send + 'static> {
     // draw buffer
     outputs: HashSet<Output>,
     buffers: HashMap<OrderedFloat<f64>, (MemoryRenderBuffer, Option<(Vec<Primitive>, Color)>)>,