3 files changed, 34 insertions, 43 deletions

diff --git a/engine-ecs/examples/with_sole.rs b/engine-ecs/examples/with_sole.rs
index b76d9dc..16a5aba 100644
--- a/engine-ecs/examples/with_sole.rs
+++ b/engine-ecs/examples/with_sole.rs
@@ -20,15 +20,17 @@ fn count_ammo(query: Query<(&Ammo,)>, mut ammo_counter: Single<AmmoCounter>)
     for (ammo,) in &query {
         println!("Found {} ammo", ammo.ammo_left);
 
-        ammo_counter.counter += ammo.ammo_left;
+        ammo_counter.get_mut().unwrap().counter += ammo.ammo_left;
     }
 }
 
 fn print_total_ammo_count(ammo_counter: Single<AmmoCounter>)
 {
-    println!("Total ammo count: {}", ammo_counter.counter);
+    let ammo_count = ammo_counter.get().unwrap().counter;
 
-    assert_eq!(ammo_counter.counter, 19);
+    println!("Total ammo count: {}", ammo_count);
+
+    assert_eq!(ammo_count, 19);
 }
 
 declare_entity!(
diff --git a/engine-ecs/src/lib.rs b/engine-ecs/src/lib.rs
index 3292bec..7bad923 100644
--- a/engine-ecs/src/lib.rs
+++ b/engine-ecs/src/lib.rs
@@ -312,13 +312,6 @@ impl World
         Some(EntityHandle::new(archetype, entity, self))
     }
 
-    pub fn get_sole<SoleT: Sole>(&self) -> Option<Single<'_, SoleT>>
-    {
-        let ent = self.get_entity(SoleT::id())?;
-
-        Some(Single::new(ent.get_with_id_mut(SoleT::id())?))
-    }
-
     pub fn event_submitter(&self) -> EventSubmitter<'_>
     {
         EventSubmitter::new(&self.data.new_events)
@@ -363,8 +356,10 @@ impl World
             return StepResult::Stop;
         }
 
-        let Some(mut stats) = self.get_sole::<Stats>() else {
-            unreachable!(); // Reason: is added in World::new
+        let mut stats = Single::<Stats>::new(self);
+
+        let Ok(stats) = stats.get_mut() else {
+            unreachable!(); // Stats is added in World::new
         };
 
         stats.current_tick += 1;
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,
 }