Scoped lock API

3 files changed, 136 insertions, 99 deletions

diff --git a/src/poisonable/error.rs b/src/poisonable/error.rs
index bff011d..b69df5d 100644
--- a/src/poisonable/error.rs
+++ b/src/poisonable/error.rs
@@ -109,7 +109,7 @@ impl<Guard> PoisonError<Guard> {
 	///
 	/// let key = ThreadKey::get().unwrap();
 	/// let p_err = mutex.lock(key).unwrap_err();
-	/// let data: &PoisonGuard<_, _> = p_err.get_ref();
+	/// let data: &PoisonGuard<_> = p_err.get_ref();
 	/// println!("recovered {} items", data.len());
 	/// ```
 	#[must_use]
@@ -154,7 +154,7 @@ impl<Guard> PoisonError<Guard> {
 
 #[mutants::skip]
 #[cfg(not(tarpaulin_include))]
-impl<G, Key> fmt::Debug for TryLockPoisonableError<'_, '_, G, Key> {
+impl<G> fmt::Debug for TryLockPoisonableError<'_, G> {
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 		match *self {
 			Self::Poisoned(..) => "Poisoned(..)".fmt(f),
@@ -163,7 +163,7 @@ impl<G, Key> fmt::Debug for TryLockPoisonableError<'_, '_, G, Key> {
 	}
 }
 
-impl<G, Key> fmt::Display for TryLockPoisonableError<'_, '_, G, Key> {
+impl<G> fmt::Display for TryLockPoisonableError<'_, G> {
 	#[cfg_attr(test, mutants::skip)]
 	#[cfg(not(tarpaulin_include))]
 	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -175,12 +175,10 @@ impl<G, Key> fmt::Display for TryLockPoisonableError<'_, '_, G, Key> {
 	}
 }
 
-impl<G, Key> Error for TryLockPoisonableError<'_, '_, G, Key> {}
+impl<G> Error for TryLockPoisonableError<'_, G> {}
 
-impl<'flag, 'key, G, Key> From<PoisonError<PoisonGuard<'flag, 'key, G, Key>>>
-	for TryLockPoisonableError<'flag, 'key, G, Key>
-{
-	fn from(value: PoisonError<PoisonGuard<'flag, 'key, G, Key>>) -> Self {
+impl<'flag, G> From<PoisonError<PoisonGuard<'flag, G>>> for TryLockPoisonableError<'flag, G> {
+	fn from(value: PoisonError<PoisonGuard<'flag, G>>) -> Self {
 		Self::Poisoned(value)
 	}
 }
diff --git a/src/poisonable/guard.rs b/src/poisonable/guard.rs
index 3f85d25..b887e2d 100644
--- a/src/poisonable/guard.rs
+++ b/src/poisonable/guard.rs
@@ -3,8 +3,6 @@ use std::hash::Hash;
 use std::marker::PhantomData;
 use std::ops::{Deref, DerefMut};
 
-use crate::Keyable;
-
 use super::{PoisonFlag, PoisonGuard, PoisonRef};
 
 impl<'a, Guard> PoisonRef<'a, Guard> {
@@ -28,26 +26,6 @@ impl<Guard> Drop for PoisonRef<'_, Guard> {
 	}
 }
 
-impl<Guard: PartialEq> PartialEq for PoisonRef<'_, Guard> {
-	fn eq(&self, other: &Self) -> bool {
-		self.guard.eq(&other.guard)
-	}
-}
-
-impl<Guard: PartialOrd> PartialOrd for PoisonRef<'_, Guard> {
-	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-		self.guard.partial_cmp(&other.guard)
-	}
-}
-
-impl<Guard: Eq> Eq for PoisonRef<'_, Guard> {}
-
-impl<Guard: Ord> Ord for PoisonRef<'_, Guard> {
-	fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-		self.guard.cmp(&other.guard)
-	}
-}
-
 #[mutants::skip] // hashing involves RNG and is hard to test
 #[cfg(not(tarpaulin_include))]
 impl<Guard: Hash> Hash for PoisonRef<'_, Guard> {
@@ -96,37 +74,9 @@ impl<Guard> AsMut<Guard> for PoisonRef<'_, Guard> {
 	}
 }
 
