1 files changed, 28 insertions, 26 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index c70e7e7..40909fb 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>
@@ -195,10 +192,24 @@ impl Storage
             });
         }
 
-        let add_edge_archetype_id = archetype_node
+        let add_edge_archetype_id = match archetype_node
             .get_or_insert_edges(component_id, ArchetypeEdges::default)
             .add
-            .unwrap_or_else(|| {
+        {
+            Some(add_edge_id) => {
+                if !self.graph.contains_archetype(add_edge_id) {
+                    let (_, add_edge_comp_ids) = self
+                        .graph
+                        .get_node_by_id(archetype_id)
+                        .expect("Archetype should exist")
+                        .make_add_edge(component_id);
+
+                    self.graph.create_node(add_edge_id, &add_edge_comp_ids);
+                }
+
+                add_edge_id
+            }
+            None => {
                 let archetype_node = self
                     .graph
                     .get_node_by_id_mut(archetype_id)
@@ -217,7 +228,8 @@ impl Storage
                 }
 
                 add_edge_id
-            });
+            }
+        };
 
         {
             let add_edge_archetype_node = self
@@ -247,13 +259,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 +335,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)