From 61dfbc23b9d68b39eee388fd0bc7df4d02da61f9 Mon Sep 17 00:00:00 2001
From: HampusM <hampus@hampusmat.com>
Date: Sat, 2 Nov 2024 20:15:20 +0100
Subject: fix(ecs): add check if entity already exists before creating it

---
 ecs/src/component/storage.rs | 42 +++++++++++++++++++++++++++++++-----------
 ecs/src/lib.rs               | 26 ++++++++++++++++++++++----
 2 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 00ba330..041c124 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -87,8 +87,12 @@ impl Storage
         &mut self,
         entity_uid: EntityUid,
         mut components: Vec<Box<dyn Component>>,
-    ) -> (ArchetypeId, EntityUid)
+    ) -> Result<(ArchetypeId, EntityUid), Error>
     {
+        if self.entity_archetype_lookup.contains_key(&entity_uid) {
+            return Err(Error::EntityAlreadyExists(entity_uid));
+        }
+
         components.sort_by_key(|component| component.id());
 
         #[cfg(feature = "debug")]
@@ -132,7 +136,7 @@ impl Storage
         self.entity_archetype_lookup
             .insert(entity_uid, archetype_id);
 
-        (archetype_id, entity_uid)
+        Ok((archetype_id, entity_uid))
     }
 
     pub fn add_components_to_entity(
@@ -171,6 +175,8 @@ impl Storage
 
         let entity = archetype.take_entity(entity_uid)?;
 
+        self.entity_archetype_lookup.remove(&entity_uid);
+
         self.push_entity(
             entity_uid,
             entity
@@ -179,7 +185,8 @@ impl Storage
                 .map(|component| component.component.into_inner())
                 .chain(components)
                 .collect(),
-        );
+        )
+        .expect("Not supposed to return Err since the entity is removed");
 
         Some(())
     }
@@ -201,6 +208,8 @@ impl Storage
 
         let component_ids_set = component_ids.into_iter().collect::<HashSet<_>>();
 
+        self.entity_archetype_lookup.remove(&entity_uid);
+
         self.push_entity(
             entity_uid,
             entity
@@ -209,7 +218,8 @@ impl Storage
                 .map(|component| component.component.into_inner())
                 .filter(|component| !component_ids_set.contains(&component.id()))
                 .collect(),
-        );
+        )
+        .expect("Not supposed to return Err since the entity is removed");
 
         Some(())
     }
@@ -315,6 +325,14 @@ impl Storage
     }
 }
 
+/// Component storage error
+#[derive(Debug, Clone, thiserror::Error)]
+pub enum Error
+{
+    #[error("Entity already exists")]
+    EntityAlreadyExists(EntityUid),
+}
+
 impl TypeName for Storage
 {
     fn type_name(&self) -> &'static str
@@ -549,13 +567,15 @@ mod tests
     {
         let mut component_storage = Storage::default();
 
-        component_storage.push_entity(
-            EntityUid::new_unique(),
-            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 }),
+                ],
+            )
+            .expect("Expected Ok");
 
         assert_eq!(component_storage.archetypes.len(), 1);
 
diff --git a/ecs/src/lib.rs b/ecs/src/lib.rs
index 452f657..3c517dd 100644
--- a/ecs/src/lib.rs
+++ b/ecs/src/lib.rs
@@ -75,17 +75,27 @@ impl World
     ///
     /// # Panics
     /// Will panic if mutable internal lock cannot be acquired.
+    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
     pub fn create_entity<Comps>(&mut self, components: Comps) -> EntityUid
     where
         Comps: ComponentSequence + TupleReduce<TypeTransformComponentsToAddedEvents>,
         Comps::Out: EventSequence,
     {
-        let (_, entity_uid) = self
+        let entity_uid = EntityUid::new_unique();
+
+        #[allow(unused_variables)]
+        if let Err(err) = self
             .data
             .component_storage
             .write_nonblock()
             .expect("Failed to acquire read-write component storage lock")
-            .push_entity(EntityUid::new_unique(), components.into_vec());
+            .push_entity(entity_uid, components.into_vec())
+        {
+            #[cfg(feature = "debug")]
+            tracing::error!("Failed to create entity: {err}");
+
+            return entity_uid;
+        };
 
         for component_added_event_id in <Comps::Out as EventSequence>::ids().iter() {
             self.emit_event_by_id(*component_added_event_id);
@@ -160,6 +170,7 @@ impl World
     ///
     /// # Panics
     /// Will panic if a mutable internal lock cannot be acquired.
+    #[cfg_attr(feature = "debug", tracing::instrument(skip_all))]
     pub fn perform_queued_actions(&self)
     {
         let mut active_action_queue = match *self.data.action_queue.active_queue.borrow()
@@ -187,8 +198,15 @@ impl World
                         .map(|component| component.id())
                         .collect::<Vec<_>>();
 
-                    component_storage_lock
-                        .push_entity(EntityUid::new_unique(), components);
+                    #[allow(unused_variables)]
+                    if let Err(err) = component_storage_lock
+                        .push_entity(EntityUid::new_unique(), components)
+                    {
+                        #[cfg(feature = "debug")]
+                        tracing::error!("Failed to create entity: {err}");
+
+                        continue;
+                    }
 
                     drop(component_storage_lock);
 
-- 
cgit v1.2.3-18-g5258