1 files changed, 25 insertions, 32 deletions
diff --git a/ecs/src/component.rs b/ecs/src/component.rs
index 6fb1230..3c2112e 100644
--- a/ecs/src/component.rs
+++ b/ecs/src/component.rs
@@ -81,31 +81,25 @@ impl<'comp, ComponentData: 'static> Handle<'comp, ComponentData>
     /// # Errors
     /// Will return `Err` if acquiring the component's lock fails.
     pub fn from_entity_component_ref(
-        entity_component_ref: EntityComponentRef<'comp>,
+        entity_component_ref: &EntityComponentRef<'comp>,
     ) -> Result<Self, HandleError>
     {
-        Ok(Self::new(
+        Self::new(
             entity_component_ref
                 .component()
                 .read_nonblock()
                 .map_err(AcquireLockError)?,
-        ))
+        )
     }
 
-    pub(crate) fn new(inner: ReadGuard<'comp, Box<dyn Any>>) -> Self
+    fn new(inner: ReadGuard<'comp, Box<dyn Any>>) -> Result<Self, HandleError>
     {
-        Self {
-            inner: inner.map(|component| {
-                component
-                    .downcast_ref::<ComponentData>()
-                    .unwrap_or_else(|| {
-                        panic!(
-                            "Failed to downcast component to type {}",
-                            type_name::<ComponentData>()
-                        );
-                    })
-            }),
-        }
+        Ok(Self {
+            inner: ReadGuard::try_map(inner, |component| {
+                component.downcast_ref::<ComponentData>()
+            })
+            .map_err(|_| HandleError::IncorrectType)?,
+        })
     }
 }
 
@@ -132,31 +126,27 @@ impl<'comp, ComponentData: 'static> HandleMut<'comp, ComponentData>
     /// # Errors
     /// Will return `Err` if acquiring the component's lock fails.
     pub fn from_entity_component_ref(
-        entity_component_ref: EntityComponentRef<'comp>,
+        entity_component_ref: &EntityComponentRef<'comp>,
     ) -> Result<Self, HandleError>
     {
-        Ok(Self::new(
+        Self::new(
             entity_component_ref
                 .component()
                 .write_nonblock()
                 .map_err(AcquireLockError)?,
-        ))
+        )
     }
 
-    pub(crate) fn new(inner: WriteGuard<'comp, Box<dyn Any>>) -> Self
+    // TODO: Make this function private
+    pub(crate) fn new(inner: WriteGuard<'comp, Box<dyn Any>>)
+        -> Result<Self, HandleError>
     {
-        Self {
-            inner: inner.map(|component| {
-                component
-                    .downcast_mut::<ComponentData>()
-                    .unwrap_or_else(|| {
-                        panic!(
-                            "Failed to downcast component to type {}",
-                            type_name::<ComponentData>()
-                        );
-                    })
-            }),
-        }
+        Ok(Self {
+            inner: WriteGuard::try_map(inner, |component| {
+                component.downcast_mut::<ComponentData>()
+            })
+            .map_err(|_| HandleError::IncorrectType)?,
+        })
     }
 }
 
@@ -183,6 +173,9 @@ pub enum HandleError
 {
     #[error(transparent)]
     AcquireLockFailed(#[from] AcquireLockError),
+
+    #[error("Incorrect component type")]
+    IncorrectType,
 }
 
 #[derive(Debug, thiserror::Error)]