1 files changed, 25 insertions, 31 deletions

diff --git a/engine-ecs/src/sole.rs b/engine-ecs/src/sole.rs
index 21dafc9..1e287d3 100644
--- a/engine-ecs/src/sole.rs
+++ b/engine-ecs/src/sole.rs
@@ -1,6 +1,5 @@
 use std::any::{type_name, Any};
 use std::fmt::Debug;
-use std::ops::{Deref, DerefMut};
 
 use crate::component::{
     HandleMut as ComponentHandleMut,
@@ -55,59 +54,54 @@ impl Debug for dyn Sole
 #[derive(Debug)]
 pub struct Single<'world, SoleT: Sole>
 {
-    sole: ComponentHandleMut<'world, SoleT>,
+    sole: Option<ComponentHandleMut<'world, SoleT>>,
 }
 
 impl<'world, SoleT> Single<'world, SoleT>
 where
     SoleT: Sole,
 {
-    pub(crate) fn new(sole: ComponentHandleMut<'world, SoleT>) -> Self
+    /// Returns a reference to `SoleT` if it exists. If not, returns `Err`.
+    pub fn get(&self) -> Result<&SoleT, DoesNotExistError>
     {
-        Self { sole }
+        self.sole
+            .as_deref()
+            .ok_or_else(|| DoesNotExistError { name: type_name::<SoleT>() })
     }
-}
 
-impl<'world, SoleT> SystemParam<'world> for Single<'world, SoleT>
-where
-    SoleT: Sole,
-{
-    type Input = ();
+    /// Returns a mutable reference to `SoleT` if it exists. If not, returns `Err`.
+    pub fn get_mut(&mut self) -> Result<&mut SoleT, DoesNotExistError>
+    {
+        self.sole
+            .as_deref_mut()
+            .ok_or_else(|| DoesNotExistError { name: type_name::<SoleT>() })
+    }
 
-    fn new(world: &'world World, _system_metadata: &SystemMetadata) -> Self
+    pub(crate) fn new(world: &'world World) -> Self
     {
         let sole = world
             .get_entity(SoleT::id())
-            .and_then(|ent| ent.get_with_id_mut::<SoleT>(SoleT::id()))
-            .unwrap_or_else(|| {
-                panic!(
-                    "Sole component {} was not found in world",
-                    type_name::<SoleT>()
-                )
-            });
-
-        Self::new(sole)
+            .and_then(|ent| ent.get_with_id_mut::<SoleT>(SoleT::id()));
+
+        Self { sole }
     }
 }
 
-impl<SoleT> Deref for Single<'_, SoleT>
+impl<'world, SoleT> SystemParam<'world> for Single<'world, SoleT>
 where
     SoleT: Sole,
 {
-    type Target = SoleT;
+    type Input = ();
 
-    fn deref(&self) -> &Self::Target
+    fn new(world: &'world World, _system_metadata: &SystemMetadata) -> Self
     {
-        &*self.sole
+        Self::new(world)
     }
 }
 
-impl<SoleT> DerefMut for Single<'_, SoleT>
-where
-    SoleT: Sole,
+#[derive(Debug, thiserror::Error)]
+#[error("Sole {name} does not exist")]
+pub struct DoesNotExistError
 {
-    fn deref_mut(&mut self) -> &mut Self::Target
-    {
-        &mut *self.sole
-    }
+    pub name: &'static str,
 }