-#[mutants::skip] // it's hard to get two guards safely
-#[cfg(not(tarpaulin_include))]
-impl<Guard: PartialEq, Key: Keyable> PartialEq for PoisonGuard<'_, '_, Guard, Key> {
-	fn eq(&self, other: &Self) -> bool {
-		self.guard.eq(&other.guard)
-	}
-}
-
-#[mutants::skip] // it's hard to get two guards safely
-#[cfg(not(tarpaulin_include))]
-impl<Guard: PartialOrd, Key: Keyable> PartialOrd for PoisonGuard<'_, '_, Guard, Key> {
-	fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
-		self.guard.partial_cmp(&other.guard)
-	}
-}
-
-#[mutants::skip] // it's hard to get two guards safely
-#[cfg(not(tarpaulin_include))]
-impl<Guard: Eq, Key: Keyable> Eq for PoisonGuard<'_, '_, Guard, Key> {}
-
-#[mutants::skip] // it's hard to get two guards safely
-#[cfg(not(tarpaulin_include))]
-impl<Guard: Ord, Key: Keyable> Ord for PoisonGuard<'_, '_, Guard, Key> {
-	fn cmp(&self, other: &Self) -> std::cmp::Ordering {
-		self.guard.cmp(&other.guard)
-	}
-}
-
 #[mutants::skip] // hashing involves RNG and is hard to test
 #[cfg(not(tarpaulin_include))]
-impl<Guard: Hash, Key: Keyable> Hash for PoisonGuard<'_, '_, Guard, Key> {
+impl<Guard: Hash> Hash for PoisonGuard<'_, Guard> {
 	fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
 		self.guard.hash(state)
 	}
@@ -134,19 +84,19 @@ impl<Guard: Hash, Key: Keyable> Hash for PoisonGuard<'_, '_, Guard, Key> {
 
 #[mutants::skip]
 #[cfg(not(tarpaulin_include))]
-impl<Guard: Debug, Key: Keyable> Debug for PoisonGuard<'_, '_, Guard, Key> {
+impl<Guard: Debug> Debug for PoisonGuard<'_, Guard> {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		Debug::fmt(&self.guard, f)
 	}
 }
 
-impl<Guard: Display, Key: Keyable> Display for PoisonGuard<'_, '_, Guard, Key> {
+impl<Guard: Display> Display for PoisonGuard<'_, Guard> {
 	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
 		Display::fmt(&self.guard, f)
 	}
 }
 
-impl<T, Guard: Deref<Target = T>, Key: Keyable> Deref for PoisonGuard<'_, '_, Guard, Key> {
+impl<T, Guard: Deref<Target = T>> Deref for PoisonGuard<'_, Guard> {
 	type Target = T;
 
 	fn deref(&self) -> &Self::Target {
@@ -155,20 +105,20 @@ impl<T, Guard: Deref<Target = T>, Key: Keyable> Deref for PoisonGuard<'_, '_, Gu
 	}
 }
 
-impl<T, Guard: DerefMut<Target = T>, Key: Keyable> DerefMut for PoisonGuard<'_, '_, Guard, Key> {
+impl<T, Guard: DerefMut<Target = T>> DerefMut for PoisonGuard<'_, Guard> {
 	fn deref_mut(&mut self) -> &mut Self::Target {
 		#[allow(clippy::explicit_auto_deref)] // fixing this results in a compiler error
 		&mut *self.guard.guard
 	}
 }
 
-impl<Guard, Key: Keyable> AsRef<Guard> for PoisonGuard<'_, '_, Guard, Key> {
+impl<Guard> AsRef<Guard> for PoisonGuard<'_, Guard> {
 	fn as_ref(&self) -> &Guard {
 		&self.guard.guard
 	}
 }
 
