feat(ecs): add action to add component(s) to entity

1 files changed, 87 insertions, 12 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 6e06ded..76e1a7e 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -12,7 +12,6 @@ use crate::component::{
     Metadata as ComponentMetadata,
 };
 use crate::entity::Uid as EntityUid;
-use crate::lock::Lock;
 use crate::type_name::TypeName;
 use crate::util::{RefCellRefMap, Sortable};
 use crate::EntityComponent;
@@ -22,6 +21,7 @@ pub struct Storage
 {
     archetypes: Vec<Archetype>,
     archetype_lookup: RefCell<HashMap<ArchetypeId, ArchetypeLookupEntry>>,
+    entity_archetype_lookup: HashMap<EntityUid, ArchetypeId>,
 }
 
 impl Storage
@@ -96,6 +96,7 @@ impl Storage
     #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
     pub fn push_entity(
         &mut self,
+        entity_uid: EntityUid,
         mut components: Vec<Box<dyn Component>>,
     ) -> (ArchetypeId, EntityUid)
     {
@@ -133,17 +134,46 @@ impl Storage
             .get_mut(archetype_index)
             .expect("Archetype is gone");
 
-        let entity_uid = EntityUid::new_unique();
-
         archetype.push_entity(entity_uid, components);
 
         archetype
             .entity_lookup
             .insert(entity_uid, archetype.entities.len() - 1);
 
+        self.entity_archetype_lookup
+            .insert(entity_uid, archetype_id);
+
         (archetype_id, entity_uid)
     }
 
+    pub fn add_components_to_entity(
+        &mut self,
+        entity_uid: EntityUid,
+        components: impl IntoIterator<Item = Box<dyn Component>>,
+    ) -> Option<()>
+    {
+        let archetype_id = self.entity_archetype_lookup.get(&entity_uid)?;
+
+        let archetype_index =
+            self.find_archetype_index_with_entity(*archetype_id, entity_uid)?;
+
+        let archetype = self.archetypes.get_mut(archetype_index)?;
+
+        let entity = archetype.take_entity(entity_uid)?;
+
+        self.push_entity(
+            entity_uid,
+            entity
+                .components
+                .into_iter()
+                .map(|component| component.component.into_inner())
+                .chain(components.into_iter().map(|component| component.into()))
+                .collect(),
+        );
+
+        Some(())
+    }
+
     fn populate_matching_archetype_lookup_entries(
         &mut self,
         comp_ids_set: &HashSet<ComponentId>,
@@ -193,6 +223,31 @@ impl Storage
         // cannot fail
         unsafe { *lookup_entry.archetype_indices.first().unwrap_unchecked() }
     }
+
+    fn find_archetype_index_with_entity(
+        &self,
+        archetype_id: ArchetypeId,
+        entity_uid: EntityUid,
+    ) -> Option<usize>
+    {
+        let archetype_lookup = self.archetype_lookup.borrow_mut();
+
+        let archetype_lookup_entry = archetype_lookup.get(&archetype_id)?;
+
+        // TODO: Also have a hashmap for precise archetype ID -> archetype index lookup.
+        // This way is slow
+        archetype_lookup_entry
+            .archetype_indices
+            .iter()
+            .find(|archetype_index| {
+                let Some(archetype) = self.archetypes.get(**archetype_index) else {
+                    return false;
+                };
+
+                archetype.has_entity(entity_uid)
+            })
+            .map(|archetype_index| *archetype_index)
+    }
 }
 
 impl TypeName for Storage
@@ -274,14 +329,30 @@ impl Archetype
             uid: entity_uid,
             components: components
                 .into_iter()
-                .map(|component| EntityComponent {
-                    id: component.id(),
-                    name: component.type_name(),
-                    component: Lock::new(component),
-                })
+                .map(|component| component.into())
                 .collect(),
         });
     }
+
+    pub fn take_entity(&mut self, entity_uid: EntityUid) -> Option<ArchetypeEntity>
+    {
+        let entity_index = self.entity_lookup.remove(&entity_uid)?;
+
+        let entity = self.entities.remove(entity_index);
+
+        for index in self.entity_lookup.values_mut() {
+            if *index > entity_index {
+                *index -= 1;
+            }
+        }
+
+        Some(entity)
+    }
+
+    fn has_entity(&self, entity_uid: EntityUid) -> bool
+    {
+        self.entity_lookup.contains_key(&entity_uid)
+    }
 }
 
 #[derive(Debug)]
@@ -385,6 +456,7 @@ mod tests
         IsOptional as ComponentIsOptional,
         Metadata as ComponentMetadata,
     };
+    use crate::entity::Uid as EntityUid;
     use crate::{self as ecs};
 
     #[derive(Debug, Component)]
@@ -419,10 +491,13 @@ mod tests
     {
         let mut component_storage = Storage::default();
 
-        component_storage.push_entity(vec![
-            Box::new(HealthPotion { _hp_restoration: 12 }),
-            Box::new(Hookshot { _range: 50 }),
-        ]);
+        component_storage.push_entity(
+            EntityUid::new_unique(),
+            vec![
+                Box::new(HealthPotion { _hp_restoration: 12 }),
+                Box::new(Hookshot { _range: 50 }),
+            ],
+        );
 
         assert_eq!(component_storage.archetypes.len(), 1);