1 files changed, 102 insertions, 9 deletions

diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index 85e7461..bab3d61 100644
--- a/ecs/src/entity.rs
+++ b/ecs/src/entity.rs
@@ -1,10 +1,22 @@
+use std::any::type_name;
+
 use linkme::distributed_slice;
 
-use crate::component::storage::archetype::{Archetype, ArchetypeEntity};
-use crate::uid::Uid;
-use crate::{EntityComponent, World};
+use crate::component::storage::archetype::{
+    Archetype,
+    Entity as ArchetypeEntity,
+    MatchingComponentIter as ArchetypeMatchingComponentIter,
+};
+use crate::component::{
+    Component,
+    Handle as ComponentHandle,
+    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>
 {
     archetype: &'a Archetype,
@@ -18,16 +30,75 @@ impl<'a> Handle<'a>
     #[must_use]
     pub fn uid(&self) -> Uid
     {
-        self.entity.uid
+        self.entity.uid()
+    }
+
+    /// Returns a reference to the specified component in this entity. `None` is
+    /// returned if the component isn't found in the entity.
+    ///
+    /// # Panics
+    /// Will panic if:
+    /// - The component's ID is not a component ID
+    /// - The component is mutably borrowed elsewhere
+    #[must_use]
+    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(component).unwrap_or_else(|err| {
+                panic!(
+                    "Taking component {} lock failed: {err}",
+                    type_name::<ComponentT>()
+                );
+            }),
+        )
+    }
+
+    /// Returns a mutable reference to the specified component in this entity. `None` is
+    /// returned if the component isn't found in the entity.
+    ///
+    /// # Panics
+    /// Will panic if:
+    /// - The component's ID is not a component ID
+    /// - The component is borrowed elsewhere
+    #[must_use]
+    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(component).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<&'a EntityComponent>
+    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(self.entity.components.get(index).unwrap())
+    pub fn component_ids(&self) -> impl Iterator<Item = Uid> + '_
+    {
+        self.archetype.component_ids_sorted()
     }
 
     pub(crate) fn new(archetype: &'a Archetype, entity: &'a ArchetypeEntity) -> Self
@@ -36,6 +107,28 @@ impl<'a> Handle<'a>
     }
 }
 
+#[derive(Debug)]
+pub struct MatchingComponentIter<'a>
+{
+    inner: ArchetypeMatchingComponentIter<'a>,
+    entity: &'a ArchetypeEntity,
+}
+
+impl<'a> Iterator for MatchingComponentIter<'a>
+{
+    type Item = EntityComponentRef<'a>;
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        let (matching_component_id, index) = self.inner.next()?;
+
+        Some(EntityComponentRef::new(
+            matching_component_id,
+            self.entity.components().get(index).unwrap(),
+        ))
+    }
+}
+
 #[allow(clippy::module_name_repetitions)]
 #[macro_export]
 macro_rules! static_entity {
@@ -53,7 +146,7 @@ macro_rules! static_entity {
                     $crate::entity::CREATE_STATIC_ENTITIES
                 )]
                 #[linkme(crate=$crate::private::linkme)]
-                static CREATE_STATIC_ENTITY: fn(&$crate::World) = |world| {
+                static CREATE_STATIC_ENTITY: fn(&mut $crate::World) = |world| {
                     world.create_entity_with_uid($components, *$ident);
                 };
             }
@@ -63,4 +156,4 @@ macro_rules! static_entity {
 
 #[distributed_slice]
 #[doc(hidden)]
-pub static CREATE_STATIC_ENTITIES: [fn(&World)];
+pub static CREATE_STATIC_ENTITIES: [fn(&mut World)];