Implemented LockGuard

4 files changed, 174 insertions, 17 deletions

diff --git a/src/guard.rs b/src/guard.rs
new file mode 100644
index 0000000..ec7020b
--- /dev/null
+++ b/src/guard.rs
@@ -0,0 +1,129 @@
+use std::ops::{Deref, DerefMut};
+
+use crate::{
+	mutex::{MutexRef, RawMutex},
+	Mutex, ThreadKey,
+};
+
+mod sealed {
+	#[allow(clippy::wildcard_imports)]
+	use super::*;
+	pub trait Sealed {}
+	impl<R: RawMutex, T> Sealed for Mutex<R, T> {}
+	impl<'a, A: Lockable<'a>, B: Lockable<'a>> Sealed for (A, B) {}
+}
+
+pub trait Lockable<'a>: sealed::Sealed {
+	/// The output of the lock
+	type Output;
+
+	/// Blocks until the lock is acquired
+	///
+	/// # Safety
+	///
+	/// It is undefined behavior to:
+	/// * Use this without ownership of 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;
+
+	/// Attempt to lock without blocking.
+	///
+	/// Returns `Ok` if successful, `None` otherwise.
+	///
+	/// # Safety
+	///
+	/// It is undefined behavior to use this without ownership of 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
+	fn unlock(guard: Self::Output);
+}
+
+impl<'a, R: RawMutex + 'a, T: 'a> Lockable<'a> for Mutex<R, T> {
+	type Output = MutexRef<'a, R, T>;
+
+	unsafe fn lock(&'a self) -> Self::Output {
+		self.lock_ref()
+	}
+
+	unsafe fn try_lock(&'a self) -> Option<Self::Output> {
+		self.try_lock_ref()
+	}
+
+	fn unlock(guard: Self::Output) {
+		drop(guard);
+	}
+}
+
+impl<'a, A: Lockable<'a>, B: Lockable<'a>> Lockable<'a> for (A, B) {
+	type Output = (A::Output, B::Output);
+
+	unsafe fn lock(&'a self) -> Self::Output {
+		loop {
+			let lock1 = self.0.lock();
+			match self.1.try_lock() {
+				Some(lock2) => return (lock1, lock2),
+				None => A::unlock(lock1),
+			}
+		}
+	}
+
+	unsafe fn try_lock(&'a self) -> Option<Self::Output> {
+		self.0.try_lock().and_then(|guard1| {
+			if let Some(guard2) = self.1.try_lock() {
+				Some((guard1, guard2))
+			} else {
+				A::unlock(guard1);
+				None
+			}
+		})
+	}
+
+	fn unlock(guard: Self::Output) {
+		A::unlock(guard.0);
+		B::unlock(guard.1);
+	}
+}
+
+pub struct LockGuard<'a, L: Lockable<'a>> {
+	guard: L::Output,
+	key: ThreadKey,
+}
+
+// TODO: docs
+impl<'a, L: Lockable<'a>> LockGuard<'a, L> {
+	pub fn lock(lock: &'a L, key: ThreadKey) -> Self {
+		Self {
+			// safety: we have the thread's key
+			guard: unsafe { lock.lock() },
+			key,
+		}
+	}
+
+	pub fn try_lock(lock: &'a L, key: ThreadKey) -> Option<Self> {
+		// safety: we have the thread's key
+		unsafe { lock.try_lock() }.map(|guard| Self { guard, key })
+	}
+
+	#[allow(clippy::missing_const_for_fn)]
+	pub fn unlock(self) -> ThreadKey {
+		L::unlock(self.guard);
+		self.key
+	}
+}
+
+impl<'a, L: Lockable<'a>> Deref for LockGuard<'a, L> {
+	type Target = L::Output;
+
+	fn deref(&self) -> &Self::Target {
+		&self.guard
+	}
+}
+
+impl<'a, L: Lockable<'a>> DerefMut for LockGuard<'a, L> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		&mut self.guard
+	}
+}
diff --git a/src/lib.rs b/src/lib.rs
index 801cbc6..fc18d7d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,12 +10,14 @@ use std::marker::PhantomData;
 use once_cell::sync::Lazy;
 use thread_local::ThreadLocal;
 
+mod guard;
 mod lock;
 mod mutex;
 
 use lock::{Key, Lock};
 use mutex::RawSpin;
 
+pub use guard::{LockGuard, Lockable};
 pub use mutex::{Mutex, MutexGuard};
 /// A spinning mutex
 pub type SpinLock<T> = Mutex<RawSpin, T>;
