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

4 files changed, 129 insertions, 14 deletions

diff --git a/ecs/src/actions.rs b/ecs/src/actions.rs
index 2ee5518..038f617 100644
--- a/ecs/src/actions.rs
+++ b/ecs/src/actions.rs
@@ -3,6 +3,7 @@ use std::marker::PhantomData;
 use std::sync::{Arc, Weak};
 
 use crate::component::{Component, Sequence as ComponentSequence};
+use crate::entity::Uid as EntityUid;
 use crate::lock::{Lock, WriteGuard};
 use crate::system::{NoInitParamFlag, Param as SystemParam, System};
 use crate::{ActionQueue, World};
@@ -23,6 +24,15 @@ impl<'world> Actions<'world>
         self.action_queue.push(Action::Spawn(components.into_vec()));
     }
 
+    /// Adds component(s) to a entity.
+    pub fn add_components<Comps>(&mut self, entity_uid: EntityUid, components: Comps)
+    where
+        Comps: ComponentSequence,
+    {
+        self.action_queue
+            .push(Action::AddComponents(entity_uid, components.into_vec()));
+    }
+
     /// Adds stopping the loop in [`Engine::event_loop`] at the next opportune time to the
     /// action queue.
     pub fn stop(&mut self)
@@ -129,6 +139,7 @@ impl<'weak_ref> Ref<'weak_ref>
 pub(crate) enum Action
 {
     Spawn(Vec<Box<dyn Component>>),
+    AddComponents(EntityUid, Vec<Box<dyn Component>>),
     Stop,
 }
 
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);
 
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 34d6b7b..df65cb3 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -71,7 +71,7 @@ impl World
             .component_storage
             .write_nonblock()
             .expect("Failed to acquire read-write component storage lock")
-            .push_entity(components.into_vec());
+            .push_entity(EntityUid::new_unique(), components.into_vec());
 
         entity_uid
     }
@@ -158,7 +158,17 @@ impl World
                             "Failed to acquire read-write component storage lock",
                         );
 
-                    component_storage_lock.push_entity(components);
+                    component_storage_lock
+                        .push_entity(EntityUid::new_unique(), components);
+                }
+                Action::AddComponents(entity_uid, components) => {
+                    let mut component_storage_lock =
+                        self.data.component_storage.write_nonblock().expect(
+                            "Failed to acquire read-write component storage lock",
+                        );
+
+                    component_storage_lock
+                        .add_components_to_entity(entity_uid, components);
                 }
                 Action::Stop => {
                     self.stop.store(true, Ordering::Relaxed);
@@ -226,6 +236,18 @@ pub struct EntityComponent
     pub component: Lock<Box<dyn Component>>,
 }
 
+impl From<Box<dyn Component>> for EntityComponent
+{
+    fn from(component: Box<dyn Component>) -> Self
+    {
+        Self {
+            id: component.id(),
+            name: component.type_name(),
+            component: Lock::new(component),
+        }
+    }
+}
+
 #[derive(Debug, Default)]
 struct ActionQueue
 {
diff --git a/ecs/src/lock.rs b/ecs/src/lock.rs
index b3c9f57..77213c9 100644
--- a/ecs/src/lock.rs
+++ b/ecs/src/lock.rs
@@ -57,6 +57,13 @@ where
 
         Ok(WriteGuard { inner: guard })
     }
+
+    pub fn into_inner(self) -> Value
+    {
+        self.inner
+            .into_inner()
+            .unwrap_or_else(|err| err.into_inner())
+    }
 }
 
 #[derive(Debug, thiserror::Error)]