Moved get_component(_unchecked_mut) from Query to QueryState

crates/bevy_ecs/src/query/state.rs

@@ -1,8 +1,9 @@
 use crate::{
     archetype::{Archetype, ArchetypeComponentId, ArchetypeGeneration, ArchetypeId},
+    change_detection::Mut,
     component::{ComponentId, Tick},
     entity::Entity,
-    prelude::FromWorld,
+    prelude::{Component, FromWorld},
     query::{
         Access, BatchingStrategy, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter,
         QueryIter, QueryParIter, WorldQuery,
@@ -13,7 +14,7 @@ use crate::{
 #[cfg(feature = "trace")]
 use bevy_utils::tracing::Instrument;
 use fixedbitset::FixedBitSet;
-use std::{borrow::Borrow, fmt, mem::MaybeUninit};
+use std::{any::TypeId, borrow::Borrow, fmt, mem::MaybeUninit};
 
 use super::{NopWorldQuery, QueryManyIter, ROQueryItem, ReadOnlyWorldQuery};
 
@@ -506,6 +507,142 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         }
     }
 
+    /// Returns a shared reference to the component `T` of the given [`Entity`].
+    ///
+    /// In case of a nonexisting entity or mismatched component, a [`QueryEntityError`] is returned instead.
+    #[inline]
+    pub(crate) fn get_component<'w, 's, 'r, T: Component>(
+        &'s self,
+        world: UnsafeWorldCell<'w>,
+        entity: Entity,
+    ) -> Result<&'r T, QueryComponentError>
+    where
+        'w: 'r,
+    {
+        let entity_ref = world
+            .get_entity(entity)
+            .ok_or(QueryComponentError::NoSuchEntity)?;
+        let component_id = world
+            .components()
+            .get_id(TypeId::of::<T>())
+            .ok_or(QueryComponentError::MissingComponent)?;
+        let archetype_component = entity_ref
+            .archetype()
+            .get_archetype_component_id(component_id)
+            .ok_or(QueryComponentError::MissingComponent)?;
+        if self
+            .archetype_component_access
+            .has_read(archetype_component)
+        {
+            // SAFETY: `world` must have access to the component `T` for this entity,
+            // since it was registered in `self`'s archetype component access set.
+            unsafe { entity_ref.get::<T>() }.ok_or(QueryComponentError::MissingComponent)
+        } else {
+            Err(QueryComponentError::MissingReadAccess)
+        }
+    }
+
+    /// Returns a shared reference to the component `T` of the given [`Entity`].
+    ///
+    /// # Panics
+    ///
+    /// If given a nonexisting entity or mismatched component, this will panic.
+    #[inline]
+    pub(crate) fn component<'w, 's, 'r, T: Component>(
+        &'s self,
+        world: UnsafeWorldCell<'w>,
+        entity: Entity,
+    ) -> &'r T
+    where
+        'w: 'r,
+    {
+        match self.get_component(world, entity) {
+            Ok(component) => component,
+            Err(error) => {
+                panic!(
+                    "Cannot get component `{:?}` from {entity:?}: {error}",
+                    TypeId::of::<T>()
+                )
+            }
+        }
+    }
+
+    /// Returns a mutable reference to the component `T` of the given entity.
+    ///
+    /// In case of a nonexisting entity or mismatched component, a [`QueryComponentError`] is returned instead.
+    ///
+    /// # Safety
+    ///
+    /// This function makes it possible to violate Rust's aliasing guarantees.
+    /// You must make sure this call does not result in multiple mutable references to the same component.
+    #[inline]
+    pub unsafe fn get_component_unchecked_mut<'w, 's, 'r, T: Component>(
+        &'s self,
+        world: UnsafeWorldCell<'w>,
+        entity: Entity,
+        last_run: Tick,
+        this_run: Tick,
+    ) -> Result<Mut<'r, T>, QueryComponentError>
+    where
+        'w: 'r,
+    {
+        let entity_ref = world
+            .get_entity(entity)
+            .ok_or(QueryComponentError::NoSuchEntity)?;
+        let component_id = world
+            .components()
+            .get_id(TypeId::of::<T>())
+            .ok_or(QueryComponentError::MissingComponent)?;
+        let archetype_component = entity_ref
+            .archetype()
+            .get_archetype_component_id(component_id)
+            .ok_or(QueryComponentError::MissingComponent)?;
+        if self
+            .archetype_component_access
+            .has_write(archetype_component)
+        {
+            // SEMI-SAFETY: It is the responsibility of the caller to ensure it is sound to get a
+            // mutable reference to this entity's component `T`.
+            let result = unsafe { entity_ref.get_mut_using_ticks::<T>(last_run, this_run) };
+
+            result.ok_or(QueryComponentError::MissingComponent)
+        } else {
+            Err(QueryComponentError::MissingWriteAccess)
+        }
+    }
+
+    /// Returns a mutable reference to the component `T` of the given entity.
+    ///
+    /// # Panics
+    ///
+    /// Panics in case of a nonexisting entity or mismatched component.
+    ///
+    /// # Safety
+    ///
+    /// This function makes it possible to violate Rust's aliasing guarantees.
+    /// You must make sure this call does not result in multiple mutable references to the same component.
+    #[inline]
+    pub(crate) unsafe fn component_unchecked_mut<'w, 's, 'r, T: Component>(
+        &'s self,
+        world: UnsafeWorldCell<'w>,
+        entity: Entity,
+        last_run: Tick,
+        this_run: Tick,
+    ) -> Mut<'r, T>
+    where
+        'w: 'r,
+    {
+        match self.get_component_unchecked_mut(world, entity, last_run, this_run) {
+            Ok(component) => component,
+            Err(error) => {
+                panic!(
+                    "Cannot get component `{:?}` from {entity:?}: {error}",
+                    TypeId::of::<T>()
+                )
+            }
+        }
+    }
+
     /// Gets the read-only query results for the given [`World`] and array of [`Entity`], where the last change and
     /// the current change tick are given.
     ///
