From 17c8786eb00153240b27203f282276bd2fe3dcc6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20Mockers?= <mockersf@gmail.com>
Date: Sat, 31 Oct 2020 03:00:27 +0100
Subject: [PATCH] only get location once in normal path

---
 crates/bevy_ecs/hecs/src/world.rs      | 34 +++++++++++++++-----------
 crates/bevy_ecs/src/system/commands.rs |  2 +-
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/crates/bevy_ecs/hecs/src/world.rs b/crates/bevy_ecs/hecs/src/world.rs
index a684d854605fa..da0a0895f5fcd 100644
--- a/crates/bevy_ecs/hecs/src/world.rs
+++ b/crates/bevy_ecs/hecs/src/world.rs
@@ -672,25 +672,20 @@ impl World {
     /// ```
     pub fn remove<T: Bundle>(&mut self, entity: Entity) -> Result<T, ComponentError> {
         self.flush();
-        let loc = self.entities.get_mut(entity)?;
-        unsafe {
-            let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
+        let to_remove = T::with_static_ids(|ids| ids.iter().copied().collect::<HashSet<_>>());
 
-            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))?;
-            match self.remove_bundle_internal(entity, to_remove) {
-                Ok(_) => Ok(bundle),
-                Err(err) => Err(err),
-            }
+        match self.remove_bundle_internal::<_, T>(entity, to_remove, true) {
+            Ok(bundle) => Ok(bundle),
+            Err(err) => Err(err),
         }
     }
 
-    fn remove_bundle_internal<S: core::hash::BuildHasher>(
+    fn remove_bundle_internal<S: core::hash::BuildHasher, T: Bundle>(
         &mut self,
         entity: Entity,
         to_remove: std::collections::HashSet<TypeId, S>,
-    ) -> Result<(), ComponentError> {
+        check_presence: bool,
+    ) -> Result<T, ComponentError> {
         use std::collections::hash_map::Entry;
 
         let loc = self.entities.get_mut(entity)?;
@@ -718,6 +713,13 @@ impl World {
                 loc.archetype as usize,
                 target as usize,
             );
+            let bundle = if check_presence {
+                Some(T::get(|ty, size| {
+                    source_arch.get_dynamic(ty, size, old_index)
+                })?)
+            } else {
+                None
+            };
             let target_index = target_arch.allocate(entity);
             loc.archetype = target;
             loc.index = target_index;
@@ -739,7 +741,7 @@ impl World {
             {
                 self.entities.get_mut(moved).unwrap().index = old_index;
             }
-            Ok(())
+            bundle.ok_or_else(|| ComponentError::MissingComponent(MissingComponent::new::<()>()))
         }
     }
 
@@ -761,7 +763,11 @@ impl World {
             if self.archetypes[loc.archetype as usize].has_dynamic(component_to_remove) {
                 let mut single_component_hashset = std::collections::HashSet::new();
                 single_component_hashset.insert(component_to_remove);
-                match self.remove_bundle_internal(entity, single_component_hashset) {
+                match self.remove_bundle_internal::<_, ((),)>(
+                    entity,
+                    single_component_hashset,
+                    false,
+                ) {
                     Ok(_) | Err(ComponentError::MissingComponent(_)) => (),
                     Err(err) => return Err(err),
                 };
diff --git a/crates/bevy_ecs/src/system/commands.rs b/crates/bevy_ecs/src/system/commands.rs
index ab07559cb6dc8..8e834883272e9 100644
--- a/crates/bevy_ecs/src/system/commands.rs
+++ b/crates/bevy_ecs/src/system/commands.rs
@@ -377,7 +377,7 @@ mod tests {
             .map(|(a, b)| (*a, *b))
             .collect::<Vec<_>>();
         assert_eq!(results_after, vec![]);
-        let results_after_u64 = world.query::<&u64>().map(|a| *a).collect::<Vec<_>>();
+        let results_after_u64 = world.query::<&u64>().copied().collect::<Vec<_>>();
         assert_eq!(results_after_u64, vec![]);
     }
 }