1 files changed, 129 insertions, 19 deletions
diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index fbc6842..c700098 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -1,6 +1,13 @@
-use std::mem::transmute;
+use std::mem::{forget, transmute};
 use std::ops::{Deref, DerefMut};
-use std::sync::{PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard, TryLockError};
+
+use parking_lot::{
+    MappedRwLockReadGuard,
+    MappedRwLockWriteGuard,
+    RwLock,
+    RwLockReadGuard,
+    RwLockWriteGuard,
+};
 
 use crate::type_name::TypeName;
 
@@ -27,12 +34,8 @@ where
     /// Returns `Err` if unavailable (A mutable handle is hold).
     pub fn read_nonblock(&self) -> Result<ReadGuard<Value>, Error>
     {
-        let guard = self.inner.try_read().or_else(|err| match err {
-            TryLockError::WouldBlock => Err(Error::Unavailable),
-            TryLockError::Poisoned(poison_err) => Ok(poison_err.into_inner()),
-        })?;
+        let guard = self.inner.try_read().ok_or(Error::ReadUnavailable)?;
 
-        #[cfg(feature = "debug")]
         tracing::trace!("Acquired lock to value of type {}", guard.type_name());
 
         Ok(ReadGuard { inner: guard })
@@ -44,12 +47,8 @@ where
     /// Returns `Err` if unavailable (A mutable or immutable handle is hold).
     pub fn write_nonblock(&self) -> Result<WriteGuard<Value>, Error>
     {
-        let guard = self.inner.try_write().or_else(|err| match err {
-            TryLockError::WouldBlock => Err(Error::Unavailable),
-            TryLockError::Poisoned(poison_err) => Ok(poison_err.into_inner()),
-        })?;
+        let guard = self.inner.try_write().ok_or(Error::WriteUnavailable)?;
 
-        #[cfg(feature = "debug")]
         tracing::trace!(
             "Acquired mutable lock to value of type {}",
             guard.type_name()
@@ -60,17 +59,18 @@ where
 
     pub fn into_inner(self) -> Value
     {
-        self.inner
-            .into_inner()
-            .unwrap_or_else(PoisonError::into_inner)
+        self.inner.into_inner()
     }
 }
 
 #[derive(Debug, thiserror::Error)]
 pub enum Error
 {
-    #[error("Lock is unavailable")]
-    Unavailable,
+    #[error("Lock is unavailable for reading")]
+    ReadUnavailable,
+
+    #[error("Lock is unavailable for writing")]
+    WriteUnavailable,
 }
 
 #[derive(Debug)]
@@ -85,6 +85,23 @@ impl<'guard, Value> ReadGuard<'guard, Value>
 where
     Value: TypeName,
 {
+    pub fn map<NewValue>(
+        self,
+        func: impl FnOnce(&Value) -> &NewValue,
+    ) -> MappedReadGuard<'guard, NewValue>
+    where
+        NewValue: TypeName,
+    {
+        // The 'inner' field cannot be moved out of ReadGuard in a normal way since
+        // ReadGuard implements Drop
+        let inner = unsafe { std::ptr::read(&self.inner) };
+        forget(self);
+
+        MappedReadGuard {
+            inner: RwLockReadGuard::map(inner, func),
+        }
+    }
+
     /// Converts the `ReadGuard` to a `ReadGuard` with a possibly longer lifetime.
     ///
     /// # Safety
@@ -115,12 +132,41 @@ where
 {
     fn drop(&mut self)
     {
-        #[cfg(feature = "debug")]
         tracing::trace!("Dropped lock to value of type {}", self.type_name());
     }
 }
 
 #[derive(Debug)]
+pub struct MappedReadGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    inner: MappedRwLockReadGuard<'guard, Value>,
+}
+
+impl<'guard, Value> Deref for MappedReadGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    type Target = Value;
+
+    fn deref(&self) -> &Self::Target
+    {
+        &self.inner
+    }
+}
+
+impl<'guard, Value> Drop for MappedReadGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    fn drop(&mut self)
+    {
+        tracing::trace!("Dropped mapped lock to value of type {}", self.type_name());
+    }
+}
+
+#[derive(Debug)]
 pub struct WriteGuard<'guard, Value>
 where
     Value: TypeName,
@@ -128,6 +174,28 @@ where
     inner: RwLockWriteGuard<'guard, Value>,
 }
 
+impl<'guard, Value> WriteGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    pub fn map<NewValue>(
+        self,
+        func: impl FnOnce(&mut Value) -> &mut NewValue,
+    ) -> MappedWriteGuard<'guard, NewValue>
+    where
+        NewValue: TypeName,
+    {
+        // The 'inner' field cannot be moved out of ReadGuard in a normal way since
+        // ReadGuard implements Drop
+        let inner = unsafe { std::ptr::read(&self.inner) };
+        forget(self);
+
+        MappedWriteGuard {
+            inner: RwLockWriteGuard::map(inner, func),
+        }
+    }
+}
+
 impl<'guard, Value> Deref for WriteGuard<'guard, Value>
 where
     Value: TypeName,
@@ -156,7 +224,49 @@ where
 {
     fn drop(&mut self)
     {
-        #[cfg(feature = "debug")]
         tracing::trace!("Dropped mutable lock to value of type {}", self.type_name());
     }
 }
+
+#[derive(Debug)]
+pub struct MappedWriteGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    inner: MappedRwLockWriteGuard<'guard, Value>,
+}
+
+impl<'guard, Value> Deref for MappedWriteGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    type Target = Value;
+
+    fn deref(&self) -> &Self::Target
+    {
+        &self.inner
+    }
+}
+
+impl<'guard, Value> DerefMut for MappedWriteGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target
+    {
+        &mut self.inner
+    }
+}
+
+impl<'guard, Value> Drop for MappedWriteGuard<'guard, Value>
+where
+    Value: TypeName,
+{
+    fn drop(&mut self)
+    {
+        tracing::trace!(
+            "Dropped mapped mutable lock to value of type {}",
+            self.type_name()
+        );
+    }
+}