1 files changed, 66 insertions, 18 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 76701ac..b227337 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -9,6 +9,7 @@ use crate::component::{
     IsOptional as ComponentIsOptional,
     Metadata as ComponentMetadata,
 };
+use crate::entity::Uid as EntityUid;
 use crate::lock::Lock;
 use crate::type_name::TypeName;
 use crate::EntityComponent;
@@ -39,8 +40,13 @@ impl Storage
     }
 
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
-    pub fn push_entity(&mut self, components: Vec<Box<dyn Component>>)
+    pub fn push_entity(
+        &mut self,
+        mut components: Vec<Box<dyn Component>>,
+    ) -> (ArchetypeId, EntityUid)
     {
+        components.sort_by_key(|component| component.id());
+
         #[cfg(feature = "debug")]
         tracing::debug!(
             "Pushing entity with components: ({})",
@@ -51,20 +57,22 @@ impl Storage
                 .join(", ")
         );
 
-        let archetype_indices = self
-            .archetype_lookup
-            .entry(ArchetypeId::from_components_metadata(
-                components
-                    .iter()
-                    .map(|component| ComponentMetadata::of(&**component)),
-            ))
-            .or_insert_with(|| {
-                self.archetypes.push(Archetype::new(
-                    components.iter().map(|component| component.id()),
-                ));
-
-                vec![self.archetypes.len() - 1]
-            });
+        let archetype_id = ArchetypeId::from_components_metadata(
+            components
+                .iter()
+                .map(|component| ComponentMetadata::of(&**component)),
+        );
+
+        let archetype_indices =
+            self.archetype_lookup
+                .entry(archetype_id)
+                .or_insert_with(|| {
+                    self.archetypes.push(Archetype::new(
+                        components.iter().map(|component| component.id()),
+                    ));
+
+                    vec![self.archetypes.len() - 1]
+                });
 
         if archetype_indices.is_empty() {
             self.archetypes.push(Archetype::new(
@@ -94,6 +102,14 @@ impl Storage
                 })
                 .collect(),
         );
+
+        let entity_uid = EntityUid::new_unique();
+
+        archetype
+            .entity_lookup
+            .insert(entity_uid, archetype.components.len() - 1);
+
+        (archetype_id, entity_uid)
     }
 
     pub fn add_archetype_lookup_entry(
@@ -127,7 +143,7 @@ impl Storage
                 .iter()
                 .enumerate()
                 .filter_map(|(index, archetype)| {
-                    if archetype.component_ids.is_superset(&ids_set) {
+                    if archetype.component_ids_is_superset(&ids_set) {
                         return Some(index);
                     }
 
@@ -157,7 +173,8 @@ impl TypeName for Storage
 #[derive(Debug)]
 pub struct Archetype
 {
-    component_ids: HashSet<ComponentId>,
+    component_ids: HashMap<ComponentId, usize>,
+    entity_lookup: HashMap<EntityUid, usize>,
     pub components: Vec<Vec<EntityComponent>>,
 }
 
@@ -166,10 +183,41 @@ impl Archetype
     fn new(component_ids: impl IntoIterator<Item = ComponentId>) -> Self
     {
         Self {
-            component_ids: component_ids.into_iter().collect(),
+            component_ids: component_ids
+                .into_iter()
+                .enumerate()
+                .map(|(index, component_id)| (component_id, index))
+                .collect(),
+            entity_lookup: HashMap::new(),
             components: Vec::new(),
         }
     }
+
+    pub fn component_ids_is_superset(
+        &self,
+        other_component_ids: &HashSet<ComponentId>,
+    ) -> bool
+    {
+        if other_component_ids.len() <= self.component_ids.len() {
+            other_component_ids
+                .iter()
+                .all(|v| self.component_ids.contains_key(v))
+        } else {
+            false
+        }
+    }
+
+    pub fn get_entity(&self, entity_uid: EntityUid) -> Option<&[EntityComponent]>
+    {
+        let entity_index = *self.entity_lookup.get(&entity_uid)?;
+
+        Some(&self.components.get(entity_index)?)
+    }
+
+    pub fn get_index_for_component(&self, component_id: &ComponentId) -> Option<usize>
+    {
+        self.component_ids.get(component_id).copied()
+    }
 }
 
 #[derive(Debug)]