Repository: Ralith/hecs
Pull request #85 (2020-10-09 17:29:48)

Add `World::resolve_unknown_gen`
This adds an unsafe method `World::resolve_unknown_gen` which allows the user to map a `u32` identifier extracted via `Entity::id` back to a potentially live `Entity`. As discussed previously it's unsafe because it may return a reference to a dead `Entity`, which will appear as though it has no components and may later spontaneously acquire components by becoming a live `Entity`.


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/entities.rs b/src/entities.rs
index a9fbd09..8e80344 100644
--- a/src/entities.rs
+++ b/src/entities.rs
@@ -236,6 +236,25 @@ impl Entities {
         Ok(meta.location)
     }
 
+    /// Returns `None` if the given id would represent an index outside of `meta`.
+    ///
+    /// # Safety
+    /// Must only be called for currently allocated `id`s.
+    pub unsafe fn resolve_unknown_gen(&self, id: u32) -> Entity {
+        let meta_len = self.meta.len();
+        if meta_len + self.pending.load(Ordering::Relaxed) as usize <= id as usize {
+            panic!("entity id is out of range");
+        } else if meta_len <= id as usize {
+            Entity { generation: 0, id }
+        } else {
+            let meta = &self.meta[id as usize];
+            Entity {
+                generation: meta.generation,
+                id,
+            }
+        }
+    }
+
     /// Allocate space for and enumerate pending entities
     #[allow(clippy::reversed_empty_ranges)]
     pub fn flush(&mut self) -> impl Iterator<Item = u32> {
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..00ca3c3 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,
@@ -305,6 +305,14 @@ impl World {
         })
     }
 
+    /// Given an id obtained from `Entity::id`, reconstruct the still-live `Entity`.
+    ///
+    /// # Safety
+    /// `id` must correspond to a currently live `Entity`. A despawned or never-allocated `id` will produce undefined behavior.
+    pub unsafe fn find_entity_from_id(&self, id: u32) -> Entity {
+        self.entities.resolve_unknown_gen(id)
+    }
+
     /// Iterate over all entities in the world
     ///
     /// Entities are yielded in arbitrary order. Prefer `World::query` for better performance when
@@ -379,8 +387,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 +407,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 +470,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 +767,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]