From 337ce18e8c731afba47900d98a5b6f35846ebe25 Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Mon, 22 Apr 2024 00:20:42 -0400
Subject: [PATCH] ecs: update existing components on entity in World::insert

---
 lyra-ecs/src/query/filter/not.rs |  2 +-
 lyra-ecs/src/query/filter/or.rs  |  2 +-
 lyra-ecs/src/world.rs            | 54 +++++++++++++++++++++++++-------
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/lyra-ecs/src/query/filter/not.rs b/lyra-ecs/src/query/filter/not.rs
index 4257261..a50550b 100644
--- a/lyra-ecs/src/query/filter/not.rs
+++ b/lyra-ecs/src/query/filter/not.rs
@@ -5,7 +5,7 @@ use crate::{query::{AsQuery, Query}, Archetype, World};
 /// This means that entities that `Q` fetches are skipped, and entities that
 /// `Q` does not fetch are not skipped.
 /// 
-/// ```rust,nobuild
+/// ```nobuild
 /// // Iterate over entities that has a transform, and are not the origin of a `ChildOf` relationship.
 /// for (en, pos, _) in world
 ///     .view::<(Entities, &Transform, Not<Has<RelationOriginComponent<ChildOf>>>)>()
diff --git a/lyra-ecs/src/query/filter/or.rs b/lyra-ecs/src/query/filter/or.rs
index fd1e8f2..3fc38b8 100644
--- a/lyra-ecs/src/query/filter/or.rs
+++ b/lyra-ecs/src/query/filter/or.rs
@@ -36,7 +36,7 @@ impl<'a, Q1: Query, Q2: Query> Fetch<'a> for OrFetch<'a, Q1, Q2> {
 /// 
 /// This checks if `Q1` can fetch before checking `Q2`.
 /// 
-/// ```rust,nobuild
+/// ```nobuild
 /// for (en, pos, _) in world
 ///     .view::<(Entities, &Transform, Or<Has<Mesh>, Has<Scene>>)>()
 ///     .iter()
diff --git a/lyra-ecs/src/world.rs b/lyra-ecs/src/world.rs
index ee6ebcd..45e85cd 100644
--- a/lyra-ecs/src/world.rs
+++ b/lyra-ecs/src/world.rs
@@ -1,4 +1,4 @@
-use std::{any::TypeId, collections::HashMap, ptr::NonNull};
+use std::{any::{Any, TypeId}, collections::HashMap, ptr::NonNull};
 
 use atomic_refcell::{AtomicRef, AtomicRefMut};
 
@@ -128,22 +128,39 @@ impl World {
         }
     }
 
-    /// Insert a bundle into an existing entity. If the components are already existing on the
-    /// entity, they will be updated, else the entity will be moved to a different Archetype
-    /// that can store the entity. That may involve creating a new Archetype.
+    /// Insert a component bundle into an existing entity.
+    /// 
+    /// If the components are already existing on the entity, they will be updated, else the
+    /// entity will be moved to a different Archetype that can store the entity. That may
+    /// involve creating a new Archetype.
     pub fn insert<B>(&mut self, entity: Entity, bundle: B)
     where
         B: Bundle
     {
-        // TODO: If the entity already has the components in `bundle`, update the values of the
-        // components with the bundle.
-
         let tick = self.tick();
 
         let record = self.entities.entity_record(entity).unwrap();
         let current_arch = self.archetypes.get(&record.id).unwrap();
         let current_arch_len = current_arch.len();
 
+        let mut contains_all = true;
+        for id in bundle.type_ids() {
+            contains_all = contains_all && current_arch.get_column(id).is_some();
+        }
+
+        if contains_all {
+            let current_arch = self.archetypes.get_mut(&record.id).unwrap();
+            let entry_idx = *current_arch.entity_indexes()
+                .get(&entity).unwrap();
+
+            bundle.take(|ptr, id, _info| {
+                let col = current_arch.get_column_mut(id).unwrap();
+                unsafe { col.set_at(entry_idx.0 as _, ptr, tick) };
+            });
+
+            return;
+        }
+
         // contains the type ids for the old component columns + the ids for the new components
         let mut combined_column_types: Vec<DynTypeId> = current_arch.columns.iter().map(|c| c.info.type_id()).collect();
         combined_column_types.extend(bundle.type_ids());
@@ -227,7 +244,7 @@ impl World {
     /// A method used for debugging implementation details of the ECS.
     /// 
     /// Here's an example of the output:
-    /// ```
+    /// ```nobuild
     /// Entities
     ///   1 in archetype 0 at 0
     ///   0 in archetype 1 at 0
@@ -271,7 +288,7 @@ impl World {
     /// the contents of entities inside the archetypes.
     /// 
     /// Below is a template of the output:
-    /// ```
+    /// ```nobuild
     /// Entities
     ///   %ENTITY_ID% in archetype %ARCHETYPE_ID% at %INDEX%
     /// Arch ID -- %ARCHETYPE_LEN% entities
@@ -618,8 +635,6 @@ mod tests {
 
         insert_and_assert(&mut world, e1, v2s[0], v3s[0]);
         println!("Entity 1 is good");
-        assert_eq!(world.archetypes.len(), 1);
-        println!("Empty archetype was removed");
     }
 
     #[test]
@@ -655,4 +670,21 @@ mod tests {
             assert!(tick >= world_tick);
         }
     }
+
+    /// Tests replacing components using World::insert
+    #[test]
+    fn entity_insert_replace() {
+        let mut world = World::new();
+        let first = world.spawn((Vec2::new(10.0, 10.0),));
+        let second = world.spawn((Vec2::new(5.0, 5.0),));
+
+        world.insert(first, Vec2::new(50.0, 50.0));
+
+        let pos = world.view_one::<&mut Vec2>(first).get().unwrap();
+        assert_eq!(*pos, Vec2::new(50.0, 50.0));
+        drop(pos);
+
+        let pos = world.view_one::<&mut Vec2>(second).get().unwrap();
+        assert_eq!(*pos, Vec2::new(5.0, 5.0));
+    }
 }
\ No newline at end of file