Keyable

9 files changed, 71 insertions, 43 deletions

diff --git a/README.md b/README.md
index 9ee6718..18c07c0 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@ let data = data.lock(&mut key);
 println!("{}", *data);
 ```
 
-Unlocking a mutex requires a mutable reference to `ThreadKey`. Each thread will be allowed to have one key at a time, but no more than that. The `ThreadKey` type is not cloneable or copyable. This means that only one thing can be locked at a time.
+Unlocking a mutex requires a `ThreadKey` or a mutable reference to `ThreadKey`. Each thread will be allowed to have one key at a time, but no more than that. The `ThreadKey` type is not cloneable or copyable. This means that only one thing can be locked at a time.
 
 To lock multiple mutexes at a time, create a `LockGuard`.
 
@@ -66,7 +66,7 @@ println!("{}", *data.1);
 
 ## Performance
 
-The `ThreadKey` is a mostly-zero cost abstraction. It doesn't use any memory, and it doesn't really exist at run-time. The only cost comes from calling `ThreadKey::lock()`, because the function has to ensure that the key hasn't already been taken. Dropping the key will also have a small cost.
+The `ThreadKey` is a mostly-zero cost abstraction. It doesn't use any memory, and it doesn't really exist at run-time. The only cost comes from calling `ThreadKey::lock()`, because the function has to ensure at runtime that the key hasn't already been taken. Dropping the key will also have a small cost.
 
 The real performance cost comes from the fact that the sets of multiple locks must be atomic. The problem is that this library must iterate through the list of locks, and not complete until every single one of them is unlocked. This also means that attempting to lock multiple mutexes gives you a lower chance of ever running. Only one needs to be locked for the operation to need a reset. This problem can be prevented by not doing that in your code. Resources should be obtained in the same order on every thread.
 
@@ -80,10 +80,10 @@ There might be some promise in trying to prevent circular wait. There could be a
 
 Currently, the mutex is implemented using a spinlock. We need to not do that. We could use parking lot, or mutexes from the operating system.
 
-A more fair system for getting sets locks would help, but I have no clue what that will look like.
+A more fair system for getting sets locks would help, but I have no clue what that looks like.
 
 A read-write lock would be very useful here, and maybe condvars?
 
 Personally, I don't like mutex poisoning, but maybe it can be worked into the library if you're into that sort of thing.
 
-More types should be lockable using a `LockGuard`.
\ No newline at end of file
+More types might be lockable using a `LockGuard`.
\ No newline at end of file
diff --git a/examples/basic.rs b/examples/basic.rs
index 8dfca84..87e7a60 100644
--- a/examples/basic.rs
+++ b/examples/basic.rs
@@ -10,13 +10,13 @@ static DATA: SpinLock<i32> = Mutex::new(0);
 fn main() {
 	for _ in 0..N {
 		thread::spawn(move || {
-			let mut key = ThreadKey::lock().unwrap();
-			let mut data = DATA.lock(&mut key);
+			let key = ThreadKey::lock().unwrap();
+			let mut data = DATA.lock(key);
 			*data += 1;
 		});
 	}
 
-	let mut key = ThreadKey::lock().unwrap();
-	let data = DATA.lock(&mut key);
+	let key = ThreadKey::lock().unwrap();
+	let data = DATA.lock(key);
 	println!("{}", *data);
 }
diff --git a/examples/double_mutex.rs b/examples/double_mutex.rs
index 033bfed..18460e4 100644
--- a/examples/double_mutex.rs
+++ b/examples/double_mutex.rs
@@ -11,17 +11,17 @@ static DATA_2: SpinLock<String> = Mutex::new(String::new());
 fn main() {
 	for _ in 0..N {
 		thread::spawn(move || {
-			let mut key = ThreadKey::lock().unwrap();
+			let key = ThreadKey::lock().unwrap();
 			let data = (&DATA_1, &DATA_2);
-			let mut guard = LockGuard::lock(&data, &mut key);
+			let mut guard = LockGuard::lock(&data, key);
 			*guard.1 = (100 - *guard.0).to_string();
 			*guard.0 += 1;
 		});
 	}
 
-	let mut key = ThreadKey::lock().unwrap();
+	let key = ThreadKey::lock().unwrap();
 	let data = (&DATA_1, &DATA_2);
-	let data = LockGuard::lock(&data, &mut key);
+	let data = LockGuard::lock(&data, key);
 	println!("{}", *data.0);
 	println!("{}", *data.1);
 }
diff --git a/examples/list.rs b/examples/list.rs
index c73b54e..448f70a 100644
--- a/examples/list.rs
+++ b/examples/list.rs
@@ -37,7 +37,7 @@ fn main() {
 			}
 
 			let data = [data[0], data[1], data[2]];
-			let mut guard = LockGuard::lock(&data, &mut key);
+			let mut guard = LockGuard::lock(&data, key);
 			*guard[0] += *guard[1];
 			*guard[1] += *guard[2];
 			*guard[2] += *guard[0];
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);
 	}
 }