Use type erasure to DRY component registration

examples/cloning.rs

@@ -1,89 +1,125 @@
-use hecs::{ColumnBatchBuilder, ColumnBatchType, Component, World};
-
-fn maybe_add_type_to_batch<T: Component>(
-    archetype: &hecs::Archetype,
-    batch_type: &mut ColumnBatchType,
-) {
-    if archetype.has::<T>() {
-        batch_type.add::<T>();
-    }
-}
+use std::any::TypeId;
 
-fn maybe_clone_column<T: Component + Clone>(
-    archetype: &hecs::Archetype,
-    batch: &mut ColumnBatchBuilder,
-) {
-    if let Some((column, mut writer)) = archetype.get::<&T>().zip(batch.writer::<T>()) {
-        for item in column.iter() {
-            if let Err(_) = writer.push(item.clone()) {
-                unreachable!("push should always succeed since batch was sized to match archetype");
-            }
-        }
-    }
+use hecs::{Archetype, ColumnBatchBuilder, ColumnBatchType, Component, TypeIdMap, TypeInfo, World};
+
+struct ComponentCloneMetadata {
+    type_info: TypeInfo,
+    func: &'static dyn Fn(&Archetype, &mut ColumnBatchBuilder),
 }
 
-/// Clones world entities along with a hardcoded list of components.
+/// Clones world entities along with registered components when [Self::clone_world()] is called.
 ///
-/// Components not included are omitted from the cloned world.
+/// Unregistered components are omitted from the cloned world. Entities containing unregistered
+/// components will still be cloned.
 ///
 /// Note that entity allocator state may differ in the cloned world - so for example a new entity
 /// spawned in each world may end up with different entity ids, entity iteration order may be
 /// different, etc.
-fn clone_world(world: &World) -> World {
-    let mut cloned = World::new();
-
-    for archetype in world.archetypes() {
-        let mut batch = ColumnBatchType::new();
-        // types have to be listed one by one here
-        maybe_add_type_to_batch::<String>(archetype, &mut batch);
-        maybe_add_type_to_batch::<i32>(archetype, &mut batch);
-
-        let mut batch = batch.into_batch(archetype.ids().len() as u32);
-        // and types need to be listed again here
-        maybe_clone_column::<String>(archetype, &mut batch);
-        maybe_clone_column::<i32>(archetype, &mut batch);
-
-        let batch = batch.build().expect("batch should be complete");
-
-        let handles = &cloned
-            .reserve_entities(archetype.ids().len() as u32)
-            .collect::<Vec<_>>();
-        cloned.flush();
-        cloned.spawn_column_batch_at(handles, batch);
+#[derive(Default)]
+struct WorldCloner {
+    registry: TypeIdMap<ComponentCloneMetadata>,
+}
+
+impl WorldCloner {
+    pub fn register<T: Component + Clone>(&mut self) {
+        self.registry.insert(TypeId::of::<T>(), ComponentCloneMetadata {
+            type_info: TypeInfo::of::<T>(),
+         func : &|src, dest| {
+            let mut column = dest.writer::<T>().unwrap();
+            for component in &*src.get::<&T>().unwrap() {
+                _ = column.push(component.clone());
+            }
+        }});
     }
 
-    cloned
+    fn clone_world(&self, world: &World) -> World {
+        let mut cloned = World::new();
+
+        for archetype in world.archetypes() {
+            let mut batch = ColumnBatchType::new();
+
+            for (&k, v) in self.registry.iter() {
+                if archetype.has_dynamic(k) {
+                    batch.add_dynamic(v.type_info);
+                }
+            }
+            let mut batch = batch.into_batch(archetype.ids().len() as u32);
+
+            for (&k, v) in self.registry.iter() {
+                if archetype.has_dynamic(k) {
+                (v.func)(archetype, &mut batch)
+                }
+            }
+
+            let batch = batch.build().expect("batch should be complete");
+
+            let handles = &cloned
+                .reserve_entities(archetype.ids().len() as u32)
+                .collect::<Vec<_>>();
+            cloned.flush();
+            cloned.spawn_column_batch_at(handles, batch);
+        }
+
+        cloned
+    }
 }
 
+
 pub fn main() {
     let int0 = 0;
     let int1 = 1;
     let str0 = "Ada".to_owned();
     let str1 = "Bob".to_owned();
+    let str2 = "Cal".to_owned();
 
     let mut world0 = World::new();
     let entity0 = world0.spawn((int0, str0));
     let entity1 = world0.spawn((int1, str1));
+    let entity2 = world0.spawn((str2,));
+    let entity3 = world0.spawn((0u8,)); // unregistered component
+
+    let mut cloner = WorldCloner::default();
+    cloner.register::<i32>();
+    cloner.register::<String>();
 
-    let world1 = clone_world(&world0);
+    let world1 = cloner.clone_world(&world0);
 
     assert_eq!(
         world0.len(),
         world1.len(),
         "cloned world should have same entity count as original world"
     );
 
-    type AllComponentsQuery = (&'static i32, &'static String);
+    // NB: unregistered components don't get cloned
+    assert!(world0.entity(entity3).expect("w0 entity3 should exist").has::<u8>(),
+            "original world entity has u8 component");
+    assert!(!world1.entity(entity3).expect("w1 entity3 should exist").has::<u8>(),
+            "cloned world entity does not have u8 component because it was not registered");
 
+    type AllRegisteredComponentsQuery = (&'static i32, &'static String);
     for entity in [entity0, entity1] {
         let w0_e = world0.entity(entity).expect("w0 entity should exist");
         let w1_e = world1.entity(entity).expect("w1 entity should exist");
-        assert!(w0_e.satisfies::<AllComponentsQuery>());
-        assert!(w1_e.satisfies::<AllComponentsQuery>());
+        assert!(w0_e.satisfies::<AllRegisteredComponentsQuery>());
+        assert!(w1_e.satisfies::<AllRegisteredComponentsQuery>());
+
+        assert_eq!(
+            w0_e.query::<AllRegisteredComponentsQuery>().get().unwrap(),
+            w1_e.query::<AllRegisteredComponentsQuery>().get().unwrap()
+        );
+    }
+
+    type SomeRegisteredComponentsQuery = (&'static String,);
+    for entity in [entity2] {
+        let w0_e = world0.entity(entity).expect("w0 entity should exist");
+        let w1_e = world1.entity(entity).expect("w1 entity should exist");
+        assert!(w0_e.satisfies::<SomeRegisteredComponentsQuery>());
+        assert!(w1_e.satisfies::<SomeRegisteredComponentsQuery>());
 
         assert_eq!(
-            w0_e.query::<AllComponentsQuery>().get().unwrap(),
-            w1_e.query::<AllComponentsQuery>().get().unwrap()
+            w0_e.query::<SomeRegisteredComponentsQuery>().get().unwrap(),
+            w1_e.query::<SomeRegisteredComponentsQuery>().get().unwrap()
         );
     }
+
 }

src/archetype.rs

@@ -454,7 +454,7 @@ struct Data {
 /// TypeId is already thoroughly hashed, so there's no reason to hash it again.
 /// Just leave the bits unchanged.
 #[derive(Default)]
-pub(crate) struct TypeIdHasher {
+pub struct TypeIdHasher {
     hash: u64,
 }
 
@@ -491,7 +491,7 @@ impl Hasher for TypeIdHasher {
 /// Because TypeId is already a fully-hashed u64 (including data in the high seven bits,
 /// which hashbrown needs), there is no need to hash it again. Instead, this uses the much
 /// faster no-op hash.
-pub(crate) type TypeIdMap<V> = HashMap<TypeId, V, BuildHasherDefault<TypeIdHasher>>;
+pub type TypeIdMap<V> = HashMap<TypeId, V, BuildHasherDefault<TypeIdHasher>>;
 
 struct OrderedTypeIdMap<V>(Box<[(TypeId, V)]>);
 

src/batch.rs

@@ -24,6 +24,12 @@ impl ColumnBatchType {
         self
     }
 
+    /// [Self::add()] but using type information determined at runtime via [TypeInfo::of()]
+    pub fn add_dynamic(&mut self, id: TypeInfo) -> &mut Self {
+        self.types.push(id);
+        self
+    }
+
     /// Construct a [`ColumnBatchBuilder`] for *exactly* `size` entities with these components
     pub fn into_batch(self, size: u32) -> ColumnBatchBuilder {
         let mut types = self.types.into_sorted_vec();

src/lib.rs

@@ -84,7 +84,7 @@ pub mod serialize;
 mod take;
 mod world;
 
-pub use archetype::{Archetype, ArchetypeColumn, ArchetypeColumnMut};
+pub use archetype::{Archetype, ArchetypeColumn, ArchetypeColumnMut, TypeIdMap, TypeInfo};
 pub use batch::{BatchIncomplete, BatchWriter, ColumnBatch, ColumnBatchBuilder, ColumnBatchType};
 pub use bundle::{
     bundle_satisfies_query, dynamic_bundle_satisfies_query, Bundle, DynamicBundle,
@@ -109,8 +109,6 @@ pub use world::{
 
 // Unstable implementation details needed by the macros
 #[doc(hidden)]
-pub use archetype::TypeInfo;
-#[doc(hidden)]
 pub use bundle::DynamicClone;
 #[doc(hidden)]
 pub use query::Fetch;