9 files changed, 91 insertions, 127 deletions

diff --git a/README.md b/README.md
index b6bbd61..aa73f09 100644
--- a/README.md
+++ b/README.md
@@ -39,7 +39,7 @@ let data = data.lock(key);
 println!("{}", *data);
 ```
 
-Unlocking a mutex requires a `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. To get the key back out of a mutex, it must be unlocked. This means that only one thing can be locked at a time.
+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.
 
 To lock multiple mutexes at a time, create a `LockGuard`.
 
diff --git a/examples/basic.rs b/examples/basic.rs
index 535b80a..8dfca84 100644
--- a/examples/basic.rs
+++ b/examples/basic.rs
@@ -1,22 +1,22 @@
 use std::thread;
 
-use happylock::mutex::Mutex;
+use happylock::mutex::{Mutex, SpinLock};
 use happylock::ThreadKey;
 
 const N: usize = 10;
 
-static DATA: Mutex<i32> = Mutex::new(0);
+static DATA: SpinLock<i32> = Mutex::new(0);
 
 fn main() {
 	for _ in 0..N {
 		thread::spawn(move || {
-			let key = ThreadKey::lock().unwrap();
-			let mut data = DATA.lock(key);
+			let mut key = ThreadKey::lock().unwrap();
+			let mut data = DATA.lock(&mut key);
 			*data += 1;
 		});
 	}
 
-	let key = ThreadKey::lock().unwrap();
-	let data = DATA.lock(key);
+	let mut key = ThreadKey::lock().unwrap();
+	let data = DATA.lock(&mut key);
 	println!("{}", *data);
 }
diff --git a/examples/double_mutex.rs b/examples/double_mutex.rs
index 76f3294..033bfed 100644
--- a/examples/double_mutex.rs
+++ b/examples/double_mutex.rs
@@ -1,27 +1,27 @@
 use std::thread;
 
-use happylock::mutex::Mutex;
+use happylock::mutex::{Mutex, SpinLock};
 use happylock::{LockGuard, ThreadKey};
 
 const N: usize = 10;
 
-static DATA_1: Mutex<i32> = Mutex::new(0);
-static DATA_2: Mutex<String> = Mutex::new(String::new());
+static DATA_1: SpinLock<i32> = Mutex::new(0);
+static DATA_2: SpinLock<String> = Mutex::new(String::new());
 
 fn main() {
 	for _ in 0..N {
 		thread::spawn(move || {
-			let key = ThreadKey::lock().unwrap();
+			let mut key = ThreadKey::lock().unwrap();
 			let data = (&DATA_1, &DATA_2);
-			let mut guard = LockGuard::lock(&data, key);
+			let mut guard = LockGuard::lock(&data, &mut key);
 			*guard.1 = (100 - *guard.0).to_string();
 			*guard.0 += 1;
 		});
 	}
 
-	let key = ThreadKey::lock().unwrap();
+	let mut key = ThreadKey::lock().unwrap();
 	let data = (&DATA_1, &DATA_2);
-	let data = LockGuard::lock(&data, key);
+	let data = LockGuard::lock(&data, &mut key);
 	println!("{}", *data.0);
 	println!("{}", *data.1);
 }
diff --git a/examples/list.rs b/examples/list.rs
index 904d9f2..c73b54e 100644
--- a/examples/list.rs
+++ b/examples/list.rs
@@ -1,11 +1,11 @@
 use std::thread;
 
-use happylock::mutex::Mutex;
+use happylock::mutex::{Mutex, SpinLock};
 use happylock::{LockGuard, ThreadKey};
 
 const N: usize = 10;
 
