From 7c6057bc69cd7263a2971d8653675a8c9c194710 Mon Sep 17 00:00:00 2001 From: Christian Hughes <9044780+ItsDoot@users.noreply.github.com> Date: Wed, 2 Oct 2024 12:02:20 -0700 Subject: [PATCH] Enable `EntityRef::get_by_id` and friends to take multiple ids and get multiple pointers back (#15593) # Objective - Closes #15577 ## Solution The following functions can now also take multiple component IDs and return multiple pointers back: - `EntityRef::get_by_id` - `EntityMut::get_by_id` - `EntityMut::into_borrow_by_id` - `EntityMut::get_mut_by_id` - `EntityMut::into_mut_by_id` - `EntityWorldMut::get_by_id` - `EntityWorldMut::into_borrow_by_id` - `EntityWorldMut::get_mut_by_id` - `EntityWorldMut::into_mut_by_id` If you pass in X, you receive Y: - give a single `ComponentId`, receive a single `Ptr`/`MutUntyped` - give a `[ComponentId; N]` (array), receive a `[Ptr; N]`/`[MutUntyped; N]` - give a `&[ComponentId; N]` (array), receive a `[Ptr; N]`/`[MutUntyped; N]` - give a `&[ComponentId]` (slice), receive a `Vec`/`Vec` - give a `&HashSet`, receive a `HashMap`/`HashMap` ## Testing - Added 4 new tests. --- ## Migration Guide - The following functions now return an `Result<_, EntityComponentError>` instead of a `Option<_>`: `EntityRef::get_by_id`, `EntityMut::get_by_id`, `EntityMut::into_borrow_by_id`, `EntityMut::get_mut_by_id`, `EntityMut::into_mut_by_id`, `EntityWorldMut::get_by_id`, `EntityWorldMut::into_borrow_by_id`, `EntityWorldMut::get_mut_by_id`, `EntityWorldMut::into_mut_by_id` --- crates/bevy_ecs/src/world/entity_ref.rs | 988 +++++++++++++++++++++--- crates/bevy_ecs/src/world/error.rs | 13 +- 2 files changed, 914 insertions(+), 87 deletions(-) diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 3eb20be577209..023db78f5dbdd 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -10,10 +10,11 @@ use crate::{ removal_detection::RemovedComponentEvents, storage::Storages, system::IntoObserverSystem, - world::{DeferredWorld, Mut, World}, + world::{error::EntityComponentError, DeferredWorld, Mut, World}, }; use bevy_ptr::{OwningPtr, Ptr}; -use core::{any::TypeId, marker::PhantomData}; +use bevy_utils::{HashMap, HashSet}; +use core::{any::TypeId, marker::PhantomData, mem::MaybeUninit}; use thiserror::Error; use super::{unsafe_world_cell::UnsafeEntityCell, Ref, ON_REMOVE, ON_REPLACE}; @@ -143,18 +144,117 @@ impl<'w> EntityRef<'w> { unsafe { self.0.get_change_ticks_by_id(component_id) } } - /// Gets the component of the given [`ComponentId`] from the entity. + /// Returns [untyped read-only reference(s)](Ptr) to component(s) for the + /// current entity, based on the given [`ComponentId`]s. /// - /// **You should prefer to use the typed API where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityRef::get`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityRef::get`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// + /// # Examples + /// + /// ## Single [`ComponentId`] + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Foo(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn(Foo(42)).id(); + /// + /// // Grab the component ID for `Foo` in whatever way you like. + /// let component_id = world.register_component::(); + /// + /// // Then, get the component by ID. + /// let ptr = world.entity(entity).get_by_id(component_id); + /// # assert_eq!(unsafe { ptr.unwrap().deref::() }, &Foo(42)); + /// ``` + /// + /// ## Array of [`ComponentId`]s + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); + /// + /// // Then, get the components by ID. You'll receive a same-sized array. + /// let Ok([x_ptr, y_ptr]) = world.entity(entity).get_by_id([x_id, y_id]) else { + /// // Up to you to handle if a component is missing from the entity. + /// # unreachable!(); + /// }; + /// # assert_eq!((unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }), (&X(42), &Y(10))); + /// ``` + /// + /// ## Slice of [`ComponentId`]s + /// + /// ``` + /// # use bevy_ecs::{prelude::*, component::ComponentId}; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); /// - /// Unlike [`EntityRef::get`], this returns a raw pointer to the component, - /// which is only valid while the `'w` borrow of the lifetime is active. + /// // Then, get the components by ID. You'll receive a vec of ptrs. + /// let ptrs = world.entity(entity).get_by_id(&[x_id, y_id] as &[ComponentId]); + /// # let ptrs = ptrs.unwrap(); + /// # assert_eq!((unsafe { ptrs[0].deref::() }, unsafe { ptrs[1].deref::() }), (&X(42), &Y(10))); + /// ``` + /// + /// ## [`HashSet`] of [`ComponentId`]s + /// + /// ``` + /// # use bevy_utils::HashSet; + /// # use bevy_ecs::{prelude::*, component::ComponentId}; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); + /// + /// // Then, get the components by ID. You'll receive a vec of ptrs. + /// let ptrs = world.entity(entity).get_by_id(&HashSet::from_iter([x_id, y_id])); + /// # let ptrs = ptrs.unwrap(); + /// # assert_eq!((unsafe { ptrs[&x_id].deref::() }, unsafe { ptrs[&y_id].deref::() }), (&X(42), &Y(10))); + /// ``` #[inline] - pub fn get_by_id(&self, component_id: ComponentId) -> Option> { + pub fn get_by_id( + &self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: We have read-only access to all components of this entity. - unsafe { self.0.get_by_id(component_id) } + unsafe { component_ids.fetch_ref(self.0) } } /// Returns read-only components for the current entity that match the query `Q`. @@ -448,59 +548,219 @@ impl<'w> EntityMut<'w> { self.as_readonly().get_change_ticks_by_id(component_id) } - /// Gets the component of the given [`ComponentId`] from the entity. + /// Returns [untyped read-only reference(s)](Ptr) to component(s) for the + /// current entity, based on the given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityWorldMut::get`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityMut::get`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityMut::get`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors /// - /// Unlike [`EntityMut::get`], this returns a raw pointer to the component, - /// which is only valid while the [`EntityMut`] is alive. + /// Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityRef::get_by_id`]. #[inline] - pub fn get_by_id(&self, component_id: ComponentId) -> Option> { - self.as_readonly().get_by_id(component_id) + pub fn get_by_id( + &self, + component_ids: F, + ) -> Result, EntityComponentError> { + self.as_readonly().get_by_id(component_ids) } - /// Consumes `self` and gets the component of the given [`ComponentId`] with - /// world `'w` lifetime from the entity. + /// Consumes `self` and returns [untyped read-only reference(s)](Ptr) to + /// component(s) with lifetime `'w` for the current entity, based on the + /// given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityWorldMut::into_borrow`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityMut::into_borrow`] + /// where possible and only use this in cases where the actual component + /// types are not known at compile time.** + /// + /// Unlike [`EntityMut::into_borrow`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityRef::get_by_id`]. #[inline] - pub fn into_borrow_by_id(self, component_id: ComponentId) -> Option> { + pub fn into_borrow_by_id( + self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: - // consuming `self` ensures that no references exist to this entity's components. - unsafe { self.0.get_by_id(component_id) } + // - We have read-only access to all components of this entity. + // - consuming `self` ensures that no references exist to this entity's components. + unsafe { component_ids.fetch_ref(self.0) } } - /// Gets a [`MutUntyped`] of the component of the given [`ComponentId`] from the entity. + /// Returns [untyped mutable reference(s)](MutUntyped) to component(s) for + /// the current entity, based on the given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityMut::get_mut`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityMut::get_mut`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityMut::get_mut`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// - Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// - Returns [`EntityComponentError::AliasedMutability`] if a component + /// is requested multiple times. + /// + /// # Examples + /// + /// ## Single [`ComponentId`] + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Foo(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn(Foo(42)).id(); + /// + /// // Grab the component ID for `Foo` in whatever way you like. + /// let component_id = world.register_component::(); + /// + /// // Then, get the component by ID. + /// let mut entity_mut = world.entity_mut(entity); + /// let mut ptr = entity_mut.get_mut_by_id(component_id) + /// # .unwrap(); + /// # assert_eq!(unsafe { ptr.as_mut().deref_mut::() }, &mut Foo(42)); + /// ``` + /// + /// ## Array of [`ComponentId`]s + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); + /// + /// // Then, get the components by ID. You'll receive a same-sized array. + /// let mut entity_mut = world.entity_mut(entity); + /// let Ok([mut x_ptr, mut y_ptr]) = entity_mut.get_mut_by_id([x_id, y_id]) else { + /// // Up to you to handle if a component is missing from the entity. + /// # unreachable!(); + /// }; + /// # assert_eq!((unsafe { x_ptr.as_mut().deref_mut::() }, unsafe { y_ptr.as_mut().deref_mut::() }), (&mut X(42), &mut Y(10))); + /// ``` + /// + /// ## Slice of [`ComponentId`]s + /// + /// ``` + /// # use bevy_ecs::{prelude::*, component::ComponentId, change_detection::MutUntyped}; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); + /// + /// // Then, get the components by ID. You'll receive a vec of ptrs. + /// let mut entity_mut = world.entity_mut(entity); + /// let ptrs = entity_mut.get_mut_by_id(&[x_id, y_id] as &[ComponentId]) + /// # .unwrap(); + /// # let [mut x_ptr, mut y_ptr]: [MutUntyped; 2] = ptrs.try_into().unwrap(); + /// # assert_eq!((unsafe { x_ptr.as_mut().deref_mut::() }, unsafe { y_ptr.as_mut().deref_mut::() }), (&mut X(42), &mut Y(10))); + /// ``` /// - /// Unlike [`EntityMut::get_mut`], this returns a raw pointer to the component, - /// which is only valid while the [`EntityMut`] is alive. + /// ## [`HashSet`] of [`ComponentId`]s + /// + /// ``` + /// # use bevy_utils::HashSet; + /// # use bevy_ecs::{prelude::*, component::ComponentId}; + /// # + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct X(i32); + /// # #[derive(Component, PartialEq, Debug)] + /// # pub struct Y(i32); + /// # let mut world = World::new(); + /// let entity = world.spawn((X(42), Y(10))).id(); + /// + /// // Grab the component IDs for `X` and `Y` in whatever way you like. + /// let x_id = world.register_component::(); + /// let y_id = world.register_component::(); + /// + /// // Then, get the components by ID. You'll receive a `HashMap` of ptrs. + /// let mut entity_mut = world.entity_mut(entity); + /// let mut ptrs = entity_mut.get_mut_by_id(&HashSet::from_iter([x_id, y_id])) + /// # .unwrap(); + /// # let [mut x_ptr, mut y_ptr] = ptrs.get_many_mut([&x_id, &y_id]).unwrap(); + /// # assert_eq!((unsafe { x_ptr.as_mut().deref_mut::() }, unsafe { y_ptr.as_mut().deref_mut::() }), (&mut X(42), &mut Y(10))); + /// ``` #[inline] - pub fn get_mut_by_id(&mut self, component_id: ComponentId) -> Option> { + pub fn get_mut_by_id( + &mut self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: // - `&mut self` ensures that no references exist to this entity's components. - // - `as_unsafe_world_cell` gives mutable permission for all components on this entity - unsafe { self.0.get_mut_by_id(component_id) } + // - We have exclusive access to all components of this entity. + unsafe { component_ids.fetch_mut(self.0) } } - /// Consumes `self` and gets a [`MutUntyped<'w>`] of the component of the given [`ComponentId`] - /// with world `'w` lifetime from the entity. + /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped) + /// to component(s) with lifetime `'w` for the current entity, based on the + /// given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityMut::into_mut`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityMut::into_mut`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityMut::into_mut`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// - Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// - Returns [`EntityComponentError::AliasedMutability`] if a component + /// is requested multiple times. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityMut::get_mut_by_id`]. #[inline] - pub fn into_mut_by_id(self, component_id: ComponentId) -> Option> { + pub fn into_mut_by_id( + self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: - // consuming `self` ensures that no references exist to this entity's components. - unsafe { self.0.get_mut_by_id(component_id) } + // - consuming `self` ensures that no references exist to this entity's components. + // - We have exclusive access to all components of this entity. + unsafe { component_ids.fetch_mut(self.0) } } } @@ -761,58 +1021,127 @@ impl<'w> EntityWorldMut<'w> { EntityRef::from(self).get_change_ticks_by_id(component_id) } - /// Gets the component of the given [`ComponentId`] from the entity. + /// Returns [untyped read-only reference(s)](Ptr) to component(s) for the + /// current entity, based on the given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityWorldMut::get`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityWorldMut::get`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityWorldMut::get`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). /// - /// Unlike [`EntityWorldMut::get`], this returns a raw pointer to the component, - /// which is only valid while the [`EntityWorldMut`] is alive. + /// # Errors + /// + /// Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityRef::get_by_id`]. #[inline] - pub fn get_by_id(&self, component_id: ComponentId) -> Option> { - EntityRef::from(self).get_by_id(component_id) + pub fn get_by_id( + &self, + component_ids: F, + ) -> Result, EntityComponentError> { + EntityRef::from(self).get_by_id(component_ids) } - /// Consumes `self` and gets the component of the given [`ComponentId`] with - /// with world `'w` lifetime from the entity. + /// Consumes `self` and returns [untyped read-only reference(s)](Ptr) to + /// component(s) with lifetime `'w` for the current entity, based on the + /// given [`ComponentId`]s. + /// + /// **You should prefer to use the typed API [`EntityWorldMut::into_borrow`] + /// where possible and only use this in cases where the actual component + /// types are not known at compile time.** + /// + /// Unlike [`EntityWorldMut::into_borrow`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. /// - /// **You should prefer to use the typed API [`EntityWorldMut::into_borrow`] where - /// possible and only use this in cases where the actual component types are not - /// known at compile time.** + /// # Examples + /// + /// For examples on how to use this method, see [`EntityRef::get_by_id`]. #[inline] - pub fn into_borrow_by_id(self, component_id: ComponentId) -> Option> { - // SAFETY: consuming `self` implies exclusive access - unsafe { self.into_unsafe_entity_cell().get_by_id(component_id) } + pub fn into_borrow_by_id( + self, + component_ids: F, + ) -> Result, EntityComponentError> { + // SAFETY: + // - We have read-only access to all components of this entity. + // - consuming `self` ensures that no references exist to this entity's components. + unsafe { component_ids.fetch_ref(self.into_unsafe_entity_cell()) } } - /// Gets a [`MutUntyped`] of the component of the given [`ComponentId`] from the entity. + /// Returns [untyped mutable reference(s)](MutUntyped) to component(s) for + /// the current entity, based on the given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityWorldMut::get_mut`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityWorldMut::get_mut`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityWorldMut::get_mut`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors /// - /// Unlike [`EntityWorldMut::get_mut`], this returns a raw pointer to the component, - /// which is only valid while the [`EntityWorldMut`] is alive. + /// - Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// - Returns [`EntityComponentError::AliasedMutability`] if a component + /// is requested multiple times. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityMut::get_mut_by_id`]. #[inline] - pub fn get_mut_by_id(&mut self, component_id: ComponentId) -> Option> { + pub fn get_mut_by_id( + &mut self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: // - `&mut self` ensures that no references exist to this entity's components. - // - `as_unsafe_world_cell` gives mutable permission for all components on this entity - unsafe { self.as_unsafe_entity_cell().get_mut_by_id(component_id) } + // - We have exclusive access to all components of this entity. + unsafe { component_ids.fetch_mut(self.as_unsafe_entity_cell()) } } - /// Consumes `self` and gets a [`MutUntyped<'w>`] of the component with the world `'w` lifetime - /// of the given [`ComponentId`] from the entity. + /// Consumes `self` and returns [untyped mutable reference(s)](MutUntyped) + /// to component(s) with lifetime `'w` for the current entity, based on the + /// given [`ComponentId`]s. /// - /// **You should prefer to use the typed API [`EntityWorldMut::into_mut`] where possible and only - /// use this in cases where the actual component types are not known at - /// compile time.** + /// **You should prefer to use the typed API [`EntityWorldMut::into_mut`] where + /// possible and only use this in cases where the actual component types + /// are not known at compile time.** + /// + /// Unlike [`EntityWorldMut::into_mut`], this returns untyped reference(s) to + /// component(s), and it's the job of the caller to ensure the correct + /// type(s) are dereferenced (if necessary). + /// + /// # Errors + /// + /// - Returns [`EntityComponentError::MissingComponent`] if the entity does + /// not have a component. + /// - Returns [`EntityComponentError::AliasedMutability`] if a component + /// is requested multiple times. + /// + /// # Examples + /// + /// For examples on how to use this method, see [`EntityMut::get_mut_by_id`]. #[inline] - pub fn into_mut_by_id(self, component_id: ComponentId) -> Option> { + pub fn into_mut_by_id( + self, + component_ids: F, + ) -> Result, EntityComponentError> { // SAFETY: - // consuming `self` ensures that no references exist to this entity's components. - unsafe { self.into_unsafe_entity_cell().get_mut_by_id(component_id) } + // - consuming `self` ensures that no references exist to this entity's components. + // - We have exclusive access to all components of this entity. + unsafe { component_ids.fetch_mut(self.into_unsafe_entity_cell()) } } /// Adds a [`Bundle`] of components to the entity. @@ -2841,17 +3170,269 @@ pub(crate) unsafe fn take_component<'a>( } } +/// Types that can be used to fetch components from an entity dynamically by +/// [`ComponentId`]s. +/// +/// Provided implementations are: +/// - [`ComponentId`]: Returns a single untyped reference. +/// - `[ComponentId; N]` and `&[ComponentId; N]`: Returns a same-sized array of untyped references. +/// - `&[ComponentId]`: Returns a [`Vec`] of untyped references. +/// - [`&HashSet`](HashSet): Returns a [`HashMap`] of IDs to untyped references. +/// +/// # Performance +/// +/// - The slice and array implementations perform an aliased mutability check in +/// [`DynamicComponentFetch::fetch_mut`] that is `O(N^2)`. +/// - The [`HashSet`] implementation performs no such check as the type itself +/// guarantees unique IDs. +/// - The single [`ComponentId`] implementation performs no such check as only +/// one reference is returned. +/// +/// # Safety +/// +/// Implementor must ensure that: +/// - No aliased mutability is caused by the returned references. +/// - [`DynamicComponentFetch::fetch_ref`] returns only read-only references. +pub unsafe trait DynamicComponentFetch { + /// The read-only reference type returned by [`DynamicComponentFetch::fetch_ref`]. + type Ref<'w>; + + /// The mutable reference type returned by [`DynamicComponentFetch::fetch_mut`]. + type Mut<'w>; + + /// Returns untyped read-only reference(s) to the component(s) with the + /// given [`ComponentId`]s, as determined by `self`. + /// + /// # Safety + /// + /// It is the caller's responsibility to ensure that: + /// - The given [`UnsafeEntityCell`] has read-only access to the fetched components. + /// - No other mutable references to the fetched components exist at the same time. + /// + /// # Errors + /// + /// - Returns [`EntityComponentError::MissingComponent`] if a component is missing from the entity. + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError>; + + /// Returns untyped mutable reference(s) to the component(s) with the + /// given [`ComponentId`]s, as determined by `self`. + /// + /// # Safety + /// + /// It is the caller's responsibility to ensure that: + /// - The given [`UnsafeEntityCell`] has mutable access to the fetched components. + /// - No other references to the fetched components exist at the same time. + /// + /// # Errors + /// + /// - Returns [`EntityComponentError::MissingComponent`] if a component is missing from the entity. + /// - Returns [`EntityComponentError::AliasedMutability`] if a component is requested multiple times. + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError>; +} + +// SAFETY: +// - No aliased mutability is caused because a single reference is returned. +// - No mutable references are returned by `fetch_ref`. +unsafe impl DynamicComponentFetch for ComponentId { + type Ref<'w> = Ptr<'w>; + type Mut<'w> = MutUntyped<'w>; + + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + // SAFETY: caller ensures that the cell has read access to the component. + unsafe { cell.get_by_id(self) }.ok_or(EntityComponentError::MissingComponent(self)) + } + + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + // SAFETY: caller ensures that the cell has mutable access to the component. + unsafe { cell.get_mut_by_id(self) }.ok_or(EntityComponentError::MissingComponent(self)) + } +} + +// SAFETY: +// - No aliased mutability is caused because the array is checked for duplicates. +// - No mutable references are returned by `fetch_ref`. +unsafe impl DynamicComponentFetch for [ComponentId; N] { + type Ref<'w> = [Ptr<'w>; N]; + type Mut<'w> = [MutUntyped<'w>; N]; + + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + <&Self>::fetch_ref(&self, cell) + } + + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + <&Self>::fetch_mut(&self, cell) + } +} + +// SAFETY: +// - No aliased mutability is caused because the array is checked for duplicates. +// - No mutable references are returned by `fetch_ref`. +unsafe impl DynamicComponentFetch for &'_ [ComponentId; N] { + type Ref<'w> = [Ptr<'w>; N]; + type Mut<'w> = [MutUntyped<'w>; N]; + + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + let mut ptrs = [const { MaybeUninit::uninit() }; N]; + for (ptr, &id) in core::iter::zip(&mut ptrs, self) { + *ptr = MaybeUninit::new( + // SAFETY: caller ensures that the cell has read access to the component. + unsafe { cell.get_by_id(id) }.ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + + // SAFETY: Each ptr was initialized in the loop above. + let ptrs = ptrs.map(|ptr| unsafe { MaybeUninit::assume_init(ptr) }); + + Ok(ptrs) + } + + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + // Check for duplicate component IDs. + for i in 0..self.len() { + for j in 0..i { + if self[i] == self[j] { + return Err(EntityComponentError::AliasedMutability(self[i])); + } + } + } + + let mut ptrs = [const { MaybeUninit::uninit() }; N]; + for (ptr, &id) in core::iter::zip(&mut ptrs, self) { + *ptr = MaybeUninit::new( + // SAFETY: caller ensures that the cell has mutable access to the component. + unsafe { cell.get_mut_by_id(id) } + .ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + + // SAFETY: Each ptr was initialized in the loop above. + let ptrs = ptrs.map(|ptr| unsafe { MaybeUninit::assume_init(ptr) }); + + Ok(ptrs) + } +} + +// SAFETY: +// - No aliased mutability is caused because the slice is checked for duplicates. +// - No mutable references are returned by `fetch_ref`. +unsafe impl DynamicComponentFetch for &'_ [ComponentId] { + type Ref<'w> = Vec>; + type Mut<'w> = Vec>; + + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + let mut ptrs = Vec::with_capacity(self.len()); + for &id in self { + ptrs.push( + // SAFETY: caller ensures that the cell has read access to the component. + unsafe { cell.get_by_id(id) }.ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + Ok(ptrs) + } + + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + // Check for duplicate component IDs. + for i in 0..self.len() { + for j in 0..i { + if self[i] == self[j] { + return Err(EntityComponentError::AliasedMutability(self[i])); + } + } + } + + let mut ptrs = Vec::with_capacity(self.len()); + for &id in self { + ptrs.push( + // SAFETY: caller ensures that the cell has mutable access to the component. + unsafe { cell.get_mut_by_id(id) } + .ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + Ok(ptrs) + } +} + +// SAFETY: +// - No aliased mutability is caused because `HashSet` guarantees unique elements. +// - No mutable references are returned by `fetch_ref`. +unsafe impl DynamicComponentFetch for &'_ HashSet { + type Ref<'w> = HashMap>; + type Mut<'w> = HashMap>; + + unsafe fn fetch_ref( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + let mut ptrs = HashMap::with_capacity(self.len()); + for &id in self { + ptrs.insert( + id, + // SAFETY: caller ensures that the cell has read access to the component. + unsafe { cell.get_by_id(id) }.ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + Ok(ptrs) + } + + unsafe fn fetch_mut( + self, + cell: UnsafeEntityCell<'_>, + ) -> Result, EntityComponentError> { + let mut ptrs = HashMap::with_capacity(self.len()); + for &id in self { + ptrs.insert( + id, + // SAFETY: caller ensures that the cell has mutable access to the component. + unsafe { cell.get_mut_by_id(id) } + .ok_or(EntityComponentError::MissingComponent(id))?, + ); + } + Ok(ptrs) + } +} + #[cfg(test)] mod tests { - use bevy_ptr::OwningPtr; + use bevy_ptr::{OwningPtr, Ptr}; use core::panic::AssertUnwindSafe; use crate::{ self as bevy_ecs, + change_detection::MutUntyped, component::ComponentId, prelude::*, system::{assert_is_system, RunSystemOnce as _}, - world::{FilteredEntityMut, FilteredEntityRef}, + world::{error::EntityComponentError, FilteredEntityMut, FilteredEntityRef}, }; use super::{EntityMutExcept, EntityRefExcept}; @@ -2935,7 +3516,7 @@ mod tests { let mut world = World::new(); let entity = world.spawn_empty().id(); let entity = world.entity(entity); - assert!(entity.get_by_id(invalid_component_id).is_none()); + assert!(entity.get_by_id(invalid_component_id).is_err()); } #[test] @@ -2944,8 +3525,8 @@ mod tests { let mut world = World::new(); let mut entity = world.spawn_empty(); - assert!(entity.get_by_id(invalid_component_id).is_none()); - assert!(entity.get_mut_by_id(invalid_component_id).is_none()); + assert!(entity.get_by_id(invalid_component_id).is_err()); + assert!(entity.get_mut_by_id(invalid_component_id).is_err()); } // regression test for https://github.com/bevyengine/bevy/pull/7387 @@ -3575,13 +4156,14 @@ mod tests { assert!(e.get_change_ticks_by_id(a_id).is_none()); } + #[derive(Component, PartialEq, Eq, Debug)] + struct X(usize); + + #[derive(Component, PartialEq, Eq, Debug)] + struct Y(usize); + #[test] fn get_components() { - #[derive(Component, PartialEq, Eq, Debug)] - struct X(usize); - - #[derive(Component, PartialEq, Eq, Debug)] - struct Y(usize); let mut world = World::default(); let e1 = world.spawn((X(7), Y(10))).id(); let e2 = world.spawn(X(8)).id(); @@ -3594,4 +4176,238 @@ mod tests { assert_eq!(None, world.entity(e2).get_components::<(&X, &Y)>()); assert_eq!(None, world.entity(e3).get_components::<(&X, &Y)>()); } + + #[test] + fn get_by_id_array() { + let mut world = World::default(); + let e1 = world.spawn((X(7), Y(10))).id(); + let e2 = world.spawn(X(8)).id(); + let e3 = world.spawn_empty().id(); + + let x_id = world.register_component::(); + let y_id = world.register_component::(); + + assert_eq!( + Ok((&X(7), &Y(10))), + world + .entity(e1) + .get_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(y_id)), + world + .entity(e2) + .get_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(x_id)), + world + .entity(e3) + .get_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + } + + #[test] + fn get_by_id_vec() { + let mut world = World::default(); + let e1 = world.spawn((X(7), Y(10))).id(); + let e2 = world.spawn(X(8)).id(); + let e3 = world.spawn_empty().id(); + + let x_id = world.register_component::(); + let y_id = world.register_component::(); + + assert_eq!( + Ok((&X(7), &Y(10))), + world + .entity(e1) + .get_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[Ptr; 2], _> = ptrs.try_into() else { + panic!("get_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(y_id)), + world + .entity(e2) + .get_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[Ptr; 2], _> = ptrs.try_into() else { + panic!("get_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(x_id)), + world + .entity(e3) + .get_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[Ptr; 2], _> = ptrs.try_into() else { + panic!("get_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.deref::() }, unsafe { y_ptr.deref::() }) + }) + ); + } + + #[test] + fn get_mut_by_id_array() { + let mut world = World::default(); + let e1 = world.spawn((X(7), Y(10))).id(); + let e2 = world.spawn(X(8)).id(); + let e3 = world.spawn_empty().id(); + + let x_id = world.register_component::(); + let y_id = world.register_component::(); + + assert_eq!( + Ok((&mut X(7), &mut Y(10))), + world + .entity_mut(e1) + .get_mut_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(y_id)), + world + .entity_mut(e2) + .get_mut_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(x_id)), + world + .entity_mut(e3) + .get_mut_by_id([x_id, y_id]) + .map(|[x_ptr, y_ptr]| { + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + + assert_eq!( + Err(EntityComponentError::AliasedMutability(x_id)), + world + .entity_mut(e1) + .get_mut_by_id([x_id, x_id]) + .map(|_| { unreachable!() }) + ); + assert_eq!( + Err(EntityComponentError::AliasedMutability(x_id)), + world + .entity_mut(e3) + .get_mut_by_id([x_id, x_id]) + .map(|_| { unreachable!() }) + ); + } + + #[test] + fn get_mut_by_id_vec() { + let mut world = World::default(); + let e1 = world.spawn((X(7), Y(10))).id(); + let e2 = world.spawn(X(8)).id(); + let e3 = world.spawn_empty().id(); + + let x_id = world.register_component::(); + let y_id = world.register_component::(); + + assert_eq!( + Ok((&mut X(7), &mut Y(10))), + world + .entity_mut(e1) + .get_mut_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[MutUntyped; 2], _> = ptrs.try_into() else { + panic!("get_mut_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(y_id)), + world + .entity_mut(e2) + .get_mut_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[MutUntyped; 2], _> = ptrs.try_into() else { + panic!("get_mut_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + assert_eq!( + Err(EntityComponentError::MissingComponent(x_id)), + world + .entity_mut(e3) + .get_mut_by_id(&[x_id, y_id] as &[ComponentId]) + .map(|ptrs| { + let Ok([x_ptr, y_ptr]): Result<[MutUntyped; 2], _> = ptrs.try_into() else { + panic!("get_mut_by_id(slice) didn't return 2 elements") + }; + + // SAFETY: components match the id they were fetched with + (unsafe { x_ptr.into_inner().deref_mut::() }, unsafe { + y_ptr.into_inner().deref_mut::() + }) + }) + ); + + assert_eq!( + Err(EntityComponentError::AliasedMutability(x_id)), + world + .entity_mut(e1) + .get_mut_by_id(&[x_id, x_id]) + .map(|_| { unreachable!() }) + ); + assert_eq!( + Err(EntityComponentError::AliasedMutability(x_id)), + world + .entity_mut(e3) + .get_mut_by_id(&[x_id, x_id]) + .map(|_| { unreachable!() }) + ); + } } diff --git a/crates/bevy_ecs/src/world/error.rs b/crates/bevy_ecs/src/world/error.rs index 326b0310ba15c..5fc4264e07cd3 100644 --- a/crates/bevy_ecs/src/world/error.rs +++ b/crates/bevy_ecs/src/world/error.rs @@ -2,7 +2,7 @@ use thiserror::Error; -use crate::schedule::InternedScheduleLabel; +use crate::{component::ComponentId, schedule::InternedScheduleLabel}; /// The error type returned by [`World::try_run_schedule`] if the provided schedule does not exist. /// @@ -10,3 +10,14 @@ use crate::schedule::InternedScheduleLabel; #[derive(Error, Debug)] #[error("The schedule with the label {0:?} was not found.")] pub struct TryRunScheduleError(pub InternedScheduleLabel); + +/// An error that occurs when dynamically retrieving components from an entity. +#[derive(Error, Debug, Clone, Copy, PartialEq, Eq)] +pub enum EntityComponentError { + /// The component with the given [`ComponentId`] does not exist on the entity. + #[error("The component with ID {0:?} does not exist on the entity.")] + MissingComponent(ComponentId), + /// The component with the given [`ComponentId`] was requested mutably more than once. + #[error("The component with ID {0:?} was requested mutably more than once.")] + AliasedMutability(ComponentId), +}