1 files changed, 10 insertions, 23 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index c70e7e7..4a00510 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -7,7 +7,8 @@ use hashbrown::HashMap;
 
 use crate::component::storage::archetype::{
     Archetype,
-    ArchetypeEntity,
+    Entity as ArchetypeEntity,
+    EntityComponent as ArchetypeEntityComponent,
     Id as ArchetypeId,
 };
 use crate::component::storage::graph::{
@@ -20,7 +21,6 @@ use crate::component::Component;
 use crate::type_name::TypeName;
 use crate::uid::{Kind as UidKind, Uid};
 use crate::util::{BorrowedOrOwned, Either, StreamingIterator, VecExt};
-use crate::EntityComponent;
 
 pub mod archetype;
 
@@ -123,17 +123,14 @@ impl Storage
 
         archetype_node
             .archetype_mut()
-            .push_entity(ArchetypeEntity { uid, components: vec![] });
+            .push_entity(ArchetypeEntity::new(uid, []));
 
         self.entity_archetype_lookup.insert(uid, empty_archetype_id);
 
         Ok(())
     }
 
-    pub fn remove_entity(
-        &mut self,
-        entity_uid: Uid,
-    ) -> Result<Vec<EntityComponent>, Error>
+    pub fn remove_entity(&mut self, entity_uid: Uid) -> Result<ArchetypeEntity, Error>
     {
         let Some(archetype_id) = self.entity_archetype_lookup.get(&entity_uid) else {
             return Err(Error::EntityDoesNotExist(entity_uid));
@@ -151,7 +148,7 @@ impl Storage
 
         self.entity_archetype_lookup.remove(&entity_uid);
 
-        Ok(entity.components)
+        Ok(entity)
     }
 
     pub fn get_entity_archetype(&self, entity_uid: Uid) -> Option<&Archetype>
@@ -247,13 +244,10 @@ impl Storage
             .expect("Add edge archetype should exist")
             .archetype_mut();
 
-        let component_index = add_edge_archetype
-            .get_index_for_component(component_id)
-            .expect("Archetype should have index for component");
-
-        entity.components.insert(
-            component_index,
-            EntityComponent::new(component_id, component),
+        entity.insert_component(
+            component_id,
+            ArchetypeEntityComponent::new(component),
+            add_edge_archetype,
         );
 
         add_edge_archetype.push_entity(entity);
@@ -326,14 +320,7 @@ impl Storage
             .remove_entity(entity_uid)
             .expect("Entity should exist in archetype");
 
-        let (comp_to_remove_index, _) = entity
-            .components
-            .iter()
-            .enumerate()
-            .find(|(_, comp)| comp.id == component_id)
-            .expect("Entity should contain component");
-
-        entity.components.remove(comp_to_remove_index);
+        entity.remove_component(component_id, archetype_node.archetype());
 
         self.graph
             .get_node_by_id_mut(remove_edge_id)