Cleaner bundle trait interfaces

macros/src/lib.rs

@@ -71,12 +71,8 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 info
             }
 
-            unsafe fn put(mut self, mut f: impl FnMut(*mut u8, std::any::TypeId, usize) -> bool) {
-                #(
-                    if f((&mut self.#fields as *mut #tys).cast::<u8>(), std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>()) {
-                        std::mem::forget(self.#fields);
-                    }
-                )*
+            fn put(mut self, dest: &mut impl ::hecs::ComponentSink) {
+                #(dest.put(self.#fields);)*
             }
         }
 
@@ -107,18 +103,14 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
                 f(&*ELEMENTS)
             }
 
-            unsafe fn get(
-                mut f: impl FnMut(std::any::TypeId, usize) -> Option<std::ptr::NonNull<u8>>,
+            fn get(
+                src: &mut impl ::hecs::ComponentSource
             ) -> Result<Self, MissingComponent> {
-                Ok(Self {
-                    #(
-                        #fields: f(std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>())
-                            .ok_or_else(MissingComponent::new::<#tys>)?
-                            .cast::<#tys>()
-                            .as_ptr()
-                            .read(),
-                    )*
-                })
+                unsafe {
+                    Ok(Self {
+                        #(#fields: src.get()?,)*
+                    })
+                }
             }
         }
     };

src/archetype.rs

@@ -15,12 +15,12 @@
 use std::alloc::{alloc, Layout};
 use std::any::TypeId;
 use std::cell::UnsafeCell;
-use std::mem::MaybeUninit;
+use std::mem::{self, MaybeUninit};
 use std::ptr::{self, NonNull};
 
 use fxhash::FxHashMap;
 