@@ -538,6 +675,28 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         Ok(values.map(|x| x.assume_init()))
     }
 
+    /// Gets the read-only query results for the given [`World`] and array of [`Entity`], where the last change and
+    /// the current change tick are given.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any item cannot be retrieved.
+    ///
+    /// # Safety
+    ///
+    /// This must be called on the same `World` that the `Query` was generated from:
+    /// use `QueryState::validate_world` to verify this.
+    pub(crate) unsafe fn many_read_only_manual<'w, const N: usize>(
+        &self,
+        world: &'w World,
+        entities: [Entity; N],
+        last_run: Tick,
+        this_run: Tick,
+    ) -> [ROQueryItem<'w, Q>; N] {
+        self.get_many_read_only_manual(world, entities, last_run, this_run)
+            .unwrap()
+    }
+
     /// Gets the query results for the given [`World`] and array of [`Entity`], where the last change and
     /// the current change tick are given.
     ///
@@ -575,6 +734,31 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
         Ok(values.map(|x| x.assume_init()))
     }
 
+    /// Gets the query results for the given [`World`] and array of [`Entity`], where the last change and
+    /// the current change tick are given.
+    ///
+    /// # Panics
+    ///
+    /// Panics if any item cannot be retrieved.
+    ///
+    /// # Safety
+    ///
+    /// This does not check for unique access to subsets of the entity-component data.
+    /// To be safe, make sure mutable queries have unique access to the components they query.
+    ///
+    /// This must be called on the same `World` that the `Query` was generated from:
+    /// use `QueryState::validate_world` to verify this.
+    pub(crate) unsafe fn many_unchecked_manual<'w, const N: usize>(
+        &self,
+        world: UnsafeWorldCell<'w>,
+        entities: [Entity; N],
+        last_run: Tick,
+        this_run: Tick,
+    ) -> [Q::Item<'w>; N] {
+        self.get_many_unchecked_manual(world, entities, last_run, this_run)
+            .unwrap()
+    }
+
     /// Returns an [`Iterator`] over the query results for the given [`World`].
     ///
     /// This can only be called for read-only queries, see [`Self::iter_mut`] for write-queries.
@@ -1343,6 +1527,99 @@ impl fmt::Display for QueryEntityError {
     }
 }
 
+/// An error that occurs when retrieving a specific [`Entity`]'s component from a [`Query`](crate::system::Query).
+#[derive(Debug, PartialEq, Eq)]
+pub enum QueryComponentError {
+    /// The [`Query`](crate::system::Query) does not have read access to the requested component.
+    ///
+    /// This error occurs when the requested component is not included in the original query.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::{prelude::*, query::QueryComponentError};
+    /// #
+    /// # #[derive(Component)]
+    /// # struct OtherComponent;
+    /// #
+    /// # #[derive(Component, PartialEq, Debug)]
+    /// # struct RequestedComponent;
+    /// #
+    /// # #[derive(Resource)]
+    /// # struct SpecificEntity {
+    /// #     entity: Entity,
+    /// # }
+    /// #
+    /// fn get_missing_read_access_error(query: Query<&OtherComponent>, res: Res<SpecificEntity>) {
+    ///     assert_eq!(
+    ///         query.get_component::<RequestedComponent>(res.entity),
+    ///         Err(QueryComponentError::MissingReadAccess),
+    ///     );
+    ///     println!("query doesn't have read access to RequestedComponent because it does not appear in Query<&OtherComponent>");
+    /// }
+    /// # bevy_ecs::system::assert_is_system(get_missing_read_access_error);
+    /// ```
+    MissingReadAccess,
+    /// The [`Query`](crate::system::Query) does not have write access to the requested component.
+    ///
+    /// This error occurs when the requested component is not included in the original query, or the mutability of the requested component is mismatched with the original query.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # use bevy_ecs::{prelude::*, query::QueryComponentError};
+    /// #
+    /// # #[derive(Component, PartialEq, Debug)]
+    /// # struct RequestedComponent;
+    /// #
+    /// # #[derive(Resource)]
+    /// # struct SpecificEntity {
+    /// #     entity: Entity,
+    /// # }
+    /// #
+    /// fn get_missing_write_access_error(mut query: Query<&RequestedComponent>, res: Res<SpecificEntity>) {
+    ///     assert_eq!(
+    ///         query.get_component::<RequestedComponent>(res.entity),
+    ///         Err(QueryComponentError::MissingWriteAccess),
+    ///     );
+    ///     println!("query doesn't have write access to RequestedComponent because it doesn't have &mut in Query<&RequestedComponent>");
+    /// }
+    /// # bevy_ecs::system::assert_is_system(get_missing_write_access_error);
+    /// ```
+    MissingWriteAccess,
+    /// The given [`Entity`] does not have the requested component.
+    MissingComponent,
+    /// The requested [`Entity`] does not exist.
+    NoSuchEntity,
+}
+
+impl std::error::Error for QueryComponentError {}
+
+impl std::fmt::Display for QueryComponentError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            QueryComponentError::MissingReadAccess => {
+                write!(
+                    f,
+                    "This query does not have read access to the requested component."
+                )
+            }
+            QueryComponentError::MissingWriteAccess => {
+                write!(
+                    f,
+                    "This query does not have write access to the requested component."
+                )
+            }
+            QueryComponentError::MissingComponent => {
+                write!(f, "The given entity does not have the requested component.")
+            }
+            QueryComponentError::NoSuchEntity => {
+                write!(f, "The requested entity does not exist.")
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{prelude::*, query::QueryEntityError};

crates/bevy_ecs/src/system/mod.rs

@@ -548,8 +548,8 @@ mod tests {
             Schedule,
         },
         system::{
-            Commands, In, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query,
-            QueryComponentError, Res, ResMut, Resource, System, SystemState,
+            Commands, In, IntoSystem, Local, NonSend, NonSendMut, ParamSet, Query, Res, ResMut,
+            Resource, System, SystemState,
         },
         world::{FromWorld, World},
     };
@@ -1759,6 +1759,8 @@ mod tests {
 
     #[test]
     fn readonly_query_get_mut_component_fails() {
+        use crate::query::QueryComponentError;
+
         let mut world = World::new();
         let entity = world.spawn(W(42u32)).id();
         run_system(&mut world, move |q: Query<&mut W<u32>>| {