1 files changed, 100 insertions, 9 deletions

diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index a43f9ce..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, 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,
+    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,
@@ -21,15 +33,72 @@ impl<'a> Handle<'a>
         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<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 fn component_ids(&self) -> impl Iterator<Item = Uid> + '_
+    {
+        self.archetype.component_ids_sorted()
     }
 
     pub(crate) fn new(archetype: &'a Archetype, entity: &'a ArchetypeEntity) -> Self
@@ -38,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 {
@@ -55,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);
                 };
             }
@@ -65,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)];