From 7108cfa365353e4768134324f85b970e4bbed1b9 Mon Sep 17 00:00:00 2001
From: Nova <technobaboo@gmail.com>
Date: Sun, 15 Jan 2023 04:04:09 -0500
Subject: [PATCH] fix(input): grab issues

---
 src/core/registry.rs     | 69 +++++++++++++++++++++++++---------------
 src/nodes/input/mod.rs   | 11 ++++---
 src/nodes/spatial/mod.rs |  5 +++
 3 files changed, 56 insertions(+), 29 deletions(-)

diff --git a/src/core/registry.rs b/src/core/registry.rs
index eb67879..30d8517 100644
--- a/src/core/registry.rs
+++ b/src/core/registry.rs
@@ -1,17 +1,20 @@
 #![allow(dead_code)]
 
+use parking_lot::{const_mutex, MappedMutexGuard, Mutex, MutexGuard};
+use rustc_hash::FxHashMap;
 use std::ptr;
 use std::sync::{Arc, Weak};
 
-use once_cell::sync::Lazy;
-use parking_lot::Mutex;
-use rustc_hash::FxHashMap;
-
-pub struct Registry<T: Send + Sync + ?Sized>(Lazy<Mutex<FxHashMap<usize, Weak<T>>>>);
+pub struct Registry<T: Send + Sync + ?Sized>(Mutex<Option<FxHashMap<usize, Weak<T>>>>);
 
 impl<T: Send + Sync + ?Sized> Registry<T> {
 	pub const fn new() -> Self {
-		Registry(Lazy::new(|| Mutex::new(FxHashMap::default())))
+		Registry(const_mutex(None))
+	}
+	fn lock(&self) -> MappedMutexGuard<FxHashMap<usize, Weak<T>>> {
+		MutexGuard::map(self.0.lock(), |r| {
+			r.get_or_insert_with(|| FxHashMap::default())
+		})
 	}
 	pub fn add(&self, t: T) -> Arc<T>
 	where
@@ -22,37 +25,52 @@ impl<T: Send + Sync + ?Sized> Registry<T> {
 		t_arc
 	}
 	pub fn add_raw(&self, t: &Arc<T>) {
-		self.0
-			.lock()
+		self.lock()
 			.insert(Arc::as_ptr(t) as *const () as usize, Arc::downgrade(t));
 	}
 	pub fn contains(&self, t: &T) -> bool {
-		self.0
-			.lock()
+		self.lock()
 			.contains_key(&(ptr::addr_of!(*t) as *const () as usize))
 	}
 	pub fn get_valid_contents(&self) -> Vec<Arc<T>> {
-		self.0
-			.lock()
+		self.lock()
 			.iter()
 			.filter_map(|pair| pair.1.upgrade())
 			.collect()
 	}
-	pub fn remove(&self, t: &T) {
+	pub fn take_valid_contents(&self) -> Vec<Arc<T>> {
 		self.0
 			.lock()
+			.take()
+			.unwrap_or_default()
+			.into_iter()
+			.filter_map(|pair| pair.1.upgrade())
+			.collect()
+	}
+	pub fn remove(&self, t: &T) {
+		self.lock()
 			.remove(&(ptr::addr_of!(*t) as *const () as usize));
 	}
 	pub fn clear(&self) {
-		self.0.lock().clear();
+		self.lock().clear();
+	}
+}
+impl<T: Send + Sync + ?Sized> Clone for Registry<T> {
+	fn clone(&self) -> Self {
+		Self(Mutex::new(self.0.lock().clone()))
 	}
 }
 
-pub struct OwnedRegistry<T: Send + Sync + ?Sized>(Lazy<Mutex<FxHashMap<usize, Arc<T>>>>);
+pub struct OwnedRegistry<T: Send + Sync + ?Sized>(Mutex<Option<FxHashMap<usize, Arc<T>>>>);
 
 impl<T: Send + Sync + ?Sized> OwnedRegistry<T> {
 	pub const fn new() -> Self {
-		OwnedRegistry(Lazy::new(|| Mutex::new(FxHashMap::default())))
+		OwnedRegistry(const_mutex(None))
+	}
+	fn lock(&self) -> MappedMutexGuard<FxHashMap<usize, Arc<T>>> {
+		MutexGuard::map(self.0.lock(), |r| {
+			r.get_or_insert_with(|| FxHashMap::default())
+		})
 	}
 	pub fn add(&self, t: T) -> Arc<T>
 	where
@@ -63,24 +81,25 @@ impl<T: Send + Sync + ?Sized> OwnedRegistry<T> {
 		t_arc
 	}
 	pub fn add_raw(&self, t: Arc<T>) {
-		self.0
-			.lock()
-			.insert(Arc::as_ptr(&t) as *const () as usize, t);
+		self.lock().insert(Arc::as_ptr(&t) as *const () as usize, t);
 	}
 	pub fn get_vec(&self) -> Vec<Arc<T>> {
-		self.0.lock().values().cloned().collect::<Vec<_>>()
+		self.lock().values().cloned().collect::<Vec<_>>()
 	}
 	pub fn contains(&self, t: &T) -> bool {
-		self.0
-			.lock()
+		self.lock()
 			.contains_key(&(ptr::addr_of!(*t) as *const () as usize))
 	}
 	pub fn remove(&self, t: &T) {
-		self.0
-			.lock()
+		self.lock()
 			.remove(&(ptr::addr_of!(*t) as *const () as usize));
 	}
 	pub fn clear(&self) {
-		self.0.lock().clear();
+		self.lock().clear();
+	}
+}
+impl<T: Send + Sync + ?Sized> Clone for OwnedRegistry<T> {
+	fn clone(&self) -> Self {
+		Self(Mutex::new(self.0.lock().clone()))
 	}
 }
diff --git a/src/nodes/input/mod.rs b/src/nodes/input/mod.rs
index 3faa0c4..a592562 100644
--- a/src/nodes/input/mod.rs
+++ b/src/nodes/input/mod.rs
@@ -215,6 +215,11 @@ impl InputHandler {
 		}
 	}
 }
+impl PartialEq for InputHandler {
+	fn eq(&self, other: &Self) -> bool {
+		self.spatial == other.spatial
+	}
+}
 impl Drop for InputHandler {
 	fn drop(&mut self) {
 		INPUT_HANDLER_REGISTRY.remove(self);
@@ -283,19 +288,17 @@ pub fn process_input() {
 			// Get the current frame
 			let frame = FRAME.load(Ordering::Relaxed);
 
+			let captures = method.captures.take_valid_contents();
 			// Iterate over the distance links and send input to them
 			for distance_link in distance_links {
 				distance_link.send_input(frame, method.datamap.lock().clone().unwrap());
 
 				// If the current distance link is in the list of captured input handlers,
 				// break out of the loop to avoid sending input to the remaining distance links
-				if method.captures.contains(&distance_link.handler) {
+				if captures.contains(&distance_link.handler) {
 					break;
 				}
 			}
-
-			// Clear the list of captured input handlers for this method
-			method.captures.clear();
 		});
 	}
 }
diff --git a/src/nodes/spatial/mod.rs b/src/nodes/spatial/mod.rs
index 28b1690..2439de9 100644
--- a/src/nodes/spatial/mod.rs
+++ b/src/nodes/spatial/mod.rs
@@ -268,6 +268,11 @@ impl Spatial {
 			.unwrap_or(f32::MAX)
 	}
 }
+impl PartialEq for Spatial {
+	fn eq(&self, other: &Self) -> bool {
+		self.uid == other.uid
+	}
+}
 impl Debug for Spatial {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		f.debug_struct("Spatial")