From 1cf9d0f8c50c5488e1da5d23e16a10c8e0f473e0 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@gmail.com>
Date: Tue, 25 Jul 2023 14:45:53 -0400
Subject: [PATCH] refactor (wayland): move seat to client

---
 src/wayland/mod.rs       | 11 ++++++++---
 src/wayland/seat.rs      |  9 +++++----
 src/wayland/state.rs     | 10 +---------
 src/wayland/xdg_shell.rs |  2 +-
 src/wayland/xwayland.rs  |  3 ++-
 5 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/wayland/mod.rs b/src/wayland/mod.rs
index 99d3680..61440f2 100644
--- a/src/wayland/mod.rs
+++ b/src/wayland/mod.rs
@@ -11,6 +11,7 @@ mod xdg_shell;
 pub mod xwayland;
 
 use self::{state::WaylandState, surface::CORE_SURFACES};
+use crate::wayland::seat::SeatData;
 use crate::{core::task, wayland::state::ClientState};
 use color_eyre::eyre::{ensure, Result};
 use global_counter::primitive::exact::CounterU32;
@@ -146,9 +147,13 @@ impl Wayland {
 					}
 					acc = listen_async.accept() => { // New client connected
 						let (stream, _) = acc?;
-						let client = dh2.insert_client(stream.into_std()?, Arc::new(ClientState::default()))?;
-
-						state.lock().new_client(client.id(), &dh2);
+						let client_state = Arc::new(ClientState {
+							compositor_state: Default::default(),
+							display: Arc::downgrade(&display),
+							seat: SeatData::new(&dh1)
+						});
+						let client = dh2.insert_client(stream.into_std()?, client_state.clone())?;
+						client_state.seat.client.set(client.id()).unwrap();
 					}
 					e = dispatch_poll_listener.readable() => { // Dispatch
 						let mut guard = e?;
diff --git a/src/wayland/seat.rs b/src/wayland/seat.rs
index 5a5886e..b17704b 100644
--- a/src/wayland/seat.rs
+++ b/src/wayland/seat.rs
@@ -253,7 +253,7 @@ impl SurfaceInfo {
 }
 
 pub struct SeatData {
-	client: ClientId,
+	pub client: OnceCell<ClientId>,
 	global_id: OnceCell<GlobalId>,
 	surfaces: Mutex<FxHashMap<ObjectId, SurfaceInfo>>,
 	pointer: OnceCell<(WlPointer, Mutex<ObjectId>)>,
@@ -261,9 +261,9 @@ pub struct SeatData {
 	touch: OnceCell<WlTouch>,
 }
 impl SeatData {
-	pub fn new(dh: &DisplayHandle, client: ClientId) -> Arc<Self> {
+	pub fn new(dh: &DisplayHandle) -> Arc<Self> {
 		let seat_data = Arc::new(SeatData {
-			client,
+			client: OnceCell::new(),
 			global_id: OnceCell::new(),
 			surfaces: Mutex::new(FxHashMap::default()),
 			pointer: OnceCell::new(),
@@ -439,7 +439,8 @@ impl GlobalDispatch<WlSeat, Arc<SeatData>, WaylandState> for WaylandState {
 	}
 
 	fn can_view(client: Client, data: &Arc<SeatData>) -> bool {
-		client.id() == data.client
+		let Some(seat_client) = data.client.get().cloned() else {return false};
+		client.id() == seat_client
 	}
 }
 
diff --git a/src/wayland/state.rs b/src/wayland/state.rs
index b29cca4..ab73ca7 100644
--- a/src/wayland/state.rs
+++ b/src/wayland/state.rs
@@ -1,6 +1,5 @@
 use crate::wayland::seat::SeatData;
 use parking_lot::Mutex;
-use rustc_hash::FxHashMap;
 use smithay::{
 	backend::{
 		allocator::dmabuf::Dmabuf,
@@ -37,10 +36,10 @@ use std::sync::{Arc, Weak};
 use tokio::sync::mpsc::UnboundedSender;
 use tracing::{info, warn};
 
-#[derive(Default)]
 pub struct ClientState {
 	pub compositor_state: CompositorClientState,
 	pub display: Weak<Mutex<Display<WaylandState>>>,
+	pub seat: Arc<SeatData>,
 }
 impl ClientState {
 	pub fn flush(&self) {
@@ -73,7 +72,6 @@ pub struct WaylandState {
 	dmabuf_state: (DmabufState, DmabufGlobal, Option<DmabufFeedback>),
 	dmabuf_tx: UnboundedSender<Dmabuf>,
 	pub output: Output,
-	pub seats: FxHashMap<ClientId, Arc<SeatData>>,
 }
 
 impl WaylandState {
@@ -166,15 +164,9 @@ impl WaylandState {
 				dmabuf_state,
 				dmabuf_tx,
 				output,
-				seats: FxHashMap::default(),
 			})
 		})
 	}
-
-	pub fn new_client(&mut self, client: ClientId, dh: &DisplayHandle) {
-		let seat_data = SeatData::new(dh, client.clone());
-		self.seats.insert(client, seat_data);
-	}
 }
 impl Drop for WaylandState {
 	fn drop(&mut self) {
diff --git a/src/wayland/xdg_shell.rs b/src/wayland/xdg_shell.rs
index 6faab23..eb8b794 100644
--- a/src/wayland/xdg_shell.rs
+++ b/src/wayland/xdg_shell.rs
@@ -278,7 +278,7 @@ impl Dispatch<XdgSurface, Mutex<XdgSurfaceData>, WaylandState> for WaylandState
 				xdg_surface.configure(SERIAL_COUNTER.inc());
 
 				let client_credentials = client.get_credentials(&state.display_handle).ok();
-				let seat_data = state.seats.get(&client.id()).unwrap().clone();
+				let Some(seat_data) = client.get_data::<ClientState>().map(|s| s.seat.clone()) else {return};
 				let Some(wl_surface) = xdg_surface_data.lock().wl_surface() else {return};
 				CoreSurface::add_to(
 					state.display_handle.clone(),
diff --git a/src/wayland/xwayland.rs b/src/wayland/xwayland.rs
index 81836e0..f60b196 100644
--- a/src/wayland/xwayland.rs
+++ b/src/wayland/xwayland.rs
@@ -1,5 +1,6 @@
 use super::{
 	seat::{KeyboardEvent, PointerEvent, SeatData},
+	state::ClientState,
 	xdg_shell::PopupData,
 };
 use crate::{
@@ -59,7 +60,7 @@ impl XWaylandState {
 							client_fd: _,
 							display: _,
 						} => {
-							handler.seat = Some(SeatData::new(&dh, client.id()));
+							handler.seat = client.get_data::<ClientState>().map(|s| s.seat.clone());
 							handler.wm =
 								X11Wm::start_wm(handle.clone(), dh.clone(), connection, client)
 									.ok();