feat(ecs): replace Relationship component with pair UID support

1 files changed, 85 insertions, 9 deletions

diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index a43f9ce..3f5a3d3 100644
--- a/ecs/src/entity.rs
+++ b/ecs/src/entity.rs
@@ -1,12 +1,26 @@
+use std::any::type_name;
+
 use linkme::distributed_slice;
 
-use crate::component::storage::archetype::{Archetype, Entity as ArchetypeEntity};
-use crate::uid::Uid;
+use crate::component::storage::archetype::{
+    Archetype,
+    Entity as ArchetypeEntity,
+    MatchingComponentIter as ArchetypeMatchingComponentIter,
+};
+use crate::component::{
+    Component,
+    Handle as ComponentHandle,
+    HandleFromEntityComponentRef,
+    HandleMut as ComponentHandleMut,
+};
+use crate::uid::{Kind as UidKind, Uid};
 use crate::{EntityComponentRef, World};
 
 /// A handle to a entity.
+#[derive(Debug)]
 pub struct Handle<'a>
 {
+    world: &'a World,
     archetype: &'a Archetype,
     entity: &'a ArchetypeEntity,
 }
@@ -21,20 +35,82 @@ impl<'a> Handle<'a>
         self.entity.uid()
     }
 
+    pub fn get<ComponentT: Component>(&self) -> Option<ComponentHandle<'_, ComponentT>>
+    {
+        assert_eq!(ComponentT::id().kind(), UidKind::Component);
+
+        let component = self.get_matching_components(ComponentT::id()).next()?;
+
+        Some(
+            ComponentHandle::from_entity_component_ref(Some(component), self.world)
+                .unwrap_or_else(|err| {
+                    panic!(
+                        "Taking component {} lock failed: {err}",
+                        type_name::<ComponentT>()
+                    );
+                }),
+        )
+    }
+
+    pub fn get_mut<ComponentT: Component>(
+        &self,
+    ) -> Option<ComponentHandleMut<'_, ComponentT>>
+    {
+        assert_eq!(ComponentT::id().kind(), UidKind::Component);
+
+        let component = self.get_matching_components(ComponentT::id()).next()?;
+
+        Some(
+            ComponentHandleMut::from_entity_component_ref(Some(component), self.world)
+                .unwrap_or_else(|err| {
+                    panic!(
+                        "Taking component {} lock failed: {err}",
+                        type_name::<ComponentT>()
+                    );
+                }),
+        )
+    }
+
     #[inline]
     #[must_use]
-    pub fn get_component(&self, component_uid: Uid) -> Option<EntityComponentRef<'a>>
+    pub fn get_matching_components(&self, component_uid: Uid)
+        -> MatchingComponentIter<'a>
     {
-        let index = self.archetype.get_index_for_component(component_uid)?;
+        MatchingComponentIter {
+            inner: self.archetype.get_matching_component_indices(component_uid),
+            entity: self.entity,
+        }
+    }
 
-        Some(EntityComponentRef::new(
-            self.entity.components().get(index).unwrap(),
-        ))
+    pub(crate) fn new(
+        world: &'a World,
+        archetype: &'a Archetype,
+        entity: &'a ArchetypeEntity,
+    ) -> Self
+    {
+        Self { world, archetype, entity }
     }
+}
+
+#[derive(Debug)]
+pub struct MatchingComponentIter<'a>
+{
+    inner: ArchetypeMatchingComponentIter<'a>,
+    entity: &'a ArchetypeEntity,
+}
+
+impl<'a> Iterator for MatchingComponentIter<'a>
+{
+    type Item = EntityComponentRef<'a>;
 
-    pub(crate) fn new(archetype: &'a Archetype, entity: &'a ArchetypeEntity) -> Self
+    fn next(&mut self) -> Option<Self::Item>
     {
-        Self { archetype, entity }
+        let (matching_component_id, index) = self.inner.next()?;
+
+        Some(EntityComponentRef::new(
+            matching_component_id,
+            self.entity.components().get(index).unwrap(),
+        ))
     }
 }