-use crate::Component;
+use crate::{Component, ComponentSink, ComponentSource, MissingComponent};
 
 /// A collection of entities having the same component types
 pub struct Archetype {
@@ -113,14 +113,14 @@ impl Archetype {
         let mut data_size = 0;
         let mut offsets = FxHashMap::default();
         for ty in &self.types {
-            data_size = align(data_size, ty.layout.align());
+            data_size = align(data_size, ty.layout().align());
             offsets.insert(ty.id, data_size);
             data_size += ty.layout.size() * count;
         }
         let alloc = alloc(
             Layout::from_size_align(
                 data_size,
-                self.types.first().map_or(1, |x| x.layout.align()),
+                self.types.first().map_or(1, |x| x.layout().align()),
             )
             .unwrap(),
         )
@@ -212,6 +212,62 @@ impl Archetype {
             .cast::<u8>();
         ptr::copy_nonoverlapping(component, ptr, size);
     }
+
+    pub(crate) unsafe fn source(&mut self, index: u32) -> impl ComponentSource + '_ {
+        struct Source<'a> {
+            parent: &'a mut Archetype,
+            index: u32,
+        }
+
+        impl<'a> ComponentSource for Source<'a> {
+            unsafe fn get<T: Component>(&mut self) -> Result<T, MissingComponent> {
+                Ok(self
+                    .parent
+                    .get::<T>(self.index)
+                    .ok_or_else(MissingComponent::new::<T>)?
+                    .as_ptr()
+                    .read())
+            }
+        }
+
+        Source {
+            parent: self,
+            index,
+        }
+    }
+
+    pub(crate) unsafe fn sink(&mut self, index: u32) -> impl ComponentSink + '_ {
+        struct Sink<'a> {
+            parent: &'a mut Archetype,
+            index: u32,
+        }
+
+        impl<'a> ComponentSink for Sink<'a> {
+            fn put<T: Component>(&mut self, x: T) {
+                unsafe {
+                    self.parent
+                        .get::<T>(self.index)
+                        .expect("no such component")
+                        .as_ptr()
+                        .write(x)
+                }
+            }
+
+            unsafe fn put_dynamic(&mut self, component: &mut dyn Component) {
+                self.parent.put_dynamic(
+                    component as *mut dyn Component as *mut u8,
+                    (*component).type_id(),
+                    mem::size_of_val(component),
+                    self.index,
+                );
+            }
+        }
+
+        Sink {
+            parent: self,
+            index,
+        }
+    }
 }
 
 impl Drop for Archetype {

src/bundle.rs

@@ -13,8 +13,7 @@
 // limitations under the License.
 
 use std::any::{type_name, TypeId};
-use std::ptr::NonNull;
-use std::{fmt, mem};
+use std::fmt;
 
 use crate::archetype::TypeInfo;
 use crate::Component;
@@ -25,23 +24,16 @@ pub trait DynamicBundle {
     fn with_ids<T>(&self, f: impl FnOnce(&[TypeId]) -> T) -> T;
     #[doc(hidden)]
     fn type_info(&self) -> Vec<TypeInfo>;
-    /// Allow a callback to move all components out of the bundle
-    ///
-    /// Must invoke `f` only with a valid pointer, its type, and the pointee's size. A `false`
-    /// return value indicates that the value was not moved and should be dropped.
     #[doc(hidden)]
-    unsafe fn put(self, f: impl FnMut(*mut u8, TypeId, usize) -> bool);
+    fn put(self, dest: &mut impl ComponentSink);
 }
 
 /// A statically typed collection of components
 pub trait Bundle: DynamicBundle {
     #[doc(hidden)]
     fn with_static_ids<T>(f: impl FnOnce(&[TypeId]) -> T) -> T;
-    /// Construct `Self` by moving components out of pointers fetched by `f`
     #[doc(hidden)]
-    unsafe fn get(
-        f: impl FnMut(TypeId, usize) -> Option<NonNull<u8>>,
-    ) -> Result<Self, MissingComponent>
+    fn get(src: &mut impl ComponentSource) -> Result<Self, MissingComponent>
     where
         Self: Sized;
 }
@@ -65,6 +57,37 @@ impl fmt::Display for MissingComponent {
 
 impl std::error::Error for MissingComponent {}
 
+/// A receiver of component data to which a bundle may be written
+pub trait ComponentSink {
+    /// Write a `T` component
+    ///
+    /// Must be called at most once per distinct `T`.
+    fn put<T: Component>(&mut self, mut x: T)
+    where
+        Self: Sized,
+    {
+        unsafe {
+            self.put_dynamic(&mut x);
+        }
+        std::mem::forget(x);
+    }
+
+    /// Write an component of any type
+    ///
+    /// # Safety
+    /// `component` will be moved out of, and *must* be forgotten immediately after calling.
+    unsafe fn put_dynamic(&mut self, component: &mut dyn Component);
+}
+
+/// A source of component data from which a bundle may be aggregated
+pub trait ComponentSource {
+    /// Obtain a `T` component
+    ///
+    /// # Safety
+    /// Must be called at most once per distinct `T`
+    unsafe fn get<T: Component>(&mut self) -> Result<T, MissingComponent>;
+}
+
 macro_rules! tuple_impl {
     ($($name: ident),*) => {
         impl<$($name: Component),*> DynamicBundle for ($($name,)*) {
@@ -79,18 +102,10 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, unused_mut)]
-            unsafe fn put(self, mut f: impl FnMut(*mut u8, TypeId, usize) -> bool) {
+            fn put(self, dest: &mut impl ComponentSink) {
                 #[allow(non_snake_case)]
                 let ($(mut $name,)*) = self;
-                $(
-                    if f(
-                        (&mut $name as *mut $name).cast::<u8>(),
-                        TypeId::of::<$name>(),
-                        mem::size_of::<$name>()
-                    ) {
-                        mem::forget($name)
-                    }
-                )*
+                $(dest.put($name);)*
             }
         }
 
@@ -107,14 +122,11 @@ macro_rules! tuple_impl {
             }
 
             #[allow(unused_variables, unused_mut)]
-            unsafe fn get(mut f: impl FnMut(TypeId, usize) -> Option<NonNull<u8>>) -> Result<Self, MissingComponent> {
-                Ok(($(
-                    f(TypeId::of::<$name>(), mem::size_of::<$name>())
-                      .ok_or_else(MissingComponent::new::<$name>)?
-                      .as_ptr()
-                      .cast::<$name>()
-                      .read(),
-                )*))
+            fn get(src: &mut impl ComponentSource) -> Result<Self, MissingComponent> {
+                #[allow(unused_unsafe)]
+                unsafe {
+                    Ok(($(src.get::<$name>()?,)*))
+                }
             }
         }
     }

src/entity_builder.rs

@@ -17,7 +17,7 @@ use std::any::TypeId;
 use std::mem::{self, MaybeUninit};
 
 use crate::archetype::TypeInfo;
-use crate::{Component, DynamicBundle};
+use crate::{Component, ComponentSink, DynamicBundle};
 
 /// Helper for incrementally constructing an entity with dynamic component types
 ///
@@ -36,7 +36,11 @@ pub struct EntityBuilder {
     storage: Box<[MaybeUninit<u8>]>,
     // Backwards from the end!
     cursor: *mut u8,
-    info: Vec<(TypeInfo, *mut u8)>,
+    info: Vec<(
+        TypeInfo,
+        *mut u8,
+        unsafe fn(*mut u8, &mut dyn ComponentSink),
+    )>,
     ids: Vec<TypeId>,
 }
 
@@ -68,7 +72,8 @@ impl EntityBuilder {
             }
             self.cursor.cast::<T>().write(component);
         }
-        self.info.push((TypeInfo::of::<T>(), self.cursor));
+        self.info
+            .push((TypeInfo::of::<T>(), self.cursor, put_helper::<T>));
         self
     }
 
