perf(ecs): make Relation get method use entity archetype lookup

2 files changed, 17 insertions, 17 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 5d9bb06..1bb7fb5 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -72,6 +72,16 @@ impl Storage
         self.iter_archetypes_by_lookup(&archetype_id)
     }
 
+    pub fn get_entity_archetype(&self, entity_uid: EntityUid) -> Option<&Archetype>
+    {
+        let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;
+
+        let archetype_index =
+            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;
+
+        self.archetypes.get(archetype_index)
+    }
+
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
     pub fn push_entity(
         &mut self,
diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index fa3f761..e49bde7 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -3,11 +3,9 @@ use std::marker::PhantomData;
 
 use crate::component::storage::Storage as ComponentStorage;
 use crate::component::{
-    is_optional as component_is_optional,
     Component,
     FromOptional as ComponentFromOptional,
     Id as ComponentId,
-    Metadata as ComponentMetadata,
 };
 use crate::entity::Uid as EntityUid;
 use crate::lock::ReadGuard;
@@ -122,21 +120,13 @@ where
     /// Retrieves the related-to component.
     pub fn get(&self) -> Option<ComponentRefMut<'_, ComponentT>>
     {
-        let mut archetype_iter =
-            self.component_storage_lock
-                .find_entities([ComponentMetadata {
-                    id: ComponentId::of::<ComponentT>(),
-                    is_optional: component_is_optional::<ComponentT>().into(),
-                }]);
-
-        let (entity, archetype) = archetype_iter.find_map(|archetype| {
-            let Some(entity) = archetype.get_entity(self.relationship_comp.entity_uid)
-            else {
-                return None;
-            };
-
-            Some((entity, archetype))
-        })?;
+        let archetype = self
+            .component_storage_lock
+            .get_entity_archetype(self.relationship_comp.entity_uid)?;
+
+        let entity = archetype
+            .get_entity(self.relationship_comp.entity_uid)
+            .expect("Entity is gone");
 
         let component_index =
             archetype.get_index_for_component(&ComponentId::of::<ComponentT>())?;