refactor(wayland): xdg surface sub-roles

src/wayland/core/surface.rs

@@ -8,7 +8,7 @@ use crate::{
 		core::buffer::BufferUsage,
 		presentation::{MonotonicTimestamp, PresentationFeedback},
 		util::{ClientExt, DoubleBuffer},
-		xdg::{popup::Popup, toplevel::Toplevel},
+		xdg::wm_base::XdgSurfaceRole,
 	},
 };
 use bevy::{
@@ -35,8 +35,9 @@ pub static WL_SURFACE_REGISTRY: Registry<Surface> = Registry::new();
 
 #[derive(Debug, Clone)]
 pub enum SurfaceRole {
-	XdgToplevel(Arc<Toplevel>),
+	Cursor,
-	XDGPopup(Arc<Popup>),
+	Subsurface,
+	Xdg(OnceLock<XdgSurfaceRole>),
 }
 
 #[derive(Debug, Clone)]
@@ -75,8 +76,7 @@ pub struct Surface {
 	pub id: ObjectId,
 	state: Mutex<DoubleBuffer<SurfaceState>>,
 	pub message_sink: MessageSink,
-	// TODO: This should probably be a OnceLock? wayland doesn't support changing the surface role
+	pub role: OnceLock<SurfaceRole>,
-	pub role: Mutex<Option<SurfaceRole>>,
 	on_commit_handlers: Mutex<Vec<OnCommitCallback>>,
 	material: OnceLock<Handle<BevyMaterial>>,
 	pending_material_applications: Registry<ModelPart>,
@@ -102,7 +102,7 @@ impl Surface {
 			id,
 			state: Default::default(),
 			message_sink: client.message_sink(),
-			role: Mutex::new(None),
+			role: OnceLock::new(),
 			on_commit_handlers: Mutex::new(Vec::new()),
 			material: OnceLock::new(),
 			pending_material_applications: Registry::new(),
@@ -416,9 +416,8 @@ impl WlSurface for Surface {
 		Ok(())
 	}
 }
-
 impl Drop for Surface {
 	fn drop(&mut self) {
-		self.role.lock().take();
+		self.role.take();
 	}
 }

src/wayland/xdg/surface.rs

@@ -1,6 +1,12 @@
 use super::{popup::Popup, positioner::Positioner, toplevel::Mapped};
-use crate::wayland::{core::surface::SurfaceRole, display::Display, xdg::toplevel::Toplevel};
+use crate::wayland::util::ClientExt;
+use crate::wayland::{
+	core::surface::SurfaceRole,
+	display::Display,
+	xdg::{toplevel::Toplevel, wm_base::XdgSurfaceRole},
+};
 use std::sync::Arc;
+use std::sync::OnceLock;
 use std::sync::Weak;
 pub use waynest::server::protocol::stable::xdg_shell::xdg_surface::*;
 use waynest::{
@@ -64,29 +70,76 @@ impl XdgSurface for Surface {
 				client.get::<Self>(sender_id).unwrap(),
 			),
 		);
+		// Set up the XDG role if not already done
+		if surface.role.get().is_none() {
+			let xdg_role = SurfaceRole::Xdg(OnceLock::new());
 
-		{
+			if surface.role.set(xdg_role).is_err() {
-			let mut surface_role = surface.role.lock();
+				return client
-
+					.protocol_error(
-			// A surface must not have any existing role when assigning a new one
+						sender_id,
-			// "A surface must not have more than one role, and a role must not be assigned to more than one
+						toplevel_id,
-			// surface at a time. However, wl_surface role-specific interfaces may reassign the role, allow
+						1, // XDG_WM_BASE_ERROR_ROLE
-			// a role to be destroyed, or allow multiple role-specific interfaces to share the same role."
+						"Failed to set surface role (race condition)".to_string(),
-			// - xdg_surface protocol doc
+					)
-			if surface_role.is_some() {
+					.await;
-				// We should send "role" error here as per xdg_wm_base.error enum
-				// But we'll ignore for now
-			} else {
-				surface_role.replace(SurfaceRole::XdgToplevel(toplevel.clone()));
 			}
 		}
 
+		// Check if the surface already has an XDG role
+		let surface_role = surface.role.get().unwrap();
+
+		// Now check if this is an XDG surface and set the sub-role
+		if let SurfaceRole::Xdg(role) = surface_role {
+			if role.get().is_some() {
+				return client
+					.protocol_error(
+						sender_id,
+						toplevel_id,
+						1, // XDG_WM_BASE_ERROR_ROLE
+						"XDG surface already has a sub-role".to_string(),
+					)
+					.await;
+			}
+
+			if role
+				.set(XdgSurfaceRole::Toplevel(toplevel.clone()))
+				.is_err()
+			{
+				return client
+					.protocol_error(
+						sender_id,
+						toplevel_id,
+						1, // XDG_WM_BASE_ERROR_ROLE
+						"Failed to set XDG sub-role (race condition)".to_string(),
+					)
+					.await;
+			}
+		} else {
+			return client
+				.protocol_error(
+					sender_id,
+					toplevel_id,
+					1, // XDG_WM_BASE_ERROR_ROLE
+					"Surface has a non-XDG role".to_string(),
+				)
+				.await;
+		}
+
 		toplevel.reconfigure(client).await?;
 
 		let pid = client.get::<Display>(ObjectId::DISPLAY).unwrap().pid;
 		let configured = self.configured.clone();
 		surface.add_commit_handler(move |surface, state| {
-			let Some(SurfaceRole::XdgToplevel(toplevel)) = &mut *surface.role.lock() else {
+			let Some(role_ref) = surface.role.get() else {
+				return true;
+			};
+
+			let SurfaceRole::Xdg(role) = role_ref else {
+				return true;
+			};
+
+			let Some(XdgSurfaceRole::Toplevel(toplevel)) = role.get() else {
 				return true;
 			};
 
@@ -119,13 +172,35 @@ impl XdgSurface for Surface {
 		parent: Option<ObjectId>,
 		positioner: ObjectId,
 	) -> Result<()> {
-		let parent = client.get::<Surface>(parent.unwrap()).unwrap();
+		let Some(parent) = parent else {
-		let panel_item = match parent.wl_surface().role.lock().as_ref().unwrap() {
+			return client
-			SurfaceRole::XdgToplevel(toplevel) => {
+				.protocol_error(
-				let toplevel_lock = toplevel.mapped.lock();
+					sender_id,
-				toplevel_lock.as_ref().unwrap()._panel_item.clone()
+					popup_id,
+					3, // INVALID_POPUP_PARENT
+					"Parent surface does not have an XDG role".to_string(),
+				)
+				.await;
+		};
+		let parent = client.get::<Surface>(parent).unwrap();
+		let panel_item = match parent.wl_surface().role.get().unwrap() {
+			SurfaceRole::Xdg(role) => match role.get().unwrap() {
+				XdgSurfaceRole::Toplevel(toplevel) => {
+					let toplevel_lock = toplevel.mapped.lock();
+					toplevel_lock.as_ref().unwrap()._panel_item.clone()
+				}
+				XdgSurfaceRole::Popup(popup) => popup.panel_item.upgrade().unwrap(),
+			},
+			_ => {
+				return client
+					.protocol_error(
+						sender_id,
+						popup_id,
+						1, // XDG_WM_BASE_ERROR_ROLE
+						"Parent surface does not have an XDG role".to_string(),
+					)
+					.await;
 			}
-			SurfaceRole::XDGPopup(popup) => popup.panel_item.upgrade().unwrap(),
 		};
 		let positioner = client.get::<Positioner>(positioner).unwrap();
 
@@ -143,15 +218,57 @@ impl XdgSurface for Surface {
 			),
 		);
 
-		{
+		// Set up the XDG role if not already done
-			let wl_surface = self.wl_surface();
+		let wl_surface = self.wl_surface();
-			let mut surface_role = wl_surface.role.lock();
+		if wl_surface.role.get().is_none() {
+			let xdg_role = SurfaceRole::Xdg(OnceLock::new());
 
-			if surface_role.is_some() {
+			if wl_surface.role.set(xdg_role).is_err() {
-				// We should send "role" error here as per xdg_wm_base.error enum
+				return client
-				// But we'll ignore for now
+					.protocol_error(
-			} else {
+						sender_id,
-				surface_role.replace(SurfaceRole::XDGPopup(popup.clone()));
+						popup_id,
+						1, // XDG_WM_BASE_ERROR_ROLE
+						"Failed to set surface role (race condition)".to_string(),
+					)
+					.await;
+			}
+		}
+
+		// Now check if this is an XDG surface and set the sub-role
+		match wl_surface.role.get().unwrap() {
+			SurfaceRole::Xdg(role) => {
+				if role.get().is_some() {
+					return client
+						.protocol_error(
+							sender_id,
+							popup_id,
+							1, // XDG_WM_BASE_ERROR_ROLE
+							"XDG surface already has a sub-role".to_string(),
+						)
+						.await;
+				}
+
+				if role.set(XdgSurfaceRole::Popup(popup.clone())).is_err() {
+					return client
+						.protocol_error(
+							sender_id,
+							popup_id,
+							1, // XDG_WM_BASE_ERROR_ROLE
+							"Failed to set XDG sub-role (race condition)".to_string(),
+						)
+						.await;
+				}
+			}
+			_ => {
+				return client
+					.protocol_error(
+						sender_id,
+						popup_id,
+						1, // XDG_WM_BASE_ERROR_ROLE
+						"Surface has a non-XDG role".to_string(),
+					)
+					.await;
 			}
 		}
 

src/wayland/xdg/wm_base.rs

@@ -1,11 +1,19 @@
+use super::popup::Popup;
+use super::positioner::Positioner;
+use super::toplevel::Toplevel;
 use crate::wayland::xdg::surface::Surface;
+use std::sync::Arc;
 pub use waynest::server::protocol::stable::xdg_shell::xdg_wm_base::*;
 use waynest::{
 	server::{Client, Dispatcher, Result},
 	wire::ObjectId,
 };
 
-use super::positioner::Positioner;
+#[derive(Debug, Clone)]
+pub enum XdgSurfaceRole {
+	Toplevel(Arc<Toplevel>),
+	Popup(Arc<Popup>),
+}
 
 #[derive(Debug, Dispatcher, Default)]
 pub struct WmBase {