feat(ecs): add funcs for getting target comp of wildcard pairs

8 files changed, 151 insertions, 88 deletions

diff --git a/ecs/examples/relationship.rs b/ecs/examples/relationship.rs
index f99d090..6ad08e7 100644
--- a/ecs/examples/relationship.rs
+++ b/ecs/examples/relationship.rs
@@ -22,10 +22,10 @@ struct Holding;
 
 fn print_player_stats(player_query: Query<(&Player, &Health, Pair<Holding, Wildcard>)>)
 {
-    for (_, health, sword_relationship) in &player_query {
+    for (_, health, target_sword) in &player_query {
         println!("Player health: {}", health.health);
 
-        if let Some(sword_ent) = sword_relationship.get_target_entity() {
+        if let Some(sword_ent) = target_sword.get_entity() {
             let sword = sword_ent
                 .get::<Sword>()
                 .expect("Sword entity is missing sword component");
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)]
diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index ef91d20..6c9ec32 100644
--- a/ecs/src/entity.rs
+++ b/ecs/src/entity.rs
@@ -48,12 +48,14 @@ impl<'a> Handle<'a>
         let component = self.get_matching_components(ComponentT::id()).next()?;
 
         Some(
-            ComponentHandle::from_entity_component_ref(component).unwrap_or_else(|err| {
-                panic!(
-                    "Taking component {} lock failed: {err}",
-                    type_name::<ComponentT>()
-                );
-            }),
+            ComponentHandle::from_entity_component_ref(&component).unwrap_or_else(
+                |err| {
+                    panic!(
+                        "Creating handle to component {} failed: {err}",
+                        type_name::<ComponentT>()
+                    );
+                },
+            ),
         )
     }
 
@@ -74,10 +76,10 @@ impl<'a> Handle<'a>
         let component = self.get_matching_components(ComponentT::id()).next()?;
 
         Some(
-            ComponentHandleMut::from_entity_component_ref(component).unwrap_or_else(
+            ComponentHandleMut::from_entity_component_ref(&component).unwrap_or_else(
                 |err| {
                     panic!(
-                        "Taking component {} lock failed: {err}",
+                        "Creating handle to component {} failed: {err}",
                         type_name::<ComponentT>()
                     );
                 },
diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index 0b36922..d21b697 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -81,21 +81,26 @@ pub struct ReadGuard<'guard, Value>
 
 impl<'guard, Value> ReadGuard<'guard, Value>
 {
-    pub fn map<NewValue>(
-        self,
-        func: impl FnOnce(&Value) -> &NewValue,
-    ) -> MappedReadGuard<'guard, NewValue>
+    pub fn try_map<NewValue>(
+        this: Self,
+        func: impl FnOnce(&Value) -> Option<&NewValue>,
+    ) -> Result<MappedReadGuard<'guard, NewValue>, Self>
     {
-        let value_type_name = self.value_type_name;
+        let value_type_name = this.value_type_name;
 
         // 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),
-            value_type_name,
+        let inner = unsafe { std::ptr::read(&this.inner) };
+        forget(this);
+
+        match RwLockReadGuard::try_map(inner, func) {
+            Ok(mapped_guard) => {
+                Ok(MappedReadGuard { inner: mapped_guard, value_type_name })
+            }
+            Err(unmapped_guard) => Err(Self {
+                inner: unmapped_guard,
+                value_type_name,
+            }),
         }
     }
 }
@@ -155,21 +160,26 @@ pub struct WriteGuard<'guard, Value>
 
 impl<'guard, Value> WriteGuard<'guard, Value>
 {
-    pub fn map<NewValue>(
-        self,
-        func: impl FnOnce(&mut Value) -> &mut NewValue,
-    ) -> MappedWriteGuard<'guard, NewValue>
+    pub fn try_map<NewValue>(
+        this: Self,
+        func: impl FnOnce(&mut Value) -> Option<&mut NewValue>,
+    ) -> Result<MappedWriteGuard<'guard, NewValue>, Self>
     {
-        let value_type_name = self.value_type_name;
+        let value_type_name = this.value_type_name;
 
         // 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),
-            value_type_name,
+        let inner = unsafe { std::ptr::read(&this.inner) };
+        forget(this);
+
+        match RwLockWriteGuard::try_map(inner, func) {
+            Ok(mapped_guard) => {
+                Ok(MappedWriteGuard { inner: mapped_guard, value_type_name })
+            }
+            Err(unmapped_guard) => Err(Self {
+                inner: unmapped_guard,
+                value_type_name,
+            }),
         }
     }
 }
diff --git a/ecs/src/pair.rs b/ecs/src/pair.rs
index 77b6da7..3dc9f36 100644
--- a/ecs/src/pair.rs
+++ b/ecs/src/pair.rs
@@ -3,6 +3,7 @@ use std::convert::Infallible;
 
 use crate::component::{
     Handle as ComponentHandle,
+    HandleError as ComponentHandleError,
     HandleMut as ComponentHandleMut,
     IntoParts as IntoComponentParts,
     Parts as ComponentParts,
@@ -17,7 +18,7 @@ use crate::query::{
     TermsBuilderInterface,
 };
 use crate::uid::{PairParams as UidPairParams, Uid, With as WithUid};
-use crate::{Component, World};
+use crate::{Component, EntityComponentRef, World};
 
 #[derive(Debug)]
 pub struct Pair<Relation, Target>
@@ -107,7 +108,7 @@ where
             .next()
             .expect("Not possible");
 
-        Self::Field::from_entity_component_ref(target_component).unwrap_or_else(|err| {
+        Self::Field::from_entity_component_ref(&target_component).unwrap_or_else(|err| {
             panic!(
                 "Creating handle to target component {} failed: {err}",
                 type_name::<Target>()
@@ -140,7 +141,7 @@ where
             .next()
             .expect("Not possible");
 
-        Self::Field::from_entity_component_ref(target_component).unwrap_or_else(|err| {
+        Self::Field::from_entity_component_ref(&target_component).unwrap_or_else(|err| {
             panic!(
                 "Creating handle to target component {} failed: {err}",
                 type_name::<Target>()
@@ -153,7 +154,7 @@ impl<Relation> QueryTermWithField for Pair<Relation, Wildcard>
 where
     Relation: Component,
 {
-    type Field<'a> = EntityTargetHandle<'a>;
+    type Field<'a> = WildcardTargetHandle<'a>;
 
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         terms_builder: &mut QueryTermsBuilder<MAX_TERM_CNT>,
@@ -172,9 +173,9 @@ where
             .next()
             .expect("Not possible");
 
-        EntityTargetHandle {
+        WildcardTargetHandle {
             world,
-            pair_uid: first_matching_comp.id(),
+            component_ref: first_matching_comp,
         }
     }
 }
@@ -224,7 +225,7 @@ impl<Relation> QueryTermWithField for &'static [Pair<Relation, Wildcard>]
 where
     Relation: Component,
 {
-    type Field<'a> = HandleIter<'a>;
+    type Field<'a> = WildcardTargetIter<'a>;
 
     fn apply_to_terms_builder<const MAX_TERM_CNT: usize>(
         terms_builder: &mut QueryTermsBuilder<MAX_TERM_CNT>,
@@ -238,7 +239,7 @@ where
         world: &'world World,
     ) -> Self::Field<'world>
     {
-        HandleIter {
+        WildcardTargetIter {
             inner: entity_handle
                 .get_matching_components(Pair::<Relation, Wildcard>::uid()),
             world,
@@ -246,50 +247,107 @@ where
     }
 }
 
-pub struct EntityTargetHandle<'world>
+pub struct WildcardTargetHandle<'world>
 {
     world: &'world World,
-    pair_uid: Uid,
+    component_ref: EntityComponentRef<'world>,
 }
 
-impl EntityTargetHandle<'_>
+impl WildcardTargetHandle<'_>
 {
+    /// Attempts to retrieve the target as a entity, returning `None` if not found.
     #[must_use]
-    pub fn get_target_entity(&self) -> Option<EntityHandle<'_>>
+    pub fn get_entity(&self) -> Option<EntityHandle<'_>>
     {
         let archetype = self
             .world
             .data
             .component_storage
-            .get_entity_archetype(self.pair_uid.target_entity())?;
+            .get_entity_archetype(self.component_ref.id().target_entity())?;
 
         let Some(archetype_entity) =
-            archetype.get_entity_by_id(self.pair_uid.target_entity())
+            archetype.get_entity_by_id(self.component_ref.id().target_entity())
         else {
             unreachable!();
         };
 
         Some(EntityHandle::new(archetype, archetype_entity))
     }
+
+    /// Attempts to retrieve the target as a component, returning `None` if the component
+    /// type is incorrect.
+    ///
+    /// # Panics
+    /// Will panic if:
+    /// - The component is mutably borrowed elsewhere
+    pub fn get_component<ComponentData>(
+        &self,
+    ) -> Option<ComponentHandle<'_, ComponentData>>
+    where
+        ComponentData: 'static,
+    {
+        ComponentHandle::<ComponentData>::from_entity_component_ref(&self.component_ref)
+            .map_or_else(
+                |err| match err {
+                    ComponentHandleError::IncorrectType => None,
+                    err @ ComponentHandleError::AcquireLockFailed(_) => {
+                        panic!(
+                            "Creating handle to component {} failed: {err}",
+                            type_name::<ComponentData>()
+                        );
+                    }
+                },
+                |handle| Some(handle),
+            )
+    }
+
+    /// Attempts to retrieve the target as a component, returning `None` if the component
+    /// type is incorrect.
+    ///
+    /// # Panics
+    /// Will panic if:
+    /// - The component is borrowed elsewhere
+    pub fn get_component_mut<ComponentData>(
+        &self,
+    ) -> Option<ComponentHandleMut<'_, ComponentData>>
+    where
+        ComponentData: 'static,
+    {
+        ComponentHandleMut::<ComponentData>::from_entity_component_ref(
+            &self.component_ref,
+        )
+        .map_or_else(
+            |err| match err {
+                ComponentHandleError::IncorrectType => None,
+                err @ ComponentHandleError::AcquireLockFailed(_) => {
+                    panic!(
+                        "Creating handle to component {} failed: {err}",
+                        type_name::<ComponentData>()
+                    );
+                }
+            },
+            |handle| Some(handle),
+        )
+    }
 }
 
-pub struct HandleIter<'a>
+pub struct WildcardTargetIter<'a>
 {
     inner: EntityMatchingComponentIter<'a>,
     world: &'a World,
 }
 
-impl<'a> Iterator for HandleIter<'a>
+impl<'a> Iterator for WildcardTargetIter<'a>
 {
-    type Item = EntityTargetHandle<'a>;
+    type Item = WildcardTargetHandle<'a>;
 
     fn next(&mut self) -> Option<Self::Item>
     {
         let matching_comp = self.inner.next()?;
 
-        Some(EntityTargetHandle {
+        Some(WildcardTargetHandle {
             world: self.world,
-            pair_uid: matching_comp.id(),
+            component_ref: matching_comp,
         })
     }
 }
diff --git a/ecs/src/query.rs b/ecs/src/query.rs
index ccb7add..dfe9dad 100644
--- a/ecs/src/query.rs
+++ b/ecs/src/query.rs
@@ -360,7 +360,7 @@ impl<ComponentT: Component> TermWithField for &ComponentT
             );
         };
 
-        Self::Field::from_entity_component_ref(component).unwrap_or_else(|err| {
+        Self::Field::from_entity_component_ref(&component).unwrap_or_else(|err| {
             panic!(
                 "Creating handle to component {} failed: {err}",
                 type_name::<ComponentT>()
@@ -401,7 +401,7 @@ impl<ComponentT: Component> TermWithField for &mut ComponentT
             );
         };
 
-        Self::Field::from_entity_component_ref(component).unwrap_or_else(|err| {
+        Self::Field::from_entity_component_ref(&component).unwrap_or_else(|err| {
             panic!(
                 "Creating handle to component {} failed: {err}",
                 type_name::<ComponentT>()
diff --git a/ecs/src/query/term.rs b/ecs/src/query/term.rs
index 9c772da..c8a96b6 100644
--- a/ecs/src/query/term.rs
+++ b/ecs/src/query/term.rs
@@ -69,7 +69,7 @@ impl<ComponentT: Component> TermWithField for Option<&ComponentT>
     {
         Some(
             ComponentHandle::<'world, ComponentT>::from_entity_component_ref(
-                entity_handle
+                &entity_handle
                     .get_matching_components(ComponentT::id())
                     .next()?,
             )
@@ -100,7 +100,7 @@ impl<ComponentT: Component> TermWithField for Option<&mut ComponentT>
     {
         Some(
             ComponentHandleMut::<'world, ComponentT>::from_entity_component_ref(
-                entity_handle
+                &entity_handle
                     .get_matching_components(ComponentT::id())
                     .next()?,
             )
diff --git a/ecs/src/system/stateful.rs b/ecs/src/system/stateful.rs
index 54f6979..2ec448c 100644
--- a/ecs/src/system/stateful.rs
+++ b/ecs/src/system/stateful.rs
@@ -116,7 +116,7 @@ macro_rules! impl_system {
                         .write_nonblock()
                         .expect("Failed to aquire read-write local component lock");
 
-                    Some(ComponentHandleMut::new(local_component))
+                    Some(ComponentHandleMut::new(local_component).unwrap())
                 }
 
                 fn set_local_component<LocalComponent: Component>(