Repository: Ralith/hecs
Pull request #82 (2020-10-09 16:11:23)

EntityBuilder::add_bundle, swap to `TypeInfo` w/ `DynamicBundle::put` and `Bundle::get`

googlebot: 
Thanks for your pull request. It looks like this may be your first contribution to a Google open source project (if not, look below for help). Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

:memo: **Please visit <https://cla.developers.google.com/> to sign.**

Once you've signed (or fixed any issues), please reply here with `@googlebot I signed it!` and we'll verify it.

----

#### What to do if you already signed the CLA

##### Individual signers

*   It's possible we don't have your GitHub username or you're using a different email address on your commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).

##### Corporate signers

*   Your company has a Point of Contact who decides which employees are authorized to participate. Ask your POC to be added to the group of authorized contributors. If you don't know who your Point of Contact is, direct the Google project maintainer to [go/cla#troubleshoot](http://go/cla#troubleshoot) ([Public version](https://opensource.google/docs/cla/#troubleshoot)).
*   The email used to register you as an authorized contributor must be the email used for the Git commit. Check [your existing CLA data](https://cla.developers.google.com/clas) and verify that your [email is set on your git commits](https://help.github.com/articles/setting-your-email-in-git/).
*   The email used to register you as an authorized contributor must also be [attached to your GitHub account](https://github.com/settings/emails).
		

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2FRalith%2Fhecs%2Fpull%2F82) for more info**.

<!-- need_sender_cla -->
sdleffler: @googlebot I signed it!
googlebot: CLAs look good, thanks!

ℹ️ **Googlers: [Go here](https://goto.google.com/prinfo/https%3A%2F%2Fgithub.com%2FRalith%2Fhecs%2Fpull%2F82) for more info**.

<!-- ok -->

Patch:
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index df7f25a..cfb7b00 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -55,9 +55,9 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             }
 
             #[allow(clippy::forget_copy)]
-            unsafe fn put(mut self, mut f: impl FnMut(*mut u8, std::any::TypeId, usize)) {
+            unsafe fn put(mut self, mut f: impl FnMut(*mut u8, ::hecs::TypeInfo)) {
                 #(
-                    f((&mut self.#fields as *mut #tys).cast::<u8>(), std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>());
+                    f((&mut self.#fields as *mut #tys).cast::<u8>(), ::hecs::TypeInfo::of::<#tys>());
                     std::mem::forget(self.#fields);
                 )*
             }
@@ -97,10 +97,10 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             }
 
             unsafe fn get(
-                mut f: impl FnMut(std::any::TypeId, usize) -> Option<std::ptr::NonNull<u8>>,
+                mut f: impl FnMut(::hecs::TypeInfo) -> Option<std::ptr::NonNull<u8>>,
             ) -> Result<Self, ::hecs::MissingComponent> {
                 #(
-                    let #fields = f(std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>())
+                    let #fields = f(::hecs::TypeInfo::of::<#tys>())
                             .ok_or_else(::hecs::MissingComponent::new::<#tys>)?
                             .cast::<#tys>()
                         .as_ptr();
diff --git a/src/archetype.rs b/src/archetype.rs
index 1f8a9c7..91e0389 100644
--- a/src/archetype.rs
+++ b/src/archetype.rs
@@ -137,6 +137,21 @@ impl Archetype {
         &self.types
     }
 
+    /// Enumerate the types of the components of entities stored in this archetype.
+    ///
+    /// Convenient for dispatching logic which needs to be performed on sets of type ids.
+    /// For example, suppose you're building a scripting system, and you want to integrate
+    /// the scripting language with your ECS. This functionality allows you to iterate
+    /// through all of the archetypes of the world with [`World::archetypes`] and extract
+    /// all possible combinations of component types which are currently stored in the `World`.
+    /// From there, you can then create a mapping of archetypes to wrapper objects for your
+    /// scripting language that provide functionality based off of the components of any
+    /// given `Entity`, and bind them onto an `Entity` when passed into your scripting
+    /// language by looking up the `Entity`'s archetype using `EntityRef::component_types`.
+    pub fn component_types(&self) -> impl Iterator<Item = TypeId> + '_ {
+        self.types.iter().map(|typeinfo| typeinfo.id)
+    }
+
     /// `index` must be in-bounds
     pub(crate) unsafe fn get_dynamic(
         &self,
diff --git a/src/bundle.rs b/src/bundle.rs
index 8d07eab..1868e2e 100644
--- a/src/bundle.rs
+++ b/src/bundle.rs
@@ -33,7 +33,7 @@ pub trait DynamicBundle {
     /// Must invoke `f` only with a valid pointer and the pointee's type and size. `put` may only be
     /// called at most once on any given value.
     #[doc(hidden)]
-    unsafe fn put(self, f: impl FnMut(*mut u8, TypeId, usize));
+    unsafe fn put(self, f: impl FnMut(*mut u8, TypeInfo));
 }
 
 /// A statically typed collection of components
@@ -52,9 +52,7 @@ pub trait Bundle: DynamicBundle {
     /// `f` must produce pointers to the expected fields. The implementation must not read from any
     /// pointers if any call to `f` returns `None`.
     #[doc(hidden)]
-    unsafe fn get(
-        f: impl FnMut(TypeId, usize) -> Option<NonNull<u8>>,
-    ) -> Result<Self, MissingComponent>
+    unsafe fn get(f: impl FnMut(TypeInfo) -> Option<NonNull<u8>>) -> Result<Self, MissingComponent>
     where
         Self: Sized;
 }
@@ -91,14 +89,13 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, unused_mut)]
-            unsafe fn put(self, mut f: impl FnMut(*mut u8, TypeId, usize)) {
+            unsafe fn put(self, mut f: impl FnMut(*mut u8, TypeInfo)) {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = self;
                 $(
                     f(
                         (&mut $name as *mut $name).cast::<u8>(),
-                        TypeId::of::<$name>(),
-                        mem::size_of::<$name>()
+                        TypeInfo::of::<$name>()
                     );
                     mem::forget($name);
                 )*
@@ -124,10 +121,10 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, unused_mut)]
-            unsafe fn get(mut f: impl FnMut(TypeId, usize) -> Option<NonNull<u8>>) -> Result<Self, MissingComponent> {
+            unsafe fn get(mut f: impl FnMut(TypeInfo) -> Option<NonNull<u8>>) -> Result<Self, MissingComponent> {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = ($(
-                    f(TypeId::of::<$name>(), mem::size_of::<$name>()).ok_or_else(MissingComponent::new::<$name>)?
+                    f(TypeInfo::of::<$name>()).ok_or_else(MissingComponent::new::<$name>)?
                         .as_ptr()
                         .cast::<$name>(),)*
                 );
diff --git a/src/entity_builder.rs b/src/entity_builder.rs
index ca21bd8..0c15f0f 100644
--- a/src/entity_builder.rs
+++ b/src/entity_builder.rs
@@ -19,7 +19,7 @@ use core::any::TypeId;
 use core::mem::{self, MaybeUninit};
 use core::ptr;
 
-use hashbrown::HashSet;
+use hashbrown::hash_map::{Entry, HashMap};
 
 use crate::archetype::TypeInfo;
 use crate::{Component, DynamicBundle};
@@ -42,7 +42,7 @@ pub struct EntityBuilder {
     cursor: usize,
     info: Vec<(TypeInfo, usize)>,
     ids: Vec<TypeId>,
-    id_set: HashSet<TypeId>,
+    indices: HashMap<TypeId, usize>,
 }
 
 impl EntityBuilder {
@@ -53,38 +53,104 @@ impl EntityBuilder {
             storage: Box::new([]),
             info: Vec::new(),
             ids: Vec::new(),
-            id_set: HashSet::new(),
+            indices: HashMap::new(),
         }
     }
 
-    /// Add `component` to the entity
+    /// Add `component` to the entity.
+    ///
+    /// If the bundle already contains a component of type `T`, it will
+    /// be dropped and replaced with the most recently added one.
     pub fn add<T: Component>(&mut self, component: T) -> &mut Self {
-        if !self.id_set.insert(TypeId::of::<T>()) {
-            return self;
-        }
-        let end = self.cursor + mem::size_of::<T>();
-        if end > self.storage.len() {
-            self.grow(end);
-        }
-        if mem::size_of::<T>() != 0 {
-            unsafe {
-                self.storage
-                    .as_mut_ptr()
-                    .add(self.cursor)
-                    .cast::<T>()
-                    .write_unaligned(component);
+        match self.indices.entry(TypeId::of::<T>()) {
+            Entry::Occupied(occupied) => {
+                let index = *occupied.get();
+                let (_, offset) = self.info[index];
+                unsafe {
+                    let storage_ptr = self
+                        .storage
+                        .as_mut_ptr()
+                        .cast::<u8>()
+                        .add(offset)
+                        .cast::<T>();
+
+                    // Drop the old value.
+                    let _ = storage_ptr.read_unaligned();
+                    // Overwrite the old value with our new one.
+                    storage_ptr.write_unaligned(component);
+                }
+                self
             }
+            Entry::Vacant(vacant) => {
+                let end = self.cursor + mem::size_of::<T>();
+                if end > self.storage.len() {
+                    Self::grow(end, self.cursor, &mut self.storage);
+                }
+                unsafe {
+                    self.storage
+                        .as_mut_ptr()
+                        .add(self.cursor)
+                        .cast::<T>()
+                        .write_unaligned(component);
+                }
+                vacant.insert(self.info.len());
+                self.info.push((TypeInfo::of::<T>(), self.cursor));
+                self.cursor += mem::size_of::<T>();
+                self
+            }
+        }
+    }
+
+    /// Add all components in `bundle` to the entity.
+    ///
+    /// If the bundle contains any component which matches the type of a component
+    /// already in the `EntityBuilder`, the newly added component from the bundle
+    /// will replace the old component and the old component will be dropped.
+    pub fn add_bundle(&mut self, bundle: impl DynamicBundle) -> &mut Self {
+        unsafe {
+            bundle.put(|ptr, ty| {
+                match self.indices.entry(ty.id()) {
+                    Entry::Occupied(occupied) => {
+                        let index = *occupied.get();
+                        let (ty, offset) = self.info[index];
+
+                        let storage_ptr = self.storage.as_mut_ptr().cast::<u8>().add(offset);
+                        // alloc a properly aligned tmp buffer and copy in the old value
+                        // so we can drop it safely
+                        let tmp = alloc(ty.layout());
+                        ptr::copy_nonoverlapping(storage_ptr, tmp, ty.layout().size());
+                        ty.drop(tmp);
+                        dealloc(tmp, ty.layout());
+                        // Overwrite the old value with our new one.
+                        ptr::copy_nonoverlapping(ptr, storage_ptr, ty.layout().size());
+                    }
+                    Entry::Vacant(vacant) => {
+                        let end = self.cursor + ty.layout().size();
+                        if end > self.storage.len() {
+                            Self::grow(end, self.cursor, &mut self.storage);
+                        }
+
+                        ptr::copy_nonoverlapping(
+                            ptr,
+                            self.storage.as_mut_ptr().add(self.cursor).cast(),
+                            ty.layout().size(),
+                        );
+
+                        vacant.insert(self.info.len());
+                        self.info.push((ty, self.cursor));
+                        self.cursor += ty.layout().size();
+                    }
+                }
+            });
         }
-        self.info.push((TypeInfo::of::<T>(), self.cursor));
-        self.cursor += mem::size_of::<T>();
         self
     }
 
-    fn grow(&mut self, min_size: usize) {
+    fn grow(min_size: usize, cursor: usize, storage: &mut Box<[MaybeUninit<u8>]>) {
         let new_len = min_size.next_power_of_two().max(64);
         let mut new_storage = vec![MaybeUninit::uninit(); new_len].into_boxed_slice();
-        new_storage[..self.cursor].copy_from_slice(&self.storage[..self.cursor]);
-        self.storage = new_storage;
+        new_storage[..cursor].copy_from_slice(&storage[..cursor]);
+        *storage = new_storage;
     }
 
     /// Construct a `Bundle` suitable for spawning
@@ -100,7 +166,7 @@ impl EntityBuilder {
     /// be called.
     pub fn clear(&mut self) {
         self.ids.clear();
-        self.id_set.clear();
+        self.indices.clear();
         self.cursor = 0;
         let max_size = self
             .info
@@ -169,10 +235,10 @@ impl DynamicBundle for BuiltEntity<'_> {
         self.builder.info.iter().map(|x| x.0).collect()
     }
 
-    unsafe fn put(self, mut f: impl FnMut(*mut u8, TypeId, usize)) {
+    unsafe fn put(self, mut f: impl FnMut(*mut u8, TypeInfo)) {
         for (ty, offset) in self.builder.info.drain(..) {
             let ptr = self.builder.storage.as_mut_ptr().add(offset).cast();
-            f(ptr, ty.id(), ty.layout().size());
+            f(ptr, ty);
         }
     }
 }
diff --git a/src/world.rs b/src/world.rs
index 6196c07..0d120e9 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -96,8 +96,8 @@ impl World {
         let archetype = &mut self.archetypes[archetype_id as usize];
         unsafe {
             let index = archetype.allocate(entity.id);
-            components.put(|ptr, ty, size| {
-                archetype.put_dynamic(ptr, ty, size, index);
+            components.put(|ptr, ty| {
+                archetype.put_dynamic(ptr, ty.id(), ty.layout().size(), index);
             });
             self.entities.meta[entity.id as usize].location = Location {
                 archetype: archetype_id,
@@ -379,8 +379,8 @@ impl World {
             if target == loc.archetype {
                 // Update components in the current archetype
                 let arch = &mut self.archetypes[loc.archetype as usize];
-                components.put(|ptr, ty, size| {
-                    arch.put_dynamic(ptr, ty, size, loc.index);
+                components.put(|ptr, ty| {
+                    arch.put_dynamic(ptr, ty.id(), ty.layout().size(), loc.index);
                 });
                 return Ok(());
             }
@@ -399,8 +399,8 @@ impl World {
             }) {
                 self.entities.meta[moved as usize].location.index = old_index;
             }
-            components.put(|ptr, ty, size| {
-                target_arch.put_dynamic(ptr, ty, size, target_index);
+            components.put(|ptr, ty| {
+                target_arch.put_dynamic(ptr, ty.id(), ty.layout().size(), target_index);
             });
         }
         Ok(())
@@ -462,7 +462,8 @@ impl World {
             };
             let old_index = loc.index;
             let source_arch = &self.archetypes[loc.archetype as usize];
-            let bundle = T::get(|ty, size| source_arch.get_dynamic(ty, size, old_index))?;
+            let bundle =
+                T::get(|ty| source_arch.get_dynamic(ty.id(), ty.layout().size(), old_index))?;
             let (source_arch, target_arch) = index2(
                 &mut self.archetypes,
                 loc.archetype as usize,
@@ -758,8 +759,9 @@ where
         let entity = self.entities.alloc();
         unsafe {
             let index = self.archetype.allocate(entity.id);
-            components.put(|ptr, ty, size| {
-                self.archetype.put_dynamic(ptr, ty, size, index);
+            components.put(|ptr, ty| {
+                self.archetype
+                    .put_dynamic(ptr, ty.id(), ty.layout().size(), index);
             });
             self.entities.meta[entity.id as usize].location = Location {
                 archetype: self.archetype_id,


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 22b62a9..2660155 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -123,11 +123,27 @@ fn build_entity() {
     entity.add("def");
     entity.add([0u8; 1024]);
     entity.add(456);
+    entity.add(789);
     let f = world.spawn(entity.build());
     assert_eq!(*world.get::<&str>(e).unwrap(), "abc");
     assert_eq!(*world.get::<i32>(e).unwrap(), 123);
     assert_eq!(*world.get::<&str>(f).unwrap(), "def");
-    assert_eq!(*world.get::<i32>(f).unwrap(), 456);
+    assert_eq!(*world.get::<i32>(f).unwrap(), 789);
+}
+
+#[test]
+fn build_entity_bundle() {
+    let mut world = World::new();
+    let mut entity = EntityBuilder::new();
+    entity.add_bundle(("abc", 123));
+    let e = world.spawn(entity.build());
+    entity.add(456);
+    entity.add_bundle(("def", [0u8; 1024], 789));
+    let f = world.spawn(entity.build());
+    assert_eq!(*world.get::<&str>(e).unwrap(), "abc");
+    assert_eq!(*world.get::<i32>(e).unwrap(), 123);
+    assert_eq!(*world.get::<&str>(f).unwrap(), "def");
+    assert_eq!(*world.get::<i32>(f).unwrap(), 789);
 }
 
 #[test]