Keyable

5 files changed, 58 insertions, 30 deletions

diff --git a/src/guard.rs b/src/guard.rs
index 3c1b636..30e7d4a 100644
--- a/src/guard.rs
+++ b/src/guard.rs
@@ -1,17 +1,17 @@
 use std::ops::{Deref, DerefMut};
 
-use crate::{lockable::Lockable, ThreadKey};
+use crate::{key::Keyable, lockable::Lockable};
 
 /// A guard for a generic [`Lockable`] type.
-pub struct LockGuard<'a, L: Lockable<'a>> {
+pub struct LockGuard<'a, L: Lockable<'a>, Key: Keyable> {
 	guard: L::Output,
-	_key: &'a mut ThreadKey,
+	_key: Key,
 }
 
-impl<'a, L: Lockable<'a>> LockGuard<'a, L> {
+impl<'a, L: Lockable<'a>, Key: Keyable> LockGuard<'a, L, Key> {
 	/// Locks the lockable type and returns a guard that can be used to access
 	/// the underlying data.
-	pub fn lock(lock: &'a L, key: &'a mut ThreadKey) -> Self {
+	pub fn lock(lock: &'a L, key: Key) -> Self {
 		Self {
 			// safety: we have the thread's key
 			guard: unsafe { lock.lock() },
@@ -22,7 +22,7 @@ impl<'a, L: Lockable<'a>> LockGuard<'a, L> {
 	/// Attempts to lock the guard without blocking. If successful, this method
 	/// returns a guard that can be used to access the data. Otherwise, the key
 	/// is given back as an error.
-	pub fn try_lock(lock: &'a L, key: &'a mut ThreadKey) -> Option<Self> {
+	pub fn try_lock(lock: &'a L, key: Key) -> Option<Self> {
 		// safety: we have the thread's key
 		unsafe { lock.try_lock() }.map(|guard| Self { guard, _key: key })
 	}
@@ -35,7 +35,7 @@ impl<'a, L: Lockable<'a>> LockGuard<'a, L> {
 	}
 }
 
-impl<'a, L: Lockable<'a>> Deref for LockGuard<'a, L> {
+impl<'a, L: Lockable<'a>, Key: Keyable> Deref for LockGuard<'a, L, Key> {
 	type Target = L::Output;
 
 	fn deref(&self) -> &Self::Target {
@@ -43,7 +43,7 @@ impl<'a, L: Lockable<'a>> Deref for LockGuard<'a, L> {
 	}
 }
 
-impl<'a, L: Lockable<'a>> DerefMut for LockGuard<'a, L> {
+impl<'a, L: Lockable<'a>, Key: Keyable> DerefMut for LockGuard<'a, L, Key> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
 		&mut self.guard
 	}
diff --git a/src/lock.rs b/src/key.rs
index 942474a..92f3b99 100644
--- a/src/lock.rs
+++ b/src/key.rs
@@ -5,7 +5,16 @@ use std::sync::atomic::{AtomicBool, Ordering};
 use once_cell::sync::Lazy;
 use thread_local::ThreadLocal;
 
-static KEY: Lazy<ThreadLocal<Lock>> = Lazy::new(ThreadLocal::new);
+use self::sealed::Sealed;
+
+mod sealed {
+	use super::ThreadKey;
+	pub trait Sealed {}
+	impl Sealed for ThreadKey {}
+	impl Sealed for &mut ThreadKey {}
+}
+
+static KEY: Lazy<ThreadLocal<AtomicLock>> = Lazy::new(ThreadLocal::new);
 
 /// The key for the current thread.
 ///
@@ -15,15 +24,26 @@ pub type ThreadKey = Key<'static>;
 
 /// A dumb lock that's just a wrapper for an [`AtomicBool`].
 #[derive(Debug, Default)]
-pub struct Lock {
+pub struct AtomicLock {
 	is_locked: AtomicBool,
 }
 
 pub struct Key<'a> {
 	phantom: PhantomData<*const ()>, // implement !Send and !Sync
-	lock: &'a Lock,
+	lock: &'a AtomicLock,
 }
 
+/// Allows the type to be used as a key for a lock
+///
+/// # Safety
+///
+/// Only one value which implements this trait may be allowed to exist at a
+/// time. Creating a new `Keyable` value requires making any other `Keyable`
+/// values invalid.
+pub unsafe trait Keyable: Sealed {}
+unsafe impl Keyable for ThreadKey {}
+unsafe impl Keyable for &mut ThreadKey {}
+
 impl Debug for ThreadKey {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 		write!(f, "ThreadKey")
@@ -56,8 +76,8 @@ impl ThreadKey {
 	}
 }
 
-impl Lock {
-	/// Create a new unlocked `Lock`.
+impl AtomicLock {
+	/// Create a new unlocked `AtomicLock`.
 	#[must_use]
 	pub const fn new() -> Self {
 		Self {
@@ -70,7 +90,7 @@ impl Lock {
 		self.is_locked.load(Ordering::Relaxed)
 	}
 
-	/// Attempt to lock the `Lock`.
+	/// Attempt to lock the `AtomicLock`.
 	///
 	/// If the lock is already locked, then this'll return false. If it is
 	/// unlocked, it'll return true. If the lock is currently unlocked, then it
@@ -86,7 +106,7 @@ impl Lock {
 		})
 	}
 
-	/// Forcibly unlocks the `Lock`.
+	/// Forcibly unlocks the `AtomicLock`.
 	///
 	/// # Safety
 	///
diff --git a/src/lib.rs b/src/lib.rs
index 457fae0..35f7bb4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,11 +4,11 @@
 #![allow(clippy::declare_interior_mutable_const)]
 
 mod guard;
-mod lock;
+mod key;
 mod lockable;
 pub mod mutex;
 
 pub use guard::LockGuard;
-pub use lock::{Key, ThreadKey};
+pub use key::{Key, ThreadKey};
 pub use lockable::Lockable;
 pub use mutex::Mutex;
diff --git a/src/lockable.rs b/src/lockable.rs
index b6ba7d6..cda4e13 100644
--- a/src/lockable.rs
+++ b/src/lockable.rs
@@ -31,6 +31,8 @@ pub unsafe trait Lockable<'a>: sealed::Sealed {
 	/// * Use this without ownership or mutable access to the [`ThreadKey`],
 	/// which should last as long as the return value is alive.
 	/// * Call this on multiple locks without unlocking first.
+	///
+	/// [`ThreadKey`]: `crate::key::ThreadKey`
 	unsafe fn lock(&'a self) -> Self::Output;
 
 	/// Attempt to lock without blocking.
@@ -42,6 +44,8 @@ pub unsafe trait Lockable<'a>: sealed::Sealed {
 	/// It is undefined behavior to use this without ownership or mutable
 	/// access to the [`ThreadKey`], which should last as long as the return
 	/// value is alive.
+	///
+	/// [`ThreadKey`]: `crate::key::ThreadKey`
 	unsafe fn try_lock(&'a self) -> Option<Self::Output>;
 
 	/// Release the lock
diff --git a/src/mutex.rs b/src/mutex.rs
index e806fac..28e1786 100644
--- a/src/mutex.rs
+++ b/src/mutex.rs
@@ -1,8 +1,7 @@
 use std::cell::UnsafeCell;
 use std::ops::{Deref, DerefMut};
 
-use crate::lock::Lock;
-use crate::ThreadKey;
+use crate::key::{AtomicLock, Keyable};
 
 /// A spinning mutex
 pub type SpinLock<T> = Mutex<T, RawSpin>;
@@ -37,11 +36,13 @@ pub unsafe trait RawMutex {
 
 /// A raw mutex which just spins
 pub struct RawSpin {
-	lock: Lock,
+	lock: AtomicLock,
 }
 
 unsafe impl RawMutex for RawSpin {
-	const INIT: Self = Self { lock: Lock::new() };
+	const INIT: Self = Self {
+		lock: AtomicLock::new(),
+	};
 
 	fn lock(&self) {
 		loop {
@@ -124,12 +125,12 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexRef<'a, T, R> {
 ///
 /// [`lock`]: `Mutex::lock`
 /// [`try_lock`]: `Mutex::try_lock`
-pub struct MutexGuard<'a, T: ?Sized + 'a, R: RawMutex = RawSpin> {
+pub struct MutexGuard<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex = RawSpin> {
 	mutex: MutexRef<'a, T, R>,
-	_thread_key: &'a mut ThreadKey,
+	_thread_key: Key,
 }
 
-impl<'a, T: ?Sized + 'a, R: RawMutex> Deref for MutexGuard<'a, T, R> {
+impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> Deref for MutexGuard<'a, T, Key, R> {
 	type Target = T;
 
 	fn deref(&self) -> &Self::Target {
@@ -137,16 +138,16 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> Deref for MutexGuard<'a, T, R> {
 	}
 }
 
-impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexGuard<'a, T, R> {
+impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> DerefMut for MutexGuard<'a, T, Key, R> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
 		&mut self.mutex
 	}
 }
 
-impl<'a, T: ?Sized + 'a, R: RawMutex> MutexGuard<'a, T, R> {
+impl<'a, T: ?Sized + 'a, Key: Keyable, R: RawMutex> MutexGuard<'a, T, Key, R> {
 	/// Create a guard to the given mutex. Undefined if multiple guards to the
 	/// same mutex exist at once.
-	unsafe fn new(mutex: &'a Mutex<T, R>, thread_key: &'a mut ThreadKey) -> Self {
+	const unsafe fn new(mutex: &'a Mutex<T, R>, thread_key: Key) -> Self {
 		Self {
 			mutex: MutexRef(mutex),
 			_thread_key: thread_key,
@@ -196,7 +197,7 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = ThreadKey::lock().unwrap();
 	/// assert_eq!(*mutex.lock(key), 10);
 	/// ```
-	pub fn lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> MutexGuard<'_, T, R> {
+	pub fn lock<'s, 'a: 's, Key: Keyable>(&'s self, key: Key) -> MutexGuard<'_, T, Key, R> {
 		self.raw.lock();
 
 		// safety: we just locked the mutex
@@ -240,7 +241,10 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = ThreadKey::lock().unwrap();
 	/// assert_eq!(*mutex.lock(key), 10);
 	/// ```
-	pub fn try_lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> Option<MutexGuard<'_, T, R>> {
+	pub fn try_lock<'s, 'a: 's, Key: Keyable>(
+		&'s self,
+		key: Key,
+	) -> Option<MutexGuard<'_, T, Key, R>> {
 		if self.raw.try_lock() {
 			// safety: we just locked the mutex
 			Some(unsafe { MutexGuard::new(self, key) })
@@ -281,7 +285,7 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = Mutex::unlock(guard);
 	/// ```
 	#[allow(clippy::missing_const_for_fn)]
-	pub fn unlock(guard: MutexGuard<'_, T, R>) {
+	pub fn unlock(guard: MutexGuard<'_, T, impl Keyable, R>) {
 		drop(guard);
 	}
 }