Add and for less branching in PackedOption

Author: NathanSWard

crates/bevy_ecs/src/archetype.rs

@@ -55,6 +55,7 @@ pub struct AddBundle {
 impl NoneValue for ArchetypeId {
     const NONE_VALUE: Self = ArchetypeId::invalid();
 
+    #[inline]
     fn is_none_value(&self) -> bool {
         *self == Self::NONE_VALUE
     }
@@ -66,6 +67,7 @@ impl NoneValue for AddBundle {
         bundle_status: Vec::new(),
     };
 
+    #[inline]
     fn is_none_value(&self) -> bool {
         self.archetype_id.is_none_value()
     }
@@ -154,6 +156,7 @@ impl NoneValue for ArchetypeComponentInfo {
         archetype_component_id: ArchetypeComponentId::NONE_VALUE,
     };
 
+    #[inline]
     fn is_none_value(&self) -> bool {
         self.archetype_component_id.is_none_value()
     }
@@ -407,6 +410,7 @@ impl SparseSetIndex for ArchetypeComponentId {
 impl NoneValue for ArchetypeComponentId {
     const NONE_VALUE: Self = Self::invalid();
 
+    #[inline]
     fn is_none_value(&self) -> bool {
         *self == Self::invalid()
     }

crates/bevy_ecs/src/storage/sparse_set.rs

@@ -200,7 +200,7 @@ impl ComponentSparseSet {
     /// it exists). It is the caller's responsibility to drop the returned ptr (if Some is
     /// returned).
     pub fn remove_and_forget(&mut self, entity: Entity) -> Option<*mut u8> {
-        self.sparse.remove(entity).expand().map(|dense_index| {
+        self.sparse.remove(entity).map_option(|dense_index| {
             // SAFE: unique access to ticks
             unsafe {
                 (*self.ticks.get()).swap_remove(dense_index);
@@ -218,7 +218,11 @@ impl ComponentSparseSet {
     }
 
     pub fn remove(&mut self, entity: Entity) -> bool {
-        if let Some(dense_index) = self.sparse.remove(entity).expand() {
+        let dense_index = self.sparse.remove(entity);
+
+        if dense_index.is_some() {
+            // SAFE: `dense_index` is a some value.
+            let dense_index = unsafe { dense_index.unwrap_unchecked() };
             self.ticks.get_mut().swap_remove(dense_index);
             self.entities.swap_remove(dense_index);
             let is_last = dense_index == self.dense.len() - 1;
@@ -360,7 +364,7 @@ impl<I: SparseSetIndex, V> SparseSet<I, V> {
     }
 
     pub fn remove(&mut self, index: I) -> Option<V> {
-        self.sparse.remove(index).expand().map(|dense_index| {
+        self.sparse.remove(index).map_option(|dense_index| {
             let is_last = dense_index == self.dense.len() - 1;
             let value = self.dense.swap_remove(dense_index);
             self.indices.swap_remove(dense_index);

crates/bevy_ecs/src/world/entity_ref.rs

@@ -303,15 +303,24 @@ impl<'w> EntityMut<'w> {
         let bundle_info = self.world.bundles.init_info::<T>(components);
         let old_location = self.location;
         let new_archetype_id = unsafe {
-            remove_bundle_from_archetype(
+            let archetype_id = remove_bundle_from_archetype(
                 archetypes,
                 storages,
                 components,
                 old_location.archetype_id,
                 bundle_info,
                 false,
-            )
-            .expand()?
+            );
+
+            if archetype_id.is_some() {
+                #[allow(unused_unsafe)]
+                // SAFE: `archetype_id` is some.
+                unsafe {
+                    archetype_id.unwrap_unchecked()
+                }
+            } else {
+                return None;
+            }
         };
 
         if new_archetype_id == old_location.archetype_id {

crates/bevy_utils/src/packed_option.rs

@@ -9,6 +9,7 @@ macro_rules! impl_max_none_value {
         $(impl NoneValue for $ty {
             const NONE_VALUE: Self = <$ty>::MAX;
 
+            #[inline]
             fn is_none_value(&self) -> bool {
                 *self == Self::NONE_VALUE
             }
@@ -48,6 +49,21 @@ impl<T: NoneValue> PackedOption<T> {
         self.expand().unwrap()
     }
 
+    /// Returns the contained value, consuming `self`,
+    /// without checking that the value is not `[None]`.
+    ///
+    /// # Safety
+    /// Calling this method on [`None`] is *[undefined behavior]*.
+    #[inline]
+    #[track_caller]
+    pub unsafe fn unwrap_unchecked(self) -> T {
+        debug_assert!(
+            self.is_some(),
+            "Trying to call `unwrap_unchecked` on a `PackedOption` in a `None` state."
+        );
+        self.0
+    }
+
     /// Returns the contained `Some` value, consuming the `self` value.
     /// Panics with `msg` if self is `None`.
     #[inline]
@@ -88,6 +104,16 @@ impl<T: NoneValue> PackedOption<T> {
         }
     }
 
+    /// Maps a `PackedOption<T>` to `Option<U>` by applying a function to the contained value.
+    #[inline]
+    pub fn map_option<U, F: FnOnce(T) -> U>(self, f: F) -> Option<U> {
+        if self.is_some() {
+            Some(f(self.0))
+        } else {
+            None
+        }
+    }
+
     /// Takes the value out of the packed option, leaveing `None` in its place.
     pub fn take(&mut self) -> PackedOption<T> {
         std::mem::take(self)