convert other get_components to return a result

Author: hymm

benches/benches/bevy_ecs/world/mod.rs

@@ -35,5 +35,6 @@ criterion_group!(
     query_get_many::<2>,
     query_get_many::<5>,
     query_get_many::<10>,
+    query_get_components_mut,
     entity_set_build_and_lookup,
 );

benches/benches/bevy_ecs/world/world_get.rs

@@ -1,11 +1,12 @@
+use benches::bench;
 use core::hint::black_box;
 
 use bevy_ecs::{
     bundle::{Bundle, NoBundleEffect},
     component::Component,
     entity::Entity,
     system::{Query, SystemState},
-    world::World,
+    world::{EntityMut, World},
 };
 use criterion::Criterion;
 use rand::{prelude::SliceRandom, SeedableRng};
@@ -357,3 +358,74 @@ pub fn query_get_many<const N: usize>(criterion: &mut Criterion) {
         });
     }
 }
+
+pub fn query_get_components_mut(criterion: &mut Criterion) {
+    let mut group = criterion.benchmark_group(bench!("world_query_get_components_mut"));
+    group.warm_up_time(core::time::Duration::from_millis(500));
+    group.measurement_time(core::time::Duration::from_secs(4));
+
+    for entity_count in RANGE.map(|i| i * 10_000) {
+        let (mut world, entities) = setup_wide::<(
+            WideTable<0>,
+            WideTable<1>,
+            WideTable<2>,
+            WideTable<3>,
+            WideTable<4>,
+            WideTable<5>,
+            WideTable<6>,
+            WideTable<7>,
+            WideTable<8>,
+            WideTable<9>,
+        )>(entity_count);
+        let mut query = world.query::<EntityMut>();
+        group.bench_function(format!("{entity_count}_entities"), |bencher| {
+            bencher.iter(|| {
+                for entity in &entities {
+                    assert!(query
+                        .get_mut(&mut world, *entity)
+                        .unwrap()
+                        .get_components_mut::<(
+                            &mut WideTable<0>,
+                            &mut WideTable<1>,
+                            &mut WideTable<2>,
+                            &mut WideTable<3>,
+                            &mut WideTable<4>,
+                            &mut WideTable<5>,
+                            &mut WideTable<6>,
+                            &mut WideTable<7>,
+                            &mut WideTable<8>,
+                            &mut WideTable<9>,
+                        )>()
+                        .is_ok());
+                }
+            });
+        });
+        group.bench_function(format!("unchecked_{entity_count}_entities"), |bencher| {
+            bencher.iter(|| {
+                for entity in &entities {
+                    // SAFETY: no duplicate components are listed
+                    unsafe {
+                        assert!(query
+                            .get_mut(&mut world, *entity)
+                            .unwrap()
+                            .get_components_mut_unchecked::<(
+                                &mut WideTable<0>,
+                                &mut WideTable<1>,
+                                &mut WideTable<2>,
+                                &mut WideTable<3>,
+                                &mut WideTable<4>,
+                                &mut WideTable<5>,
+                                &mut WideTable<6>,
+                                &mut WideTable<7>,
+                                &mut WideTable<8>,
+                                &mut WideTable<9>,
+                            )>()
+                            .is_ok());
+                    }
+                }
+            });
+        });
+    }
+
+    group.finish();
+}

crates/bevy_ecs/src/event/trigger.rs

@@ -301,7 +301,7 @@ unsafe impl<
                 return;
             }
             if let Ok(entity) = world.get_entity(current_entity)