@@ -107,7 +112,7 @@ impl EntityBuilder {
     pub fn clear(&mut self) {
         self.ids.clear();
         unsafe {
-            for (ty, component) in self.info.drain(..) {
+            for (ty, component, _) in self.info.drain(..) {
                 ty.drop(component);
             }
             self.cursor = self.storage.as_mut_ptr().add(self.storage.len()).cast();
@@ -130,6 +135,10 @@ pub struct BuiltEntity<'a> {
     builder: &'a mut EntityBuilder,
 }
 
+unsafe fn put_helper<T: Component>(ptr: *mut u8, sink: &mut dyn ComponentSink) {
+    sink.put_dynamic(&mut *ptr.cast::<T>());
+}
+
 impl DynamicBundle for BuiltEntity<'_> {
     fn with_ids<T>(&self, f: impl FnOnce(&[TypeId]) -> T) -> T {
         f(&self.builder.ids)
@@ -140,10 +149,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) -> bool) {
-        for (ty, component) in self.builder.info.drain(..) {
-            if !f(component, ty.id(), ty.layout().size()) {
-                ty.drop(component);
+    fn put(self, dst: &mut impl ComponentSink) {
+        for (_, component, put) in self.builder.info.drain(..) {
+            unsafe {
+                put(component, dst);
             }
         }
     }

src/lib.rs

@@ -46,7 +46,7 @@ mod query;
 mod world;
 
 pub use borrow::{EntityRef, Ref, RefMut};
-pub use bundle::{Bundle, DynamicBundle, MissingComponent};
+pub use bundle::{Bundle, DynamicBundle, MissingComponent, ComponentSink, ComponentSource};
 pub use entity_builder::{BuiltEntity, EntityBuilder};
 pub use query::{Query, QueryIter};
 pub use world::{Component, ComponentError, Entity, Iter, NoSuchEntity, World};

src/world.rs

@@ -76,10 +76,7 @@ impl World {
         unsafe {
             let index = archetype.allocate(entity.id);
             self.entities.meta[entity.id as usize].location.index = index;
-            components.put(|ptr, ty, size| {
-                archetype.put_dynamic(ptr, ty, size, index);
-                true
-            });
+            components.put(&mut archetype.sink(index));
         }
         entity
     }
@@ -233,10 +230,7 @@ impl World {
             };
             if target == loc.archetype {
                 let arch = &mut self.archetypes[loc.archetype as usize];
-                components.put(|ptr, ty, size| {
-                    arch.put_dynamic(ptr, ty, size, loc.index);
-                    true
-                });
+                components.put(&mut arch.sink(loc.index));
                 return Ok(());
             }
 
@@ -249,10 +243,7 @@ impl World {
             source_arch.move_to(loc.index, |ptr, ty, size| {
                 target_arch.put_dynamic(ptr, ty, size, target_index);
             });
-            components.put(|ptr, ty, size| {
-                target_arch.put_dynamic(ptr, ty, size, target_index);
-                true
-            });
+            components.put(&mut target_arch.sink(target_index));
             loc.archetype = target;
             loc.index = target_index;
         }
@@ -308,7 +299,7 @@ impl World {
                 target as usize,
             );
             let target_index = target_arch.allocate(entity.id);
-            let x = T::get(|ty, size| source_arch.get_dynamic(ty, size, loc.index))?;
+            let x = T::get(&mut source_arch.source(loc.index))?;
             source_arch.move_to(loc.index, |src, ty, size| {
                 // Only move the components present in the target archetype, i.e. the non-removed ones.
                 if let Some(dst) = target_arch.get_dynamic(ty, size, target_index) {
@@ -393,7 +384,12 @@ impl fmt::Display for NoSuchEntity {
 impl Error for NoSuchEntity {}
 
 /// Types that can be components (implemented automatically)
-pub trait Component: Send + Sync + 'static {}
+pub trait Component: Send + Sync + 'static {
+    /// Look up the type of the component
+    fn type_id(&self) -> TypeId {
+        TypeId::of::<Self>()
+    }
+}
 impl<T: Send + Sync + 'static> Component for T {}
 
 /// Lightweight unique ID of an entity