3 files changed, 114 insertions, 40 deletions

diff --git a/ecs/src/component/local.rs b/ecs/src/component/local.rs
index d4e414e..0f6f641 100644
--- a/ecs/src/component/local.rs
+++ b/ecs/src/component/local.rs
@@ -1,14 +1,14 @@
 use std::ops::{Deref, DerefMut};
 
-use crate::component::Component;
-use crate::system::{ComponentRefMut, Param as SystemParam, System};
+use crate::component::{Component, HandleMut as ComponentHandleMut};
+use crate::system::{Param as SystemParam, System};
 use crate::World;
 
 /// Holds a component which is local to a single system.
 #[derive(Debug)]
 pub struct Local<'world, LocalComponent: Component>
 {
-    local_component: ComponentRefMut<'world, LocalComponent>,
+    local_component: ComponentHandleMut<'world, LocalComponent>,
 }
 
 impl<'world, LocalComponent> SystemParam<'world> for Local<'world, LocalComponent>
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)
diff --git a/ecs/src/component/storage/archetype.rs b/ecs/src/component/storage/archetype.rs
index f6f8132..5306cf9 100644
--- a/ecs/src/component/storage/archetype.rs
+++ b/ecs/src/component/storage/archetype.rs
@@ -3,16 +3,16 @@ use std::slice::Iter as SliceIter;
 
 use hashbrown::HashMap;
 
-use crate::component::Metadata as ComponentMetadata;
+use crate::component::{Component, Metadata as ComponentMetadata};
+use crate::lock::Lock;
 use crate::uid::{Kind as UidKind, Uid};
 use crate::util::HashMapExt;
-use crate::EntityComponent;
 
 #[derive(Debug)]
 pub struct Archetype
 {
     id: Id,
-    entities: Vec<ArchetypeEntity>,
+    entities: Vec<Entity>,
     entity_index_lookup: HashMap<Uid, usize>,
     component_index_lookup: HashMap<Uid, usize>,
     component_ids: Vec<Uid>,
@@ -53,7 +53,7 @@ impl Archetype
             .keys_is_subset(&other.component_index_lookup)
     }
 
-    pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&ArchetypeEntity>
+    pub fn get_entity_by_id(&self, entity_uid: Uid) -> Option<&Entity>
     {
         let index = *self.entity_index_lookup.get(&entity_uid)?;
 
@@ -64,7 +64,7 @@ impl Archetype
         }))
     }
 
-    pub fn push_entity(&mut self, entity: ArchetypeEntity)
+    pub fn push_entity(&mut self, entity: Entity)
     {
         self.entity_index_lookup
             .insert(entity.uid, self.entities.len());
@@ -72,7 +72,7 @@ impl Archetype
         self.entities.push(entity);
     }
 
-    pub fn remove_entity(&mut self, entity_uid: Uid) -> Option<ArchetypeEntity>
+    pub fn remove_entity(&mut self, entity_uid: Uid) -> Option<Entity>
     {
         //debug_assert_eq!(entity_uid.kind(), UidKind::Entity);
 
@@ -142,12 +142,12 @@ impl Archetype
 #[derive(Debug)]
 pub struct EntityIter<'archetype>
 {
-    iter: SliceIter<'archetype, ArchetypeEntity>,
+    iter: SliceIter<'archetype, Entity>,
 }
 
 impl<'archetype> Iterator for EntityIter<'archetype>
 {
-    type Item = &'archetype ArchetypeEntity;
+    type Item = &'archetype Entity;
 
     fn next(&mut self) -> Option<Self::Item>
     {
@@ -156,10 +156,82 @@ impl<'archetype> Iterator for EntityIter<'archetype>
 }
 
 #[derive(Debug)]
-pub struct ArchetypeEntity
+pub struct Entity
 {
-    pub uid: Uid,
-    pub components: Vec<EntityComponent>,
+    uid: Uid,
+    components: Vec<EntityComponent>,
+}
+
+impl Entity
+{
+    pub fn new(uid: Uid, components: impl IntoIterator<Item = EntityComponent>) -> Self
+    {
+        Self {
+            uid,
+            components: components.into_iter().collect(),
+        }
+    }
+
+    pub fn uid(&self) -> Uid
+    {
+        self.uid
+    }
+
+    pub fn components(&self) -> &[EntityComponent]
+    {
+        &self.components
+    }
+
+    pub fn remove_component(&mut self, component_id: Uid, archetype: &Archetype)
+    {
+        let index = archetype
+            .get_index_for_component(component_id)
+            .expect("Archetype should contain component");
+
+        self.components.remove(index);
+    }
+
+    pub fn insert_component(
+        &mut self,
+        component_id: Uid,
+        component: EntityComponent,
+        archetype: &Archetype,
+    )
+    {
+        let index = archetype
+            .get_index_for_component(component_id)
+            .expect("Archetype should contain component");
+
+        self.components.insert(index, component);
+    }
+}
+
+#[derive(Debug)]
+pub struct EntityComponent
+{
+    name: &'static str,
+    component: Lock<Box<dyn Component>>,
+}
+
+impl EntityComponent
+{
+    pub fn new(component: Box<dyn Component>) -> Self
+    {
+        Self {
+            name: component.type_name(),
+            component: Lock::new(component),
+        }
+    }
+
+    pub fn name(&self) -> &str
+    {
+        self.name
+    }
+
+    pub fn component(&self) -> &Lock<Box<dyn Component>>
+    {
+        &self.component
+    }
 }
 
 /// Archetype ID.