1 files changed, 103 insertions, 10 deletions

diff --git a/src/collection/utils.rs b/src/collection/utils.rs
index 1d96e5c..59a68da 100644
--- a/src/collection/utils.rs
+++ b/src/collection/utils.rs
@@ -1,7 +1,8 @@
 use std::cell::Cell;
 
 use crate::handle_unwind::handle_unwind;
-use crate::lockable::{Lockable, RawLock};
+use crate::lockable::{Lockable, RawLock, Sharable};
+use crate::Keyable;
 
 #[must_use]
 pub fn get_locks<L: Lockable>(data: &L) -> Vec<&dyn RawLock> {
@@ -32,18 +33,18 @@ pub fn ordered_contains_duplicates(l: &[&dyn RawLock]) -> bool {
 }
 
 /// Lock a set of locks in the given order. It's UB to call this without a `ThreadKey`
-pub unsafe fn ordered_lock(locks: &[&dyn RawLock]) {
+pub unsafe fn ordered_write(locks: &[&dyn RawLock]) {
 	// these will be unlocked in case of a panic
 	let locked = Cell::new(0);
 
 	handle_unwind(
 		|| {
 			for lock in locks {
-				lock.raw_lock();
+				lock.raw_write();
 				locked.set(locked.get() + 1);
 			}
 		},
-		|| attempt_to_recover_locks_from_panic(&locks[0..locked.get()]),
+		|| attempt_to_recover_writes_from_panic(&locks[0..locked.get()]),
 	)
 }
 
@@ -65,19 +66,19 @@ pub unsafe fn ordered_read(locks: &[&dyn RawLock]) {
 /// Locks the locks in the order they are given. This causes deadlock if the
 /// locks contain duplicates, or if this is called by multiple threads with the
 /// locks in different orders.
-pub unsafe fn ordered_try_lock(locks: &[&dyn RawLock]) -> bool {
+pub unsafe fn ordered_try_write(locks: &[&dyn RawLock]) -> bool {
 	let locked = Cell::new(0);
 
 	handle_unwind(
 		|| unsafe {
 			for (i, lock) in locks.iter().enumerate() {
 				// safety: we have the thread key
-				if lock.raw_try_lock() {
+				if lock.raw_try_write() {
 					locked.set(locked.get() + 1);
 				} else {
 					for lock in &locks[0..i] {
 						// safety: this lock was already acquired
-						lock.raw_unlock();
+						lock.raw_unlock_write();
 					}
 					return false;
 				}
@@ -87,7 +88,7 @@ pub unsafe fn ordered_try_lock(locks: &[&dyn RawLock]) -> bool {
 		},
 		||
 		// safety: everything in locked is locked
-		attempt_to_recover_locks_from_panic(&locks[0..locked.get()]),
+		attempt_to_recover_writes_from_panic(&locks[0..locked.get()]),
 	)
 }
 
@@ -120,12 +121,104 @@ pub unsafe fn ordered_try_read(locks: &[&dyn RawLock]) -> bool {
 	)
 }
 
+pub fn scoped_write<'a, L: RawLock + Lockable, R>(
+	collection: &'a L,
+	key: impl Keyable,
+	f: impl FnOnce(L::DataMut<'a>) -> R,
+) -> R {
+	unsafe {
+		// safety: we have the key
+		collection.raw_write();
+
+		// safety: we just locked this
+		let r = f(collection.data_mut());
+
+		// this ensures the key is held long enough
+		drop(key);
+
+		// safety: we've locked already, and aren't using the data again
+		collection.raw_unlock_write();
+
+		r
+	}
+}
+
+pub fn scoped_try_write<'a, L: RawLock + Lockable, Key: Keyable, R>(
+	collection: &'a L,
+	key: Key,
+	f: impl FnOnce(L::DataMut<'a>) -> R,
+) -> Result<R, Key> {
+	unsafe {
+		// safety: we have the key
+		if !collection.raw_try_write() {
+			return Err(key);
+		}
+
+		// safety: we just locked this
+		let r = f(collection.data_mut());
+
+		// this ensures the key is held long enough
+		drop(key);
+
+		// safety: we've locked already, and aren't using the data again
+		collection.raw_unlock_write();
+
+		Ok(r)
+	}
+}
+
+pub fn scoped_read<'a, L: RawLock + Sharable, R>(
+	collection: &'a L,
+	key: impl Keyable,
+	f: impl FnOnce(L::DataRef<'a>) -> R,
+) -> R {
+	unsafe {
+		// safety: we have the key
+		collection.raw_read();
+
+		// safety: we just locked this
+		let r = f(collection.data_ref());
+
+		// this ensures the key is held long enough
+		drop(key);
+
+		// safety: we've locked already, and aren't using the data again
+		collection.raw_unlock_read();
+
+		r
+	}
+}
+
+pub fn scoped_try_read<'a, L: RawLock + Sharable, Key: Keyable, R>(
+	collection: &'a L,
+	key: Key,
+	f: impl FnOnce(L::DataRef<'a>) -> R,
+) -> Result<R, Key> {
+	unsafe {
+		// safety: we have the key
+		if !collection.raw_try_read() {
+			return Err(key);
+		}
+
+		// safety: we just locked this
+		let r = f(collection.data_ref());
+
+		// this ensures the key is held long enough
+		drop(key);
+
+		// safety: we've locked already, and aren't using the data again
+		collection.raw_unlock_read();
+
+		Ok(r)
+	}
+}
+
 /// Unlocks the already locked locks in order to recover from a panic
-pub unsafe fn attempt_to_recover_locks_from_panic(locks: &[&dyn RawLock]) {
+pub unsafe fn attempt_to_recover_writes_from_panic(locks: &[&dyn RawLock]) {
 	handle_unwind(
 		|| {
 			// safety: the caller assumes that these are already locked
-			locks.iter().for_each(|lock| lock.raw_unlock());
+			locks.iter().for_each(|lock| lock.raw_unlock_write());
 		},
 		// if we get another panic in here, we'll just have to poison what remains
 		|| locks.iter().for_each(|l| l.poison()),