-impl<Guard, Key: Keyable> AsMut<Guard> for PoisonGuard<'_, '_, Guard, Key> {
+impl<Guard> AsMut<Guard> for PoisonGuard<'_, Guard> {
 	fn as_mut(&mut self) -> &mut Guard {
 		&mut self.guard.guard
 	}
diff --git a/src/poisonable/poisonable.rs b/src/poisonable/poisonable.rs
index 2dac4bb..efe4ed0 100644
--- a/src/poisonable/poisonable.rs
+++ b/src/poisonable/poisonable.rs
@@ -1,10 +1,9 @@
-use std::marker::PhantomData;
 use std::panic::{RefUnwindSafe, UnwindSafe};
 
 use crate::lockable::{
 	Lockable, LockableGetMut, LockableIntoInner, OwnedLockable, RawLock, Sharable,
 };
-use crate::Keyable;
+use crate::{Keyable, ThreadKey};
 
 use super::{
 	PoisonError, PoisonFlag, PoisonGuard, PoisonRef, PoisonResult, Poisonable,
@@ -49,6 +48,11 @@ unsafe impl<L: Lockable> Lockable for Poisonable<L> {
 	where
 		Self: 'g;
 
+	type DataMut<'a>
+		= PoisonResult<L::DataMut<'a>>
+	where
+		Self: 'a;
+
 	fn get_ptrs<'a>(&'a self, ptrs: &mut Vec<&'a dyn RawLock>) {
 		self.inner.get_ptrs(ptrs)
 	}
@@ -62,6 +66,14 @@ unsafe impl<L: Lockable> Lockable for Poisonable<L> {
 			Ok(ref_guard)
 		}
 	}
+
+	unsafe fn data_mut(&self) -> Self::DataMut<'_> {
+		if self.is_poisoned() {
+			Err(PoisonError::new(self.inner.data_mut()))
+		} else {
+			Ok(self.inner.data_mut())
+		}
+	}
 }
 
 unsafe impl<L: Sharable> Sharable for Poisonable<L> {
@@ -70,6 +82,11 @@ unsafe impl<L: Sharable> Sharable for Poisonable<L> {
 	where
 		Self: 'g;
 
+	type DataRef<'a>
+		= PoisonResult<L::DataRef<'a>>
+	where
+		Self: 'a;
+
 	unsafe fn read_guard(&self) -> Self::ReadGuard<'_> {
 		let ref_guard = PoisonRef::new(&self.poisoned, self.inner.read_guard());
 
@@ -79,6 +96,14 @@ unsafe impl<L: Sharable> Sharable for Poisonable<L> {
 			Ok(ref_guard)
 		}
 	}
+
+	unsafe fn data_ref(&self) -> Self::DataRef<'_> {
+		if self.is_poisoned() {
+			Err(PoisonError::new(self.inner.data_ref()))
+		} else {
+			Ok(self.inner.data_ref())
+		}
+	}
 }
 
 unsafe impl<L: OwnedLockable> OwnedLockable for Poisonable<L> {}