-static DATA: [Mutex<usize>; 6] = [
+static DATA: [SpinLock<usize>; 6] = [
 	Mutex::new(0),
 	Mutex::new(1),
 	Mutex::new(2),
@@ -14,39 +14,38 @@ static DATA: [Mutex<usize>; 6] = [
 	Mutex::new(5),
 ];
 
-static SEED: Mutex<u32> = Mutex::new(42);
+static SEED: SpinLock<u32> = Mutex::new(42);
 
-fn random(key: ThreadKey) -> (ThreadKey, usize) {
+fn random(key: &mut ThreadKey) -> usize {
 	let mut seed = SEED.lock(key);
 	let x = *seed;
 	let x = x ^ (x << 13);
 	let x = x ^ (x >> 17);
 	let x = x ^ (x << 5);
 	*seed = x;
-	(Mutex::unlock(seed), x as usize)
+	x as usize
 }
 
 fn main() {
 	for _ in 0..N {
 		thread::spawn(move || {
 			let mut key = ThreadKey::lock().unwrap();
-			let mut rand;
 			let mut data = Vec::new();
 			for _ in 0..3 {
-				(key, rand) = random(key);
+				let rand = random(&mut key);
 				data.push(&DATA[rand % 6]);
 			}
 
 			let data = [data[0], data[1], data[2]];
-			let mut guard = LockGuard::lock(&data, key);
+			let mut guard = LockGuard::lock(&data, &mut key);
 			*guard[0] += *guard[1];
 			*guard[1] += *guard[2];
 			*guard[2] += *guard[0];
 		});
 	}
 
-	let key = ThreadKey::lock().unwrap();
-	let data = LockGuard::lock(&DATA, key);
+	let mut key = ThreadKey::lock().unwrap();
+	let data = LockGuard::lock(&DATA, &mut key);
 	for val in &*data {
 		println!("{}", **val);
 	}
diff --git a/src/guard.rs b/src/guard.rs
index 5f847d5..3c1b636 100644
--- a/src/guard.rs
+++ b/src/guard.rs
@@ -5,37 +5,33 @@ use crate::{lockable::Lockable, ThreadKey};
 /// A guard for a generic [`Lockable`] type.
 pub struct LockGuard<'a, L: Lockable<'a>> {
 	guard: L::Output,
-	key: ThreadKey,
+	_key: &'a mut ThreadKey,
 }
 
 impl<'a, L: Lockable<'a>> LockGuard<'a, L> {
 	/// Locks the lockable type and returns a guard that can be used to access
 	/// the underlying data.
-	pub fn lock(lock: &'a L, key: ThreadKey) -> Self {
+	pub fn lock(lock: &'a L, key: &'a mut ThreadKey) -> Self {
 		Self {
 			// safety: we have the thread's key
 			guard: unsafe { lock.lock() },
-			key,
+			_key: key,
 		}
 	}
 
 	/// 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: ThreadKey) -> Result<Self, ThreadKey> {
+	pub fn try_lock(lock: &'a L, key: &'a mut ThreadKey) -> Option<Self> {
 		// safety: we have the thread's key
-		match unsafe { lock.try_lock() } {
-			Some(guard) => Ok(Self { guard, key }),
-			None => Err(key),
-		}
+		unsafe { lock.try_lock() }.map(|guard| Self { guard, _key: key })
 	}
 
 	/// Unlocks the underlying lockable data type, returning the key that's
 	/// associated with it.
 	#[allow(clippy::missing_const_for_fn)]
-	pub fn unlock(self) -> ThreadKey {
+	pub fn unlock(self) {
 		L::unlock(self.guard);
-		self.key
 	}
 }
 
diff --git a/src/lib.rs b/src/lib.rs
index dd7354c..457fae0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,60 +3,12 @@
 #![allow(clippy::module_name_repetitions)]
 #![allow(clippy::declare_interior_mutable_const)]
 
-use std::any::type_name;
-use std::fmt::{self, Debug};
-use std::marker::PhantomData;
-
-use once_cell::sync::Lazy;
-use thread_local::ThreadLocal;
-
 mod guard;
 mod lock;
 mod lockable;
 pub mod mutex;
 
-use lock::{Key, Lock};
-
 pub use guard::LockGuard;
+pub use lock::{Key, ThreadKey};
 pub use lockable::Lockable;
 pub use mutex::Mutex;
-
-static KEY: Lazy<ThreadLocal<Lock>> = Lazy::new(ThreadLocal::new);
-
-/// The key for the current thread.
-///
-/// Only one of these exist per thread. To get the current thread's key, call
-/// [`ThreadKey::lock`].
-pub struct ThreadKey {
-	phantom: PhantomData<*const ()>, // implement !Send and !Sync
-	_key: Key<'static>,
-}
-
-impl Debug for ThreadKey {
-	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-		type_name::<Self>().fmt(f)
-	}
-}
-
-impl ThreadKey {
-	/// Get the current thread's `ThreadKey`, if it's not already taken.
-	///
-	/// The first time this is called, it will successfully return a
-	/// `ThreadKey`. However, future calls to this function will return
-	/// [`None`], unless the key is dropped or unlocked first.
-	#[must_use]
-	pub fn lock() -> Option<Self> {
-		KEY.get_or_default().try_lock().map(|key| Self {
-			phantom: PhantomData,
-			_key: key,
-		})
-	}
-
-	/// Unlocks the `ThreadKey`.
-	///
-	/// After this method is called, a call to [`ThreadKey::lock`] will return
-	/// this `ThreadKey`.
-	pub fn unlock(key: Self) {
-		drop(key);
-	}
-}
diff --git a/src/lock.rs b/src/lock.rs
index f310099..942474a 100644
--- a/src/lock.rs
+++ b/src/lock.rs
@@ -1,34 +1,61 @@
+use std::fmt::{self, Debug};
+use std::marker::PhantomData;
 use std::sync::atomic::{AtomicBool, Ordering};
 
+use once_cell::sync::Lazy;
+use thread_local::ThreadLocal;
+
+static KEY: Lazy<ThreadLocal<Lock>> = Lazy::new(ThreadLocal::new);
+
+/// The key for the current thread.
+///
+/// Only one of these exist per thread. To get the current thread's key, call
+/// [`ThreadKey::lock`]. If the `ThreadKey` is dropped, it can be reobtained.
+pub type ThreadKey = Key<'static>;
+
 /// A dumb lock that's just a wrapper for an [`AtomicBool`].
 #[derive(Debug, Default)]
 pub struct Lock {
 	is_locked: AtomicBool,
 }
 
-/// A key for a lock.
-///
-/// This key is needed in order to unlock a [`Lock`]. The [`Lock`] is
-/// automatically unlocked if this key is dropped.
-#[derive(Debug)]
 pub struct Key<'a> {
+	phantom: PhantomData<*const ()>, // implement !Send and !Sync
 	lock: &'a Lock,
 }
 
-impl<'a> Key<'a> {
-	/// Create a key to a lock.
-	const unsafe fn new(lock: &'a Lock) -> Self {
-		Self { lock }
+impl Debug for ThreadKey {
+	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+		write!(f, "ThreadKey")
 	}
 }
 
 impl<'a> Drop for Key<'a> {
 	fn drop(&mut self) {
-		// safety: this key will soon be destroyed
 		unsafe { self.lock.force_unlock() }
 	}
 }
 
+impl ThreadKey {
+	/// Get the current thread's `ThreadKey`, if it's not already taken.
+	///
+	/// The first time this is called, it will successfully return a
+	/// `ThreadKey`. However, future calls to this function will return
+	/// [`None`], unless the key is dropped or unlocked first.
+	#[must_use]
+	pub fn lock() -> Option<Self> {
+		KEY.get_or_default().try_lock()
+	}
+
+	/// Unlocks the `ThreadKey`.
+	///
+	/// After this method is called, a call to [`ThreadKey::lock`] will return
+	/// this `ThreadKey`.
+	pub fn unlock(key: Self) {
+		drop(key);
+	}
+}
+
 impl Lock {
 	/// Create a new unlocked `Lock`.
 	#[must_use]
@@ -53,7 +80,10 @@ impl Lock {
 	/// repeatedly in a loop.
 	pub fn try_lock(&self) -> Option<Key> {
 		// safety: we just acquired the lock
-		(!self.is_locked.swap(true, Ordering::Acquire)).then_some(unsafe { Key::new(self) })
+		(!self.is_locked.swap(true, Ordering::Acquire)).then_some(Key {
+			phantom: PhantomData,
+			lock: self,
+		})
 	}
 
 	/// Forcibly unlocks the `Lock`.
@@ -65,9 +95,4 @@ impl Lock {
 	pub unsafe fn force_unlock(&self) {
 		self.is_locked.store(false, Ordering::Release);
 	}
-
-	/// Unlock the lock, consuming its key.
-	pub fn unlock(key: Key) {
-		drop(key);
-	}
 }
diff --git a/src/lockable.rs b/src/lockable.rs
index 9b754a3..b6ba7d6 100644
--- a/src/lockable.rs
+++ b/src/lockable.rs
@@ -28,8 +28,8 @@ pub unsafe trait Lockable<'a>: sealed::Sealed {
 	/// # Safety
 	///
 	/// It is undefined behavior to:
-	/// * Use this without ownership of the [`ThreadKey`], which should last as
-	/// long as the return value is alive.
+	/// * 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.
 	unsafe fn lock(&'a self) -> Self::Output;
 
@@ -39,8 +39,9 @@ pub unsafe trait Lockable<'a>: sealed::Sealed {
 	///
 	/// # Safety
 	///
-	/// It is undefined behavior to use this without ownership of the
-	/// [`ThreadKey`], which should last as long as the return value is alive.
+	/// 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.
 	unsafe fn try_lock(&'a self) -> Option<Self::Output>;
 
 	/// Release the lock
diff --git a/src/mutex.rs b/src/mutex.rs
index fa7bf36..e806fac 100644
--- a/src/mutex.rs
+++ b/src/mutex.rs
@@ -5,7 +5,7 @@ use crate::lock::Lock;
 use crate::ThreadKey;
 
 /// A spinning mutex
-pub type SpinLock<T> = Mutex<RawSpin, T>;
+pub type SpinLock<T> = Mutex<T, RawSpin>;
 
 /// Implements a raw C-like mutex.
 ///
@@ -78,19 +78,16 @@ unsafe impl RawMutex for RawSpin {
 ///
 /// Locking the mutex on a thread that already locked it is impossible, due to
 /// the requirement of the [`ThreadKey`]. Therefore, this will never deadlock.
-/// When the [`MutexGuard`] is dropped, the [`ThreadKey`] can be reobtained by
-/// calling [`ThreadKey::lock`]. You can also get it by calling
-/// [`Mutex::unlock`].
 ///
 /// [`lock`]: `Mutex::lock`
 /// [`try_lock`]: `Mutex::try_lock`
-pub struct Mutex<T: ?Sized, R = RawSpin> {
+pub struct Mutex<T: ?Sized, R> {
 	raw: R,
 	value: UnsafeCell<T>,
 }
 
 /// A reference to a mutex that unlocks it when dropped
-pub struct MutexRef<'a, T: ?Sized + 'a, R: RawMutex = RawSpin>(&'a Mutex<T, R>);
+pub struct MutexRef<'a, T: ?Sized + 'a, R: RawMutex>(&'a Mutex<T, R>);
 
 impl<'a, T: ?Sized + 'a, R: RawMutex> Drop for MutexRef<'a, T, R> {
 	fn drop(&mut self) {
@@ -129,7 +126,7 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexRef<'a, T, R> {
 /// [`try_lock`]: `Mutex::try_lock`
 pub struct MutexGuard<'a, T: ?Sized + 'a, R: RawMutex = RawSpin> {
 	mutex: MutexRef<'a, T, R>,
-	thread_key: ThreadKey,
+	_thread_key: &'a mut ThreadKey,
 }
 
 impl<'a, T: ?Sized + 'a, R: RawMutex> Deref for MutexGuard<'a, T, R> {
@@ -149,10 +146,10 @@ impl<'a, T: ?Sized + 'a, R: RawMutex> DerefMut for MutexGuard<'a, T, R> {
 impl<'a, T: ?Sized + 'a, R: RawMutex> MutexGuard<'a, T, R> {
 	/// Create a guard to the given mutex. Undefined if multiple guards to the
 	/// same mutex exist at once.
-	const unsafe fn new(mutex: &'a Mutex<T, R>, thread_key: ThreadKey) -> Self {
+	unsafe fn new(mutex: &'a Mutex<T, R>, thread_key: &'a mut ThreadKey) -> Self {
 		Self {
 			mutex: MutexRef(mutex),
-			thread_key,
+			_thread_key: thread_key,
 		}
 	}
 }
@@ -199,16 +196,16 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = ThreadKey::lock().unwrap();
 	/// assert_eq!(*mutex.lock(key), 10);
 	/// ```
-	pub fn lock(&self, key: ThreadKey) -> MutexGuard<'_, T, R> {
+	pub fn lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> MutexGuard<'_, T, R> {
 		self.raw.lock();
 
 		// safety: we just locked the mutex
 		unsafe { MutexGuard::new(self, key) }
 	}
 
-	/// Lock without a [`ThreadKey`]. You must own the [`ThreadKey`] as long as
-	/// the [`MutexRef`] is alive. This may cause deadlock if called multiple
-	/// times without unlocking first.
+	/// Lock without a [`ThreadKey`]. You must mutually own the [`ThreadKey`] as
+	/// long as the [`MutexRef`] is alive. This may cause deadlock if called
+	/// multiple times without unlocking first.
 	pub(crate) unsafe fn lock_ref(&self) -> MutexRef<'_, T, R> {
 		self.raw.lock();
 
@@ -243,12 +240,12 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = ThreadKey::lock().unwrap();
 	/// assert_eq!(*mutex.lock(key), 10);
 	/// ```
-	pub fn try_lock(&self, key: ThreadKey) -> Result<MutexGuard<'_, T, R>, ThreadKey> {
+	pub fn try_lock<'s, 'a: 's>(&'s self, key: &'a mut ThreadKey) -> Option<MutexGuard<'_, T, R>> {
 		if self.raw.try_lock() {
 			// safety: we just locked the mutex
-			Ok(unsafe { MutexGuard::new(self, key) })
+			Some(unsafe { MutexGuard::new(self, key) })
 		} else {
-			Err(key)
+			None
 		}
 	}
 
@@ -270,11 +267,6 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 
 	/// Consumes the [`MutexGuard`], and consequently unlocks its `Mutex`.
 	///
-	/// This returns the [`ThreadKey`] that was used to lock the `Mutex`, which
-	/// means that [`ThreadKey::lock`] does not need to be called, and will in
-	/// fact return [`None`] if the [`ThreadKey`] returned by this function is
-	/// not dropped.
-	///
 	/// # Examples
 	///
 	/// ```
@@ -289,9 +281,8 @@ impl<T: ?Sized, R: RawMutex> Mutex<T, R> {
 	/// let key = Mutex::unlock(guard);
 	/// ```
 	#[allow(clippy::missing_const_for_fn)]
-	#[must_use]
-	pub fn unlock(guard: MutexGuard<'_, T, R>) -> ThreadKey {
-		guard.thread_key
+	pub fn unlock(guard: MutexGuard<'_, T, R>) {
+		drop(guard);
 	}
 }