diff --git a/src/lock.rs b/src/lock.rs
index 101a061..a8662d0 100644
--- a/src/lock.rs
+++ b/src/lock.rs
@@ -17,7 +17,7 @@ pub struct Key<'a> {
 
 impl<'a> Key<'a> {
 	/// Create a key to a lock.
-	const fn new(lock: &'a Lock) -> Self {
+	const unsafe fn new(lock: &'a Lock) -> Self {
 		Self { lock }
 	}
 }
@@ -52,7 +52,8 @@ impl Lock {
 	/// This is not a fair lock. It is not recommended to call this function
 	/// repeatedly in a loop.
 	pub fn try_lock(&self) -> Option<Key> {
-		(!self.is_locked.fetch_or(true, Ordering::Acquire)).then_some(Key::new(self))
+		// safety: we just acquired the lock
+		(!self.is_locked.fetch_or(true, Ordering::Acquire)).then_some(unsafe { Key::new(self) })
 	}
 
 	/// Forcibly unlocks the `Lock`.
diff --git a/src/mutex.rs b/src/mutex.rs
index 6c2e254..e09d47b 100644
--- a/src/mutex.rs
+++ b/src/mutex.rs
@@ -81,7 +81,7 @@ pub struct Mutex<R, T: ?Sized> {
 }
 
 /// A reference to a mutex that unlocks it when dropped
-struct MutexRef<'a, R: RawMutex, T: ?Sized + 'a>(&'a Mutex<R, T>);
+pub struct MutexRef<'a, R: RawMutex, T: ?Sized + 'a>(&'a Mutex<R, T>);
 
 impl<'a, R: RawMutex, T: ?Sized + 'a> Drop for MutexRef<'a, R, T> {
 	fn drop(&mut self) {
@@ -91,6 +91,26 @@ impl<'a, R: RawMutex, T: ?Sized + 'a> Drop for MutexRef<'a, R, T> {
 	}
 }
 
+impl<'a, R: RawMutex, T: ?Sized + 'a> Deref for MutexRef<'a, R, T> {
+	type Target = T;
+
+	fn deref(&self) -> &Self::Target {
+		// safety: this is the only type that can use `value`, and there's
+		//         a reference to this type, so there cannot be any mutable
+		//         references to this value.
+		unsafe { &*self.0.value.get() }
+	}
+}
+
+impl<'a, R: RawMutex, T: ?Sized + 'a> DerefMut for MutexRef<'a, R, T> {
+	fn deref_mut(&mut self) -> &mut Self::Target {
+		// safety: this is the only type that can use `value`, and we have a
+		//         mutable reference to this type, so there cannot be any other
+		//         references to this value.
+		unsafe { &mut *self.0.value.get() }
+	}
+}
+
 /// An RAII implementation of a “scoped lock” of a mutex. When this structure
 /// is dropped (falls out of scope), the lock will be unlocked.
 ///
@@ -107,35 +127,25 @@ impl<'a, R: RawMutex, T: ?Sized + 'a> Deref for MutexGuard<'a, R, T> {
 	type Target = T;
 
 	fn deref(&self) -> &Self::Target {
-		// safety: this is the only type that can use `value`, and there's
-		//         a reference to this type, so there cannot be any mutable
-		//         references to this value.
-		unsafe { &*self.mutex().value.get() }
+		&self.mutex
 	}
 }
 
 impl<'a, R: RawMutex, T: ?Sized + 'a> DerefMut for MutexGuard<'a, R, T> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
-		// safety: this is the only type that can use `value`, and we have a
-		//         mutable reference to this type, so there cannot be any other
-		//         references to this value.
-		unsafe { &mut *self.mutex().value.get() }
+		&mut self.mutex
 	}
 }
 
 impl<'a, R: RawMutex, T: ?Sized + 'a> MutexGuard<'a, R, T> {
-	/// Create a guard to the given mutex
+	/// 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<R, T>, thread_key: ThreadKey) -> Self {
 		Self {
 			mutex: MutexRef(mutex),
 			thread_key,
 		}
 	}
-
-	/// Get a reference to the mutex this is guarding
-	const fn mutex(&self) -> &'a Mutex<R, T> {
-		self.mutex.0
-	}
 }
 
 impl<R: RawMutex, T> Mutex<R, T> {
@@ -193,6 +203,15 @@ impl<R: RawMutex, T: ?Sized> Mutex<R, T> {
 		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.
+	pub(crate) unsafe fn lock_ref(&self) -> MutexRef<'_, R, T> {
+		self.raw.lock();
+
+		MutexRef(self)
+	}
+
 	/// Attempts to lock the `Mutex` without blocking.
 	///
 	/// # Errors
@@ -230,6 +249,12 @@ impl<R: RawMutex, T: ?Sized> Mutex<R, T> {
 		}
 	}
 
+	/// Lock without a [`ThreadKey`]. It is undefined behavior to do this without
+	/// owning the [`ThreadKey`].
+	pub(crate) unsafe fn try_lock_ref(&self) -> Option<MutexRef<'_, R, T>> {
+		self.raw.try_lock().then_some(MutexRef(self))
+	}
+
 	/// Forcibly unlocks the `Lock`.
 	///
 	/// # Safety