Repository: Ralith/hecs
Pull request #25 (2019-12-29 00:08:59)

Fix panic on attempting to remove missing component


Patch:
diff --git a/macros/src/lib.rs b/macros/src/lib.rs
index 5b109bf..48c7887 100644
--- a/macros/src/lib.rs
+++ b/macros/src/lib.rs
@@ -94,15 +94,13 @@ pub fn derive_bundle(input: TokenStream) -> TokenStream {
             unsafe fn get(
                 mut f: impl FnMut(std::any::TypeId, usize) -> Option<std::ptr::NonNull<u8>>,
             ) -> Result<Self, MissingComponent> {
-                Ok(Self {
-                    #(
-                        #fields: f(std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>())
+                #(
+                    let #fields = f(std::any::TypeId::of::<#tys>(), std::mem::size_of::<#tys>())
                             .ok_or_else(MissingComponent::new::<#tys>)?
                             .cast::<#tys>()
-                            .as_ptr()
-                            .read(),
-                    )*
-                })
+                        .as_ptr();
+                )*
+                Ok(Self { #( #fields: #fields.read(), )* })
             }
         }
     };
diff --git a/src/archetype.rs b/src/archetype.rs
index a8433d4..0510ea5 100644
--- a/src/archetype.rs
+++ b/src/archetype.rs
@@ -69,7 +69,11 @@ impl Archetype {
     }
 
     pub(crate) fn has<T: Component>(&self) -> bool {
-        self.state.contains_key(&TypeId::of::<T>())
+        self.has_dynamic(TypeId::of::<T>())
+    }
+
+    pub(crate) fn has_dynamic(&self, id: TypeId) -> bool {
+        self.state.contains_key(&id)
     }
 
     pub(crate) fn get<T: Component>(&self) -> Option<NonNull<T>> {
diff --git a/src/bundle.rs b/src/bundle.rs
index cabfb9c..e546f7f 100644
--- a/src/bundle.rs
+++ b/src/bundle.rs
@@ -39,6 +39,11 @@ 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`
+    ///
+    /// # Safety
+    ///
+    /// `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>>,
@@ -110,13 +115,13 @@ 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(),
-                )*))
+                #[allow(non_snake_case)]
+                let ($(mut $name,)*) = ($(
+                    f(TypeId::of::<$name>(), mem::size_of::<$name>()).ok_or_else(MissingComponent::new::<$name>)?
+                        .as_ptr()
+                        .cast::<$name>(),)*
+                );
+                Ok(($($name.read(),)*))
             }
         }
     }
diff --git a/src/world.rs b/src/world.rs
index 5903dc3..279b926 100644
--- a/src/world.rs
+++ b/src/world.rs
@@ -339,6 +339,9 @@ impl World {
                     index
                 }
             };
+            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 (source_arch, target_arch) = index2(
                 &mut self.archetypes,
                 loc.archetype as usize,
@@ -346,8 +349,7 @@ impl World {
             );
             let target_index = target_arch.allocate(entity.id);
             loc.archetype = target;
-            let old_index = mem::replace(&mut loc.index, target_index);
-            let x = T::get(|ty, size| source_arch.get_dynamic(ty, size, old_index))?;
+            loc.index = target_index;
             if let Some(moved) = source_arch.move_to(old_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) {
@@ -356,7 +358,7 @@ impl World {
             }) {
                 self.entities.meta[moved as usize].location.index = old_index;
             }
-            Ok(x)
+            Ok(bundle)
         }
     }
 


Test patch:
diff --git a/tests/tests.rs b/tests/tests.rs
index 634a251..51ee246 100644
--- a/tests/tests.rs
+++ b/tests/tests.rs
@@ -267,3 +267,10 @@ fn alias() {
     let _a = q.iter().collect::<Vec<_>>();
     let _b = q.iter().collect::<Vec<_>>();
 }
+
+#[test]
+fn remove_missing() {
+    let mut world = World::new();
+    let e = world.spawn(("abc", 123));
+    assert!(world.remove_one::<bool>(e).is_err());
+}