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 +++++++++++++++++++++++++++++++-----------
 1 file changed, 31 insertions(+), 11 deletions(-)

(limited to 'ecs/src/component/storage.rs')

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);
 
-- 
cgit v1.2.3-18-g5258