Repository: Ralith/hecs
Pull request #47 (2020-01-25 03:20:53)

Batched spawn and storage preallocation


Patch:
diff --git a/benches/bench.rs b/benches/bench.rs
index 8ddea7c..685dbb0 100644
--- a/benches/bench.rs
+++ b/benches/bench.rs
@@ -41,6 +41,25 @@ fn spawn_static(b: &mut Bencher) {
     });
 }
 
+fn spawn_batch(b: &mut Bencher) {
+    #[derive(Bundle)]
+    struct Bundle {
+        pos: Position,
+        vel: Velocity,
+    };
+
+    let mut world = World::new();
+    b.iter(|| {
+        world
+            .spawn_batch((0..1_000).map(|_| Bundle {
+                pos: Position(0.0),
+                vel: Velocity(0.0),
+            }))
+            .for_each(|_| {});
+        world.clear();
+    });
+}
+
 fn iterate_100k(b: &mut Bencher) {
     let mut world = World::new();
     for i in 0..100_000 {
@@ -62,5 +81,12 @@ fn build(b: &mut Bencher) {
     });
 }
 
-benchmark_group!(benches, spawn_tuple, spawn_static, iterate_100k, build);
+benchmark_group!(
+    benches,
+    spawn_tuple,
+    spawn_static,
+    spawn_batch,
+    iterate_100k,
+    build
+);
 benchmark_main!(benches);
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index 48c7887..ae18870 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -50,9 +50,7 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             }
 
             fn type_info(&self) -> Vec<::hecs::TypeInfo> {
-                let mut info = vec![#(::hecs::TypeInfo::of::<#tys>()),*];
-                info.sort_unstable();
-                info
+                Self::static_type_info()
             }
 
             unsafe fn put(mut self, mut f: impl FnMut(*mut u8, std::any::TypeId, usize) -> bool) {
@@ -91,12 +89,18 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 f(&*ELEMENTS)
             }
 
+            fn static_type_info() -> Vec<::hecs::TypeInfo> {
+                let mut info = vec![#(::hecs::TypeInfo::of::<#tys>()),*];
+                info.sort_unstable();
+                info
+            }
+
             unsafe fn get(
                 mut f: impl FnMut(std::any::TypeId, usize) -> Option<std::ptr::NonNull<u8>>,
-            ) -> Result<Self, MissingComponent> {
+            ) -> Result<Self, ::hecs::MissingComponent> {
                 #(
                     let #fields = f(std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>())
-                            .ok_or_else(MissingComponent::new::<#tys>)?
+                            .ok_or_else(::hecs::MissingComponent::new::<#tys>)?
                             .cast::<#tys>()
                         .as_ptr();
                 )*
diff --git a/src/archetype.rs b/src/archetype.rs
index 0510ea5..adc41a8 100644
--- a/src/archetype.rs
+++ b/src/archetype.rs
@@ -117,8 +117,8 @@ impl Archetype {
         }
     }
 
-    pub(crate) fn len(&self) -> usize {
-        self.len as usize
+    pub(crate) fn len(&self) -> u32 {
+        self.len
     }
 
     pub(crate) fn entities(&self) -> NonNull<u32> {
@@ -151,55 +151,67 @@ impl Archetype {
 
     /// Every type must be written immediately after this call
     pub(crate) unsafe fn allocate(&mut self, id: u32) -> u32 {
-        if (self.len as usize) < self.entities.len() {
-            self.entities[self.len as usize] = id;
-            self.len += 1;
-            return self.len - 1;
+        if self.len as usize == self.entities.len() {
+            self.grow(self.len.max(64));
         }
 
-        // At this point we need to allocate more storage.
-        let old_count = self.entities.len();
-        let count = if old_count == 0 { 64 } else { old_count * 2 };
-        let mut new_entities = vec![!0; count].into_boxed_slice();
-        new_entities[0..old_count].copy_from_slice(&self.entities);
-        self.entities = new_entities;
+        self.entities[self.len as usize] = id;
+        self.len += 1;
+        self.len - 1
+    }
 
-        let old_data_size = mem::replace(&mut self.data_size, 0);
-        let mut state = HashMap::with_capacity(self.types.len());
-        for ty in &self.types {
-            self.data_size = align(self.data_size, ty.layout.align());
-            state.insert(ty.id, TypeState::new(self.data_size));
-            self.data_size += ty.layout.size() * count;
+    pub(crate) fn reserve(&mut self, additional: u32) {
+        if additional > (self.capacity() - self.len()) {
+            self.grow(additional - (self.capacity() - self.len()));
         }
-        let new_data = if self.data_size == 0 {
-            NonNull::dangling()
-        } else {
-            NonNull::new(alloc(
-                Layout::from_size_align(
-                    self.data_size,
-                    self.types.first().map_or(1, |x| x.layout.align()),
-                )
-                .unwrap(),
-            ))
-            .unwrap()
-        };
-        if old_data_size != 0 {
+    }
+
+    pub(crate) fn capacity(&self) -> u32 {
+        self.entities.len() as u32
+    }
+
+    fn grow(&mut self, increment: u32) {
+        unsafe {
+            let old_count = self.len as usize;
+            let count = old_count + increment as usize;
+            let mut new_entities = vec![!0; count].into_boxed_slice();
+            new_entities[0..old_count].copy_from_slice(&self.entities[0..old_count]);
+            self.entities = new_entities;
+
+            let old_data_size = mem::replace(&mut self.data_size, 0);
+            let mut state = HashMap::with_capacity(self.types.len());
             for ty in &self.types {
-                let old_off = self.state.get(&ty.id).unwrap().offset;
-                let new_off = state.get(&ty.id).unwrap().offset;
-                ptr::copy_nonoverlapping(
-                    (*self.data.get()).as_ptr().add(old_off),
-                    new_data.as_ptr().add(new_off),
-                    ty.layout.size() * old_count,
-                );
+                self.data_size = align(self.data_size, ty.layout.align());
+                state.insert(ty.id, TypeState::new(self.data_size));
+                self.data_size += ty.layout.size() * count;
+            }
+            let new_data = if self.data_size == 0 {
+                NonNull::dangling()
+            } else {
+                NonNull::new(alloc(
+                    Layout::from_size_align(
+                        self.data_size,
+                        self.types.first().map_or(1, |x| x.layout.align()),
+                    )
+                    .unwrap(),
+                ))
+                .unwrap()
+            };
+            if old_data_size != 0 {
+                for ty in &self.types {
+                    let old_off = self.state.get(&ty.id).unwrap().offset;
+                    let new_off = state.get(&ty.id).unwrap().offset;
+                    ptr::copy_nonoverlapping(
+                        (*self.data.get()).as_ptr().add(old_off),
+                        new_data.as_ptr().add(new_off),
+                        ty.layout.size() * old_count,
+                    );
+                }
             }
-        }
 
-        self.data = UnsafeCell::new(new_data);
-        self.state = state;
-        self.entities[self.len as usize] = id;
-        self.len += 1;
-        self.len - 1
+            self.data = UnsafeCell::new(new_data);
+            self.state = state;
+        }
     }
 
     /// Returns the ID of the entity moved into `index`, if any
diff --git a/src/bundle.rs b/src/bundle.rs
index dbce023..88c1bb4 100644
--- a/src/bundle.rs
+++ b/src/bundle.rs
@@ -40,6 +40,11 @@ pub trait DynamicBundle {
 pub trait Bundle: DynamicBundle {
     #[doc(hidden)]
     fn with_static_ids<T>(f: impl FnOnce(&[TypeId]) -> T) -> T;
+
+    /// Obtain the fields' TypeInfos, sorted by descending alignment then id
+    #[doc(hidden)]
+    fn static_type_info() -> Vec<TypeInfo>;
+
     /// Construct `Self` by moving components out of pointers fetched by `f`
     ///
     /// # Safety
@@ -82,9 +87,7 @@ macro_rules! tuple_impl {
             }
 
             fn type_info(&self) -> Vec<TypeInfo> {
-                let mut xs = vec![$(TypeInfo::of::<$name>()),*];
-                xs.sort_unstable();
-                xs
+                Self::static_type_info()
             }
 
             #[allow(unused_variables, unused_mut)]
@@ -115,6 +118,12 @@ macro_rules! tuple_impl {
                 f(&ids)
             }
 
+            fn static_type_info() -> Vec<TypeInfo> {
+                let mut xs = vec![$(TypeInfo::of::<$name>()),*];
+                xs.sort_unstable();
+                xs
+            }
+
             #[allow(unused_variables, unused_mut)]
             unsafe fn get(mut f: impl FnMut(TypeId, usize) -> Option<NonNull<u8>>) -> Result<Self, MissingComponent> {
                 #[allow(non_snake_case)]
diff --git a/src/entities.rs b/src/entities.rs
index e35d56e..57ce0a1 100644
--- a/src/entities.rs
+++ b/src/entities.rs
@@ -74,7 +74,7 @@ impl Entities {
     /// Storage for entity generation and location is lazily allocated by calling `flush`. Locations
     /// can be determined by the return value of `flush` and by iterating through the `reserved`
     /// accessors, and should all be written immediately after flushing.
-    pub fn reserve(&self) -> Entity {
+    pub fn reserve_entity(&self) -> Entity {
         loop {
             let index = self.free_cursor.load(Ordering::Relaxed);
             match index.checked_sub(1) {
@@ -172,6 +172,19 @@ impl Entities {
         Ok(loc)
     }
 
+    /// Ensure `n` at least allocations can succeed without reallocating
+    pub fn reserve(&mut self, additional: u32) {
+        debug_assert_eq!(
+            self.pending.load(Ordering::Relaxed),
+            0,
+            "allocator must be flushed before potentially growing"
+        );
+        let free = self.free_cursor.load(Ordering::Relaxed);
+        if additional > free {
+            self.grow(additional - free);
+        }
+    }
+
     pub fn contains(&self, entity: Entity) -> bool {
         if entity.id >= self.meta.len() as u32 {
             return true;
@@ -225,11 +238,11 @@ impl Entities {
 
     /// Allocate space for and enumerate pending entities
     pub fn flush(&mut self) -> impl Iterator<Item = u32> {
-        let pending = self.pending.swap(0, Ordering::Relaxed); // Not racey due to &mut self
+        let pending = self.pending.load(Ordering::Relaxed); // Not racey due to &mut self
 
         if pending != 0 {
             let first = self.meta.len() as u32;
-            self.grow(pending);
+            self.grow(0);
             first..(first + pending)
         } else {
             0..0
@@ -253,9 +266,10 @@ impl Entities {
     }
 
     /// Expand storage and mark all but the first `pending` of the new slots as free
-    fn grow(&mut self, pending: u32) {
-        let new_len = (self.meta.len() * 2)
-            .max(self.meta.len() + pending as usize)
+    fn grow(&mut self, increment: u32) {
+        let pending = self.pending.swap(0, Ordering::Relaxed);
+        let new_len = (self.meta.len() + pending as usize + increment as usize)
+            .max(self.meta.len() * 2)
             .max(1024);
         let mut new_meta = Vec::with_capacity(new_len);
         new_meta.extend_from_slice(&self.meta);
@@ -291,7 +305,6 @@ impl Entities {
             new_reserved.push(AtomicU32::new(x.load(Ordering::Relaxed)));
         }
         new_reserved.resize_with(new_len, || AtomicU32::new(0));
-        for _ in 0..new_len {}
         self.reserved = new_reserved.into();
     }
 }
diff --git a/src/query.rs b/src/query.rs
index 6d00cf2..09c6f95 100644
--- a/src/query.rs
+++ b/src/query.rs
@@ -371,7 +371,7 @@ impl<'q, 'w, Q: Query> IntoIterator for &'q mut QueryBorrow<'w, Q> {
 struct ChunkIter<Q: Query> {
     entities: NonNull<u32>,
     fetch: Q::Fetch,
-    len: usize,
+    len: u32,
 }
 
 /// Iterator over the set of entities with the components in `Q`
@@ -433,7 +433,7 @@ impl<'q, 'w, Q: Query> ExactSizeIterator for QueryIter<'q, 'w, Q> {
             .archetypes
             .iter()
             .filter(|&x| Q::Fetch::wants(x))
-            .map(|x| x.len())
+            .map(|x| x.len() as usize)
             .sum()
     }
 }
diff --git a/src/world.rs b/src/world.rs
index 8d8748b..d5db717 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -14,6 +14,7 @@
 
 use crate::alloc::vec::Vec;
 use core::any::TypeId;
+use core::convert::TryFrom;
 use core::{fmt, mem, ptr};
 
 #[cfg(feature = "std")]
@@ -101,6 +102,45 @@ impl World {
         entity
     }
 
+    /// Efficiently spawn a large number of entities with the same components
+    ///
+    /// Faster than calling `spawn` repeatedly with the same components. The resulting iterator
+    /// *must* be driven to completion for entities to actually be spawned.
+    ///
+    /// # Example
+    /// ```
+    /// # use hecs::*;
+    /// let mut world = World::new();
+    /// let entities = world.spawn_batch((0..1_000).map(|i| (i, "abc"))).collect::<Vec<_>>();
+    /// for i in 0..1_000 {
+    ///     assert_eq!(*world.get::<i32>(entities[i]).unwrap(), i as i32);
+    /// }
+    /// ```
+    pub fn spawn_batch<'a, T: Bundle>(
+        &'a mut self,
+        iter: impl IntoIterator<Item = T> + 'a,
+    ) -> impl Iterator<Item = Entity> + 'a {
+        let iter = iter.into_iter();
+        let (lower, upper) = iter.size_hint();
+        let archetype_id = self
+            .reserve_inner::<T>(u32::try_from(upper.unwrap_or(lower)).expect("iterator too large"));
+
+        iter.map(move |components| unsafe {
+            let archetype = &mut self.archetypes[archetype_id as usize];
+            let entity = self.entities.alloc();
+            let index = archetype.allocate(entity.id);
+            components.put(|ptr, ty, size| {
+                archetype.put_dynamic(ptr, ty, size, index);
+                true
+            });
+            self.entities.meta[entity.id as usize].location = Location {
+                archetype: archetype_id,
+                index,
+            };
+            entity
+        })
+    }
+
     /// Allocate an entity ID concurrently
     ///
     /// Unlike `spawn`, this can be called simultaneously to other operations on the `World` such as
@@ -111,8 +151,8 @@ impl World {
     /// explicitly by calling `flush`.
     ///
     /// Useful for reserving an ID that will later have components attached to it with `insert`.
-    pub fn reserve(&self) -> Entity {
-        self.entities.reserve()
+    pub fn reserve_entity(&self) -> Entity {
+        self.entities.reserve_entity()
     }
 
     /// Destroy an entity and all its components
@@ -125,6 +165,28 @@ impl World {
         Ok(())
     }
 
+    /// Ensure `additional` entities with exact components `T` can be spawned without reallocating
+    pub fn reserve<T: Bundle>(&mut self, additional: u32) {
+        self.reserve_inner::<T>(additional);
+    }
+
+    fn reserve_inner<T: Bundle>(&mut self, additional: u32) -> u32 {
+        self.flush();
+        self.entities.reserve(additional);
+
+        let archetype_id = T::with_static_ids(|ids| {
+            self.index.get(ids).copied().unwrap_or_else(|| {
+                let x = self.archetypes.len() as u32;
+                self.archetypes.push(Archetype::new(T::static_type_info()));
+                self.index.insert(ids.to_vec(), x);
+                x
+            })
+        });
+
+        self.archetypes[archetype_id as usize].reserve(additional);
+        archetype_id
+    }
+
     /// Despawn all entities
     ///
     /// Preserves allocated storage for reuse.


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 4b19bef..833a7d3 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -280,8 +280,8 @@ fn remove_missing() {
 #[test]
 fn reserve() {
     let mut world = World::new();
-    let a = world.reserve();
-    let b = world.reserve();
+    let a = world.reserve_entity();
+    let b = world.reserve_entity();
 
     assert_eq!(world.query::<()>().iter().count(), 0);