@@ -266,14 +291,10 @@ impl<L> Poisonable<L> {
 
 impl<L: Lockable> Poisonable<L> {
 	/// Creates a guard for the poisonable, without locking it
-	unsafe fn guard<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> PoisonResult<PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>> {
+	unsafe fn guard(&self, key: ThreadKey) -> PoisonResult<PoisonGuard<'_, L::Guard<'_>>> {
 		let guard = PoisonGuard {
 			guard: PoisonRef::new(&self.poisoned, self.inner.guard()),
 			key,
-			_phantom: PhantomData,
 		};
 
 		if self.is_poisoned() {
@@ -285,6 +306,50 @@ impl<L: Lockable> Poisonable<L> {
 }
 
 impl<L: Lockable + RawLock> Poisonable<L> {
+	pub fn scoped_lock<'a, R>(
+		&'a self,
+		key: impl Keyable,
+		f: impl Fn(<Self as Lockable>::DataMut<'a>) -> R,
+	) -> R {
+		unsafe {
+			// safety: we have the thread key
+			self.raw_lock();
+
+			// safety: the data was just locked
+			let r = f(self.data_mut());
+
+			// safety: the collection is still locked
+			self.raw_unlock();
+
+			drop(key); // ensure the key stays alive long enough
+
+			r
+		}
+	}
+
+	pub fn scoped_try_lock<'a, Key: Keyable, R>(
+		&'a self,
+		key: Key,
+		f: impl Fn(<Self as Lockable>::DataMut<'a>) -> R,
+	) -> Result<R, Key> {
+		unsafe {
+			// safety: we have the thread key
+			if !self.raw_try_lock() {
+				return Err(key);
+			}
+
+			// safety: we just locked the collection
+			let r = f(self.data_mut());
+
+			// safety: the collection is still locked
+			self.raw_unlock();
+
+			drop(key); // ensures the key stays valid long enough
+
+			Ok(r)
+		}
+	}
+
 	/// Acquires the lock, blocking the current thread until it is ok to do so.
 	///
 	/// This function will block the current thread until it is available to
@@ -316,10 +381,7 @@ impl<L: Lockable + RawLock> Poisonable<L> {
 	/// let key = ThreadKey::get().unwrap();
 	/// assert_eq!(*mutex.lock(key).unwrap(), 10);
 	/// ```
-	pub fn lock<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> PoisonResult<PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>> {
+	pub fn lock(&self, key: ThreadKey) -> PoisonResult<PoisonGuard<'_, L::Guard<'_>>> {
 		unsafe {
 			self.inner.raw_lock();
 			self.guard(key)
@@ -370,10 +432,7 @@ impl<L: Lockable + RawLock> Poisonable<L> {
 	///
 	/// [`Poisoned`]: `TryLockPoisonableError::Poisoned`
 	/// [`WouldBlock`]: `TryLockPoisonableError::WouldBlock`
-	pub fn try_lock<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> TryLockPoisonableResult<'flag, 'key, L::Guard<'flag>, Key> {
+	pub fn try_lock(&self, key: ThreadKey) -> TryLockPoisonableResult<'_, L::Guard<'_>> {
 		unsafe {
 			if self.inner.raw_try_lock() {
 				Ok(self.guard(key)?)
@@ -398,23 +457,17 @@ impl<L: Lockable + RawLock> Poisonable<L> {
 	///
 	/// let key = Poisonable::<Mutex<_>>::unlock(guard);
 	/// ```
-	pub fn unlock<'flag, 'key, Key: Keyable + 'key>(
-		guard: PoisonGuard<'flag, 'key, L::Guard<'flag>, Key>,
-	) -> Key {
+	pub fn unlock<'flag>(guard: PoisonGuard<'flag, L::Guard<'flag>>) -> ThreadKey {
 		drop(guard.guard);
 		guard.key
 	}
 }
 
 impl<L: Sharable + RawLock> Poisonable<L> {
-	unsafe fn read_guard<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> PoisonResult<PoisonGuard<'flag, 'key, L::ReadGuard<'flag>, Key>> {
+	unsafe fn read_guard(&self, key: ThreadKey) -> PoisonResult<PoisonGuard<'_, L::ReadGuard<'_>>> {
 		let guard = PoisonGuard {
 			guard: PoisonRef::new(&self.poisoned, self.inner.read_guard()),
 			key,
-			_phantom: PhantomData,
 		};
 
 		if self.is_poisoned() {
@@ -424,6 +477,50 @@ impl<L: Sharable + RawLock> Poisonable<L> {
 		Ok(guard)
 	}
 
+	pub fn scoped_read<'a, R>(
+		&'a self,
+		key: impl Keyable,
+		f: impl Fn(<Self as Sharable>::DataRef<'a>) -> R,
+	) -> R {
+		unsafe {
+			// safety: we have the thread key
+			self.raw_read();
+
+			// safety: the data was just locked
+			let r = f(self.data_ref());
+
+			// safety: the collection is still locked
+			self.raw_unlock_read();
+
+			drop(key); // ensure the key stays alive long enough
+
+			r
+		}
+	}
+
+	pub fn scoped_try_read<'a, Key: Keyable, R>(
+		&'a self,
+		key: Key,
+		f: impl Fn(<Self as Sharable>::DataRef<'a>) -> R,
+	) -> Result<R, Key> {
+		unsafe {
+			// safety: we have the thread key
+			if !self.raw_try_read() {
+				return Err(key);
+			}
+
+			// safety: we just locked the collection
+			let r = f(self.data_ref());
+
+			// safety: the collection is still locked
+			self.raw_unlock_read();
+
+			drop(key); // ensures the key stays valid long enough
+
+			Ok(r)
+		}
+	}
+
 	/// Locks with shared read access, blocking the current thread until it can
 	/// be acquired.
 	///
@@ -457,10 +554,7 @@ impl<L: Sharable + RawLock> Poisonable<L> {
 	///     assert!(c_lock.read(key).is_ok());
 	/// }).join().expect("thread::spawn failed");
 	/// ```
-	pub fn read<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> PoisonResult<PoisonGuard<'flag, 'key, L::ReadGuard<'flag>, Key>> {
+	pub fn read(&self, key: ThreadKey) -> PoisonResult<PoisonGuard<'_, L::ReadGuard<'_>>> {
 		unsafe {
 			self.inner.raw_read();
 			self.read_guard(key)
@@ -504,10 +598,7 @@ impl<L: Sharable + RawLock> Poisonable<L> {
 	///
 	/// [`Poisoned`]: `TryLockPoisonableError::Poisoned`
 	/// [`WouldBlock`]: `TryLockPoisonableError::WouldBlock`
-	pub fn try_read<'flag, 'key, Key: Keyable + 'key>(
-		&'flag self,
-		key: Key,
-	) -> TryLockPoisonableResult<'flag, 'key, L::ReadGuard<'flag>, Key> {
+	pub fn try_read(&self, key: ThreadKey) -> TryLockPoisonableResult<'_, L::ReadGuard<'_>> {
 		unsafe {
 			if self.inner.raw_try_read() {
 				Ok(self.read_guard(key)?)
@@ -530,9 +621,7 @@ impl<L: Sharable + RawLock> Poisonable<L> {
 	/// let mut guard = lock.read(key).unwrap();
 	/// let key = Poisonable::<RwLock<_>>::unlock_read(guard);
 	/// ```
-	pub fn unlock_read<'flag, 'key, Key: Keyable + 'key>(
-		guard: PoisonGuard<'flag, 'key, L::ReadGuard<'flag>, Key>,
-	) -> Key {
+	pub fn unlock_read<'flag>(guard: PoisonGuard<'flag, L::ReadGuard<'flag>>) -> ThreadKey {
 		drop(guard.guard);
 		guard.key
 	}