1 files changed, 77 insertions, 19 deletions

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,
         })
     }
 }