-                && let Some(item) = entity.get_components::<T>()
+                && let Ok(item) = entity.get_components::<T>()
                 && let Some(traverse_to) = T::traverse(item, event)
             {
                 current_entity = traverse_to;

crates/bevy_ecs/src/query/access_iter.rs

@@ -4,7 +4,7 @@ use crate::{
 };
 
 /// The data storage type that is being accessed.
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum EcsAccessType {
     /// Accesses [`Component`](crate::prelude::Component) data
     Component(EcsAccessLevel),
@@ -14,7 +14,7 @@ pub enum EcsAccessType {
 
 /// The way the data will be accessed and whether we take access on all the components on
 /// an entity or just one component.
-#[derive(Clone, Copy, Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum EcsAccessLevel {
     /// Reads [`Component`](crate::prelude::Component) with [`ComponentId`]
     Read(ComponentId),
@@ -47,7 +47,7 @@ pub enum EcsAccessLevel {
 }
 
 /// Access level needed by [`QueryData`] fetch to the resource.
-#[derive(Copy, Clone, Debug)]
+#[derive(Copy, Clone, Debug, PartialEq)]
 pub enum ResourceAccessLevel {
     /// Reads the resource with [`ComponentId`]
     Read(ComponentId),
@@ -207,37 +207,40 @@ impl EcsAccessType {
 }
 
 /// Error returned from [`has_conflicts`].
-#[derive(Clone, Copy, Debug)]
-pub enum AccessError {
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum QueryAccessError {
+    /// Component was not registered on world
     ComponentNotRegistered,
+    /// The [`EcsAccessType`]'s conflict with each other
     Conflict(EcsAccessType, EcsAccessType),
+    /// Entity did not have the requested components
     EntityDoesNotMatch,
 }
 
-impl core::error::Error for AccessError {}
+impl core::error::Error for QueryAccessError {}
 
-impl core::fmt::Display for AccessError {
+impl core::fmt::Display for QueryAccessError {
     fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
         match *self {
-            AccessError::ComponentNotRegistered => {
+            QueryAccessError::ComponentNotRegistered => {
                 write!(
                     f,
                     "At least one component in Q was not registered in world. 
                     Consider calling `World::register_component`"
                 )
             }
-            AccessError::Conflict(ecs_access, ecs_access1) => {
+            QueryAccessError::Conflict(ecs_access, ecs_access1) => {
                 write!(f, "Conflict between {ecs_access:?} and {ecs_access1:?}")
             }
-            AccessError::EntityDoesNotMatch => {
+            QueryAccessError::EntityDoesNotMatch => {
                 write!(f, "Entity does not match Q")
             }
         }
     }
 }
 
 /// Check if `Q` has any internal conflicts.
-pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), AccessError> {
+pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), QueryAccessError> {
     let mut index_outer = 0;
     for (i, access) in Q::iter_access(components, &mut index_outer).enumerate() {
         let mut index_inner = 0;
@@ -251,7 +254,7 @@ pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), Access
                 continue;
             }
             let (Some(access), Some(access_other)) = (access, access_other) else {
-                return Err(AccessError::ComponentNotRegistered);
+                return Err(QueryAccessError::ComponentNotRegistered);
             };
 
             // if we're in an except sequence, check if the sequence has ended
@@ -264,7 +267,7 @@ pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), Access
 
                 if sequence_ended {
                     if !except_compatible {
-                        return Err(AccessError::Conflict(access, access_other));
+                        return Err(QueryAccessError::Conflict(access, access_other));
                     }
                     except_compatible = false;
                     except_index = None;
@@ -279,7 +282,7 @@ pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), Access
                     except_index = Some(index);
                     except_compatible = true;
                 },
-                AccessCompatible::Conflicts => return Err(AccessError::Conflict(access, access_other)),
+                AccessCompatible::Conflicts => return Err(QueryAccessError::Conflict(access, access_other)),
                 AccessCompatible::ConflictsExceptSecond(index) => {
                     except_index = Some(index);
                 }
@@ -288,7 +291,7 @@ pub fn has_conflicts<Q: QueryData>(components: &Components) -> Result<(), Access
 
         if except_index.is_some() && !except_compatible {
             if let (Some(access), Some(access_other)) = (access, last_access) {
-                return Err(AccessError::Conflict(access, access_other));
+                return Err(QueryAccessError::Conflict(access, access_other));
             }
         }
     }
@@ -324,15 +327,15 @@ mod tests {
         // Conflicts
         assert!(matches!(
             has_conflicts::<(&C1, &mut C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C1, &C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C1, &mut C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
     }
 
@@ -351,35 +354,35 @@ mod tests {
         // Conflicts
         assert!(matches!(
             has_conflicts::<(EntityRef, &mut C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C1, EntityRef)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityMut, &C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&C1, EntityMut)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityMut, &mut C1)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C1, EntityMut)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityMut, EntityRef)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityRef, EntityMut)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
     }
 
@@ -403,11 +406,11 @@ mod tests {
         // Conflicts
         assert!(matches!(
             has_conflicts::<(EntityRefExcept<C1>, &mut C2)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C2, EntityRefExcept<C1>)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
     }
 
@@ -430,19 +433,19 @@ mod tests {
         // Conflicts
         assert!(matches!(
             has_conflicts::<(&C2, EntityMutExcept<C1>)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityMutExcept<C1>, &C2)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(EntityMutExcept<C1>, &mut C2)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
         assert!(matches!(
             has_conflicts::<(&mut C2, EntityMutExcept<C1>)>(c),
-            Err(AccessError::Conflict(_, _))
+            Err(QueryAccessError::Conflict(_, _))
         ));
     }
 }

crates/bevy_ecs/src/world/entity_access/entity_mut.rs

@@ -3,7 +3,7 @@ use crate::{
     change_detection::{ComponentTicks, MaybeLocation, Tick},
     component::{Component, ComponentId, Mutable},
     entity::{ContainsEntity, Entity, EntityEquivalent, EntityLocation},
-    query::{has_conflicts, Access, AccessError, ReadOnlyQueryData, ReleaseStateQueryData},
+    query::{has_conflicts, Access, QueryAccessError, ReadOnlyQueryData, ReleaseStateQueryData},
     world::{
         error::EntityComponentError, unsafe_world_cell::UnsafeEntityCell, DynamicComponentFetch,
         EntityRef, FilteredEntityMut, FilteredEntityRef, Mut, Ref,
@@ -152,7 +152,7 @@ impl<'w> EntityMut<'w> {
     /// or `None` if the entity does not have the components required by the query `Q`.
     pub fn get_components<Q: ReadOnlyQueryData + ReleaseStateQueryData>(
         &self,
-    ) -> Option<Q::Item<'_, 'static>> {
+    ) -> Result<Q::Item<'_, 'static>, QueryAccessError> {
         self.as_readonly().get_components::<Q>()
     }
 
@@ -186,18 +186,40 @@ impl<'w> EntityMut<'w> {
     /// the `QueryData` does not provide aliasing mutable references to the same component.
     pub unsafe fn get_components_mut_unchecked<Q: ReleaseStateQueryData>(
         &mut self,
-    ) -> Option<Q::Item<'_, 'static>> {
+    ) -> Result<Q::Item<'_, 'static>, QueryAccessError> {
         // SAFETY: Caller ensures the `QueryData` does not provide aliasing mutable references to the same component
         unsafe { self.reborrow().into_components_mut_unchecked::<Q>() }
     }
 
-    /// returns None if component wasn't registered, or if the access is not compatible bewteen terms
+    /// Returns components for the current entity that match the query `Q`.
+    /// In the case of conflicting [`QueryData`](crate::query::QueryData), unregistered components, or missing components,
+    /// this will return a [`QueryAccessError`]
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::prelude::*;
+    /// #
+    /// #[derive(Component)]
+    /// struct X(usize);
+    /// #[derive(Component)]
+    /// struct Y(usize);
+    ///
+    /// # let mut world = World::default();
+    /// let mut entity = world.spawn((X(0), Y(0))).into_mutable();
+    /// // Get mutable access to two components at once
+    /// // SAFETY: X and Y are different components
+    /// let (mut x, mut y) = entity.get_components_mut::<(&mut X, &mut Y)>().unwrap();
+    /// ```
+    ///
+    /// Note that this does a O(n^2) check that the [`QueryData`](crate::query::QueryData) does not conflict. If performance is a
+    /// consideration you should use [`Self::get_components_mut_unchecked`] instead.
     pub fn get_components_mut<Q: ReleaseStateQueryData>(
         &mut self,
-    ) -> Result<Q::Item<'_, 'static>, AccessError> {
+    ) -> Result<Q::Item<'_, 'static>, QueryAccessError> {
         has_conflicts::<Q>(self.cell.world().components())?;
         // SAFETY: we checked that there were not conflicting components above
-        unsafe { self.get_components_mut_unchecked::<Q>() }.ok_or(AccessError::EntityDoesNotMatch)
+        unsafe { self.get_components_mut_unchecked::<Q>() }
     }
 
     /// Consumes self and returns components for the current entity that match the query `Q` for the world lifetime `'w`,
@@ -230,7 +252,7 @@ impl<'w> EntityMut<'w> {
     /// the `QueryData` does not provide aliasing mutable references to the same component.
     pub unsafe fn into_components_mut_unchecked<Q: ReleaseStateQueryData>(
         self,
-    ) -> Option<Q::Item<'w, 'static>> {
+    ) -> Result<Q::Item<'w, 'static>, QueryAccessError> {
         // SAFETY:
         // - We have mutable access to all components of this entity.
         // - Caller asserts the `QueryData` does not provide aliasing mutable references to the same component