fix(ecs): make ComponentRef & ComponentRefMut deref not able to panic

2 files changed, 170 insertions, 28 deletions

diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index 74b0415..c700098 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -1,7 +1,13 @@
-use std::mem::transmute;
+use std::mem::{forget, transmute};
 use std::ops::{Deref, DerefMut};
 
-use parking_lot::{RwLock, RwLockReadGuard, RwLockWriteGuard};
+use parking_lot::{
+    MappedRwLockReadGuard,
+    MappedRwLockWriteGuard,
+    RwLock,
+    RwLockReadGuard,
+    RwLockWriteGuard,
+};
 
 use crate::type_name::TypeName;
 
@@ -79,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
@@ -114,6 +137,36 @@ where
 }
 
 #[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,
@@ -121,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,
@@ -152,3 +227,46 @@ where
         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()
+        );
+    }
+}
diff --git a/ecs/src/system.rs b/ecs/src/system.rs
index 48c2723..b410d8f 100644
--- a/ecs/src/system.rs
+++ b/ecs/src/system.rs
@@ -13,7 +13,14 @@ use crate::component::{
     FromOptional as FromOptionalComponent,
     FromOptionalMut as FromOptionalMutComponent,
 };
-use crate::lock::{Error as LockError, Lock, ReadGuard, WriteGuard};
+use crate::lock::{
+    Error as LockError,
+    Lock,
+    MappedReadGuard,
+    MappedWriteGuard,
+    ReadGuard,
+    WriteGuard,
+};
 use crate::tuple::{ReduceElement as TupleReduceElement, Tuple};
 use crate::World;
 
@@ -190,7 +197,7 @@ impl<Accumulator> TupleReduceElement<Accumulator, ParamWithInputFilter> for ()
 #[derive(Debug)]
 pub struct ComponentRefMut<'a, ComponentT: Component>
 {
-    inner: WriteGuard<'a, Box<dyn Component>>,
+    inner: MappedWriteGuard<'a, ComponentT>,
     _ph: PhantomData<ComponentT>,
 }
 
@@ -198,7 +205,19 @@ impl<'a, ComponentT: Component> ComponentRefMut<'a, ComponentT>
 {
     pub(crate) fn new(inner: WriteGuard<'a, Box<dyn Component>>) -> Self
     {
-        Self { inner, _ph: PhantomData }
+        Self {
+            inner: inner.map(|component| {
+                let component_type_name = component.type_name();
+
+                component.downcast_mut::<ComponentT>().unwrap_or_else(|| {
+                    panic!(
+                        "Cannot downcast component {component_type_name} to type {}",
+                        type_name::<ComponentT>()
+                    );
+                })
+            }),
+            _ph: PhantomData,
+        }
     }
 }
 
@@ -210,15 +229,12 @@ impl<'component, ComponentT: Component> FromOptionalMutComponent<'component>
         _world: &'component World,
     ) -> Self
     {
-        Self {
-            inner: inner.unwrap_or_else(|| {
-                panic!(
-                    "Component {} was not found in entity",
-                    type_name::<ComponentT>()
-                );
-            }),
-            _ph: PhantomData,
-        }
+        Self::new(inner.unwrap_or_else(|| {
+            panic!(
+                "Component {} was not found in entity",
+                type_name::<ComponentT>()
+            );
+        }))
     }
 }
 
@@ -275,7 +291,7 @@ impl<'a, ComponentT: Component> Deref for ComponentRefMut<'a, ComponentT>
 
     fn deref(&self) -> &Self::Target
     {
-        self.inner.downcast_ref().unwrap()
+        &self.inner
     }
 }
 
@@ -283,14 +299,14 @@ impl<'a, ComponentT: Component> DerefMut for ComponentRefMut<'a, ComponentT>
 {
     fn deref_mut(&mut self) -> &mut Self::Target
     {
-        self.inner.downcast_mut().unwrap()
+        &mut self.inner
     }
 }
 
 #[derive(Debug)]
 pub struct ComponentRef<'a, ComponentT: Component>
 {
-    inner: ReadGuard<'a, Box<dyn Component>>,
+    inner: MappedReadGuard<'a, ComponentT>,
     _ph: PhantomData<ComponentT>,
 }
 
@@ -298,7 +314,18 @@ impl<'a, ComponentT: Component> ComponentRef<'a, ComponentT>
 {
     pub(crate) fn new(inner: ReadGuard<'a, Box<dyn Component>>) -> Self
     {
-        Self { inner, _ph: PhantomData }
+        Self {
+            inner: inner.map(|component| {
+                component.downcast_ref::<ComponentT>().unwrap_or_else(|| {
+                    panic!(
+                        "Cannot downcast component {} to type {}",
+                        component.type_name(),
+                        type_name::<ComponentT>()
+                    );
+                })
+            }),
+            _ph: PhantomData,
+        }
     }
 }
 
@@ -310,15 +337,12 @@ impl<'component, ComponentT: Component> FromOptionalComponent<'component>
         _world: &'component World,
     ) -> Self
     {
-        Self {
-            inner: inner.unwrap_or_else(|| {
-                panic!(
-                    "Component {} was not found in entity",
-                    type_name::<ComponentT>()
-                );
-            }),
-            _ph: PhantomData,
-        }
+        Self::new(inner.unwrap_or_else(|| {
+            panic!(
+                "Component {} was not found in entity",
+                type_name::<ComponentT>()
+            );
+        }))
     }
 }
 
@@ -375,7 +399,7 @@ impl<'a, ComponentT: Component> Deref for ComponentRef<'a, ComponentT>
 
     fn deref(&self) -> &Self::Target
     {
-        self.inner.downcast_ref().unwrap()
+        &self.inner
     }
 }