Moved get_component(_unchecked_mut) from Query to QueryState

crates/bevy_ecs/src/query/error.rs

@@ -0,0 +1,151 @@
+use std::fmt;
+
+use crate::entity::Entity;
+
+/// An error that occurs when retrieving a specific [`Entity`]'s query result from [`Query`](crate::system::Query) or [`QueryState`](crate::query::QueryState).
+// TODO: return the type_name as part of this error
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
+pub enum QueryEntityError {
+    /// The given [`Entity`]'s components do not match the query.
+    ///
+    /// Either it does not have a requested component, or it has a component which the query filters out.
+    QueryDoesNotMatch(Entity),
+    /// The given [`Entity`] does not exist.
+    NoSuchEntity(Entity),
+    /// The [`Entity`] was requested mutably more than once.
+    ///
+    /// See [`QueryState::get_many_mut`](crate::query::QueryState::get_many_mut) for an example.
+    AliasedMutability(Entity),
+}
+
+impl std::error::Error for QueryEntityError {}
+
+impl fmt::Display for QueryEntityError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            QueryEntityError::QueryDoesNotMatch(_) => {
+                write!(f, "The given entity's components do not match the query.")
+            }
+            QueryEntityError::NoSuchEntity(_) => write!(f, "The requested entity does not exist."),
+            QueryEntityError::AliasedMutability(_) => {
+                write!(f, "The entity was requested mutably more than once.")
+            }
+        }
+    }
+}
+
+/// 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.")
+            }
+        }
+    }
+}
+
+/// An error that occurs when evaluating a [`Query`](crate::system::Query) or [`QueryState`](crate::query::QueryState) as a single expected result via
+/// [`get_single`](crate::system::Query::get_single) or [`get_single_mut`](crate::system::Query::get_single_mut).
+#[derive(Debug)]
+pub enum QuerySingleError {
+    /// No entity fits the query.
+    NoEntities(&'static str),
+    /// Multiple entities fit the query.
+    MultipleEntities(&'static str),
+}
+
+impl std::error::Error for QuerySingleError {}
+
+impl std::fmt::Display for QuerySingleError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            QuerySingleError::NoEntities(query) => write!(f, "No entities fit the query {query}"),
+            QuerySingleError::MultipleEntities(query) => {
+                write!(f, "Multiple entities fit the query {query}!")
+            }
+        }
+    }
+}

crates/bevy_ecs/src/query/mod.rs

@@ -1,13 +1,15 @@
 //! Contains APIs for retrieving component data from the world.
 
 mod access;
+mod error;
 mod fetch;
 mod filter;
 mod iter;
 mod par_iter;
 mod state;
 
 pub use access::*;
+pub use error::*;
 pub use fetch::*;
 pub use filter::*;
 pub use iter::*;

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,9 +14,12 @@ 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};
+use super::{
+    NopWorldQuery, QueryComponentError, QueryEntityError, QueryManyIter, QuerySingleError,
+    ROQueryItem, ReadOnlyWorldQuery,
+};
 
 /// Provides scoped access to a [`World`] state according to a given [`WorldQuery`] and query filter.
 #[repr(C)]
@@ -506,6 +510,110 @@ 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)
+        {
+            // 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)
+        }
+    }
+
     /// 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.
     ///
@@ -1196,7 +1304,10 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     #[track_caller]
     #[inline]
     pub fn single<'w>(&mut self, world: &'w World) -> ROQueryItem<'w, Q> {
-        self.get_single(world).unwrap()
+        match self.get_single(world) {
+            Ok(items) => items,
+            Err(error) => panic!("Cannot get single mutable query result: {error}"),
+        }
     }
 
     /// Returns a single immutable query result when there is exactly one entity matching
@@ -1235,7 +1346,10 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     #[inline]
     pub fn single_mut<'w>(&mut self, world: &'w mut World) -> Q::Item<'w> {
         // SAFETY: query has unique world access
-        self.get_single_mut(world).unwrap()
+        match self.get_single_mut(world) {
+            Ok(items) => items,
+            Err(error) => panic!("Cannot get single query result: {error}"),
+        }
     }
 
     /// Returns a single mutable query result when there is exactly one entity matching
@@ -1311,38 +1425,6 @@ impl<Q: WorldQuery, F: ReadOnlyWorldQuery> QueryState<Q, F> {
     }
 }
 
-/// An error that occurs when retrieving a specific [`Entity`]'s query result from [`Query`](crate::system::Query) or [`QueryState`].
-// TODO: return the type_name as part of this error
-#[derive(Debug, PartialEq, Eq, Clone, Copy)]
-pub enum QueryEntityError {
-    /// The given [`Entity`]'s components do not match the query.
-    ///
-    /// Either it does not have a requested component, or it has a component which the query filters out.
-    QueryDoesNotMatch(Entity),
-    /// The given [`Entity`] does not exist.
-    NoSuchEntity(Entity),
-    /// The [`Entity`] was requested mutably more than once.
-    ///
-    /// See [`QueryState::get_many_mut`] for an example.
-    AliasedMutability(Entity),
-}
-
-impl std::error::Error for QueryEntityError {}
-
-impl fmt::Display for QueryEntityError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        match self {
-            QueryEntityError::QueryDoesNotMatch(_) => {
-                write!(f, "The given entity's components do not match the query.")
-            }
-            QueryEntityError::NoSuchEntity(_) => write!(f, "The requested entity does not exist."),
-            QueryEntityError::AliasedMutability(_) => {
-                write!(f, "The entity was requested mutably more than once.")
-            }
-        }
-    }
-}
-
 #[cfg(test)]
 mod tests {
     use crate::{prelude::*, query::QueryEntityError};
@@ -1451,26 +1533,3 @@ mod tests {
         let _panics = query_state.get_many_mut(&mut world_2, []);
     }
 }
-
-/// An error that occurs when evaluating a [`Query`](crate::system::Query) or [`QueryState`] as a single expected result via
-/// [`get_single`](crate::system::Query::get_single) or [`get_single_mut`](crate::system::Query::get_single_mut).
-#[derive(Debug)]
-pub enum QuerySingleError {
-    /// No entity fits the query.
-    NoEntities(&'static str),
-    /// Multiple entities fit the query.
-    MultipleEntities(&'static str),
-}
-
-impl std::error::Error for QuerySingleError {}
-
-impl std::fmt::Display for QuerySingleError {
-    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        match self {
-            QuerySingleError::NoEntities(query) => write!(f, "No entities fit the query {query}"),
-            QuerySingleError::MultipleEntities(query) => {
-                write!(f, "Multiple entities fit the query {query}!")
-            }
-        }
-    }
-}