1 files changed, 49 insertions, 12 deletions

diff --git a/ecs/src/entity.rs b/ecs/src/entity.rs
index 34fd19f..ca867f9 100644
--- a/ecs/src/entity.rs
+++ b/ecs/src/entity.rs
@@ -21,6 +21,7 @@ pub struct Handle<'a>
 {
     archetype: &'a Archetype,
     entity: &'a ArchetypeEntity,
+    world: &'a World,
 }
 
 impl<'a> Handle<'a>
@@ -41,7 +42,7 @@ impl<'a> Handle<'a>
     /// - 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>>
+    pub fn get<ComponentT: Component>(&self) -> Option<ComponentHandle<'a, ComponentT>>
     {
         assert_eq!(ComponentT::id().kind(), UidKind::Component);
 
@@ -76,12 +77,44 @@ impl<'a> Handle<'a>
         let component = self.get_matching_components(ComponentT::id()).next()?;
 
         Some(
-            ComponentHandleMut::from_entity_component_ref(&component).unwrap_or_else(
-                |err| {
+            ComponentHandleMut::from_entity_component_ref(&component, self.world)
+                .unwrap_or_else(|err| {
                     panic!(
                         "Creating handle to component {} failed: {err}",
                         type_name::<ComponentT>()
                     );
+                }),
+        )
+    }
+
+    /// Returns a reference to the component with the ID `id` in this entity.
+    /// `None` is returned if the component isn't found.
+    ///
+    /// # Panics
+    /// Will panic if:
+    /// - The ID is not a component/pair ID
+    /// - The component is borrowed mutably elsewhere
+    /// - The component type is incorrect
+    #[must_use]
+    pub fn get_with_id<ComponentDataT: 'static>(
+        &self,
+        id: Uid,
+    ) -> Option<ComponentHandle<'a, ComponentDataT>>
+    {
+        assert!(
+            matches!(id.kind(), UidKind::Component | UidKind::Pair),
+            "ID {id:?} is not a component/pair ID"
+        );
+
+        let component = self.get_matching_components(id).next()?;
+
+        Some(
+            ComponentHandle::from_entity_component_ref(&component).unwrap_or_else(
+                |err| {
+                    panic!(
+                        "Creating handle to component {} failed: {err}",
+                        type_name::<ComponentDataT>()
+                    );
                 },
             ),
         )
@@ -96,10 +129,10 @@ impl<'a> Handle<'a>
     /// - The component is borrowed elsewhere
     /// - The component type is incorrect
     #[must_use]
-    pub fn get_with_id_mut<ComponentData: 'static>(
+    pub fn get_with_id_mut<ComponentDataT: 'static>(
         &self,
         id: Uid,
-    ) -> Option<ComponentHandleMut<'a, ComponentData>>
+    ) -> Option<ComponentHandleMut<'a, ComponentDataT>>
     {
         assert!(
             matches!(id.kind(), UidKind::Component | UidKind::Pair),
@@ -109,14 +142,13 @@ impl<'a> Handle<'a>
         let component = self.get_matching_components(id).next()?;
 
         Some(
-            ComponentHandleMut::from_entity_component_ref(&component).unwrap_or_else(
-                |err| {
+            ComponentHandleMut::from_entity_component_ref(&component, self.world)
+                .unwrap_or_else(|err| {
                     panic!(
                         "Creating handle to component {} failed: {err}",
-                        type_name::<ComponentData>()
+                        type_name::<ComponentDataT>()
                     );
-                },
-            ),
+                }),
         )
     }
 
@@ -131,9 +163,13 @@ impl<'a> Handle<'a>
         }
     }
 
-    pub(crate) fn new(archetype: &'a Archetype, entity: &'a ArchetypeEntity) -> Self
+    pub(crate) fn new(
+        archetype: &'a Archetype,
+        entity: &'a ArchetypeEntity,
+        world: &'a World,
+    ) -> Self
     {
-        Self { archetype, entity }
+        Self { archetype, entity, world }
     }
 }
 
@@ -155,6 +191,7 @@ impl<'a> Iterator for MatchingComponentIter<'a>
         Some(EntityComponentRef::new(
             matching_component_id,
             self.entity.components().get(index).unwrap(),
+            self.entity.uid(),
         ))
     }
 }