diff --git a/Cargo.toml b/Cargo.toml index c3f5d205408c8..892c22caaf499 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1352,6 +1352,17 @@ description = "Groups commonly used compound queries and query filters into a si category = "ECS (Entity Component System)" wasm = false +[[example]] +name = "dynamic" +path = "examples/ecs/dynamic.rs" +doc-scrape-examples = true + +[package.metadata.example.dynamic] +name = "Dynamic ECS" +description = "Dynamically create components, spawn entities with those components and query those components" +category = "ECS (Entity Component System)" +wasm = false + [[example]] name = "event" path = "examples/ecs/event.rs" diff --git a/crates/bevy_ecs/macros/src/world_query.rs b/crates/bevy_ecs/macros/src/world_query.rs index 5196d25deb192..7444a55c2a0f0 100644 --- a/crates/bevy_ecs/macros/src/world_query.rs +++ b/crates/bevy_ecs/macros/src/world_query.rs @@ -165,22 +165,18 @@ pub(crate) fn world_query_impl( #( <#field_types>::update_component_access(&state.#named_field_idents, _access); )* } - fn update_archetype_component_access( - state: &Self::State, - _archetype: &#path::archetype::Archetype, - _access: &mut #path::query::Access<#path::archetype::ArchetypeComponentId> - ) { - #( - <#field_types>::update_archetype_component_access(&state.#named_field_idents, _archetype, _access); - )* - } - fn init_state(world: &mut #path::world::World) -> #state_struct_name #user_ty_generics { #state_struct_name { #(#named_field_idents: <#field_types>::init_state(world),)* } } + fn get_state(world: &#path::world::World) -> Option<#state_struct_name #user_ty_generics> { + Some(#state_struct_name { + #(#named_field_idents: <#field_types>::get_state(world)?,)* + }) + } + fn matches_component_set(state: &Self::State, _set_contains_id: &impl Fn(#path::component::ComponentId) -> bool) -> bool { true #(&& <#field_types>::matches_component_set(&state.#named_field_idents, _set_contains_id))* } diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index 9b5fee48106be..c262c8b3ce2b1 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/crates/bevy_ecs/src/lib.rs @@ -36,7 +36,7 @@ pub mod prelude { component::Component, entity::Entity, event::{Event, EventReader, EventWriter, Events}, - query::{Added, AnyOf, Changed, Has, Or, QueryState, With, Without}, + query::{Added, AnyOf, Changed, Has, Or, QueryBuilder, QueryState, With, Without}, removal_detection::RemovedComponents, schedule::{ apply_deferred, apply_state_transition, common_conditions::*, Condition, diff --git a/crates/bevy_ecs/src/query/access.rs b/crates/bevy_ecs/src/query/access.rs index 808d31d2c48e7..d7253f073ed7f 100644 --- a/crates/bevy_ecs/src/query/access.rs +++ b/crates/bevy_ecs/src/query/access.rs @@ -157,6 +157,12 @@ impl Access { self.writes_all } + /// Removes all writes. + pub fn clear_writes(&mut self) { + self.writes_all = false; + self.writes.clear(); + } + /// Removes all accesses. pub fn clear(&mut self) { self.reads_all = false; @@ -198,6 +204,29 @@ impl Access { && other.writes.is_disjoint(&self.reads_and_writes) } + /// Returns `true` if the set is a subset of another, i.e. `other` contains + /// at least all the values in `self`. + pub fn is_subset(&self, other: &Access) -> bool { + if self.writes_all { + return other.writes_all; + } + + if other.writes_all { + return true; + } + + if self.reads_all { + return other.reads_all; + } + + if other.reads_all { + return self.writes.is_subset(&other.writes); + } + + self.reads_and_writes.is_subset(&other.reads_and_writes) + && self.writes.is_subset(&other.writes) + } + /// Returns a vector of elements that the access and `other` cannot access at the same time. pub fn get_conflicts(&self, other: &Access) -> Vec { let mut conflicts = FixedBitSet::default(); @@ -267,16 +296,18 @@ impl Access { /// See comments the [`WorldQuery`](super::WorldQuery) impls of [`AnyOf`](super::AnyOf)/`Option`/[`Or`](super::Or) for more information. #[derive(Debug, Clone, Eq, PartialEq)] pub struct FilteredAccess { - access: Access, + pub(crate) access: Access, + pub(crate) required: FixedBitSet, // An array of filter sets to express `With` or `Without` clauses in disjunctive normal form, for example: `Or<(With, With)>`. // Filters like `(With, Or<(With, Without)>` are expanded into `Or<((With, With), (With, Without))>`. - filter_sets: Vec>, + pub(crate) filter_sets: Vec>, } impl Default for FilteredAccess { fn default() -> Self { Self { access: Access::default(), + required: FixedBitSet::default(), filter_sets: vec![AccessFilters::default()], } } @@ -306,15 +337,23 @@ impl FilteredAccess { /// Adds access to the element given by `index`. pub fn add_read(&mut self, index: T) { self.access.add_read(index.clone()); + self.add_required(index.clone()); self.and_with(index); } /// Adds exclusive access to the element given by `index`. pub fn add_write(&mut self, index: T) { self.access.add_write(index.clone()); + self.add_required(index.clone()); self.and_with(index); } + fn add_required(&mut self, index: T) { + let index = index.sparse_set_index(); + self.required.grow(index + 1); + self.required.insert(index); + } + /// Adds a `With` filter: corresponds to a conjunction (AND) operation. /// /// Suppose we begin with `Or<(With, With)>`, which is represented by an array of two `AccessFilter` instances. @@ -391,6 +430,7 @@ impl FilteredAccess { /// `Or<((With, With), (With, Without), (Without, With), (Without, Without))>`. pub fn extend(&mut self, other: &FilteredAccess) { self.access.extend(&other.access); + self.required.union_with(&other.required); // We can avoid allocating a new array of bitsets if `other` contains just a single set of filters: // in this case we can short-circuit by performing an in-place union for each bitset. @@ -423,12 +463,18 @@ impl FilteredAccess { pub fn write_all(&mut self) { self.access.write_all(); } + + /// Returns `true` if the set is a subset of another, i.e. `other` contains + /// at least all the values in `self`. + pub fn is_subset(&self, other: &FilteredAccess) -> bool { + self.required.is_subset(&other.required) && self.access().is_subset(other.access()) + } } #[derive(Clone, Eq, PartialEq)] -struct AccessFilters { - with: FixedBitSet, - without: FixedBitSet, +pub(crate) struct AccessFilters { + pub(crate) with: FixedBitSet, + pub(crate) without: FixedBitSet, _index_type: PhantomData, } diff --git a/crates/bevy_ecs/src/query/builder.rs b/crates/bevy_ecs/src/query/builder.rs new file mode 100644 index 0000000000000..67db644a4fcb4 --- /dev/null +++ b/crates/bevy_ecs/src/query/builder.rs @@ -0,0 +1,401 @@ +use std::marker::PhantomData; + +use crate::{component::ComponentId, prelude::*}; + +use super::{FilteredAccess, QueryData, QueryFilter}; + +/// Builder struct to create [`QueryState`] instances at runtime. +/// +/// ``` +/// # use bevy_ecs::prelude::*; +/// # +/// # #[derive(Component)] +/// # struct A; +/// # +/// # #[derive(Component)] +/// # struct B; +/// # +/// # #[derive(Component)] +/// # struct C; +/// # +/// let mut world = World::new(); +/// let entity_a = world.spawn((A, B)).id(); +/// let entity_b = world.spawn((A, C)).id(); +/// +/// // Instantiate the builder using the type signature of the iterator you will consume +/// let mut query = QueryBuilder::<(Entity, &B)>::new(&mut world) +/// // Add additional terms through builder methods +/// .with::() +/// .without::() +/// .build(); +/// +/// // Consume the QueryState +/// let (entity, b) = query.single(&world); +///``` +pub struct QueryBuilder<'w, D: QueryData = (), F: QueryFilter = ()> { + access: FilteredAccess, + world: &'w mut World, + or: bool, + first: bool, + _marker: PhantomData<(D, F)>, +} + +impl<'w, D: QueryData, F: QueryFilter> QueryBuilder<'w, D, F> { + /// Creates a new builder with the accesses required for `Q` and `F` + pub fn new(world: &'w mut World) -> Self { + let fetch_state = D::init_state(world); + let filter_state = F::init_state(world); + + let mut access = FilteredAccess::default(); + D::update_component_access(&fetch_state, &mut access); + + // Use a temporary empty FilteredAccess for filters. This prevents them from conflicting with the + // main Query's `fetch_state` access. Filters are allowed to conflict with the main query fetch + // because they are evaluated *before* a specific reference is constructed. + let mut filter_access = FilteredAccess::default(); + F::update_component_access(&filter_state, &mut filter_access); + + // Merge the temporary filter access with the main access. This ensures that filter access is + // properly considered in a global "cross-query" context (both within systems and across systems). + access.extend(&filter_access); + + Self { + access, + world, + or: false, + first: false, + _marker: PhantomData, + } + } + + /// Returns a reference to the world passed to [`Self::new`]. + pub fn world(&self) -> &World { + self.world + } + + /// Returns a mutable reference to the world passed to [`Self::new`]. + pub fn world_mut(&mut self) -> &mut World { + self.world + } + + /// Adds access to self's underlying [`FilteredAccess`] respecting [`Self::or`] and [`Self::and`] + pub fn extend_access(&mut self, mut access: FilteredAccess) { + if self.or { + if self.first { + access.required.clear(); + self.access.extend(&access); + self.first = false; + } else { + self.access.append_or(&access); + } + } else { + self.access.extend(&access); + } + } + + /// Adds accesses required for `T` to self. + pub fn data(&mut self) -> &mut Self { + let state = T::init_state(self.world); + let mut access = FilteredAccess::default(); + T::update_component_access(&state, &mut access); + self.extend_access(access); + self + } + + /// Adds filter from `T` to self. + pub fn filter(&mut self) -> &mut Self { + let state = T::init_state(self.world); + let mut access = FilteredAccess::default(); + T::update_component_access(&state, &mut access); + self.extend_access(access); + self + } + + /// Adds [`With`] to the [`FilteredAccess`] of self. + pub fn with(&mut self) -> &mut Self { + self.filter::>(); + self + } + + /// Adds [`With`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`]. + pub fn with_id(&mut self, id: ComponentId) -> &mut Self { + let mut access = FilteredAccess::default(); + access.and_with(id); + self.extend_access(access); + self + } + + /// Adds [`Without`] to the [`FilteredAccess`] of self. + pub fn without(&mut self) -> &mut Self { + self.filter::>(); + self + } + + /// Adds [`Without`] to the [`FilteredAccess`] of self from a runtime [`ComponentId`]. + pub fn without_id(&mut self, id: ComponentId) -> &mut Self { + let mut access = FilteredAccess::default(); + access.and_without(id); + self.extend_access(access); + self + } + + /// Adds `&T` to the [`FilteredAccess`] of self. + pub fn ref_id(&mut self, id: ComponentId) -> &mut Self { + self.with_id(id); + self.access.add_read(id); + self + } + + /// Adds `&mut T` to the [`FilteredAccess`] of self. + pub fn mut_id(&mut self, id: ComponentId) -> &mut Self { + self.with_id(id); + self.access.add_write(id); + self + } + + /// Takes a function over mutable access to a [`QueryBuilder`], calls that function + /// on an empty builder and then adds all accesses from that builder to self as optional. + pub fn optional(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self { + let mut builder = QueryBuilder::new(self.world); + f(&mut builder); + self.access.extend_access(builder.access()); + self + } + + /// Takes a function over mutable access to a [`QueryBuilder`], calls that function + /// on an empty builder and then adds all accesses from that builder to self. + /// + /// Primarily used when inside a [`Self::or`] closure to group several terms. + pub fn and(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self { + let mut builder = QueryBuilder::new(self.world); + f(&mut builder); + let access = builder.access().clone(); + self.extend_access(access); + self + } + + /// Takes a function over mutable access to a [`QueryBuilder`], calls that function + /// on an empty builder, all accesses added to that builder will become terms in an or expression. + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component)] + /// # struct A; + /// # + /// # #[derive(Component)] + /// # struct B; + /// # + /// # let mut world = World::new(); + /// # + /// QueryBuilder::::new(&mut world).or(|builder| { + /// builder.with::(); + /// builder.with::(); + /// }); + /// // is equivalent to + /// QueryBuilder::::new(&mut world).filter::, With)>>(); + /// ``` + pub fn or(&mut self, f: impl Fn(&mut QueryBuilder)) -> &mut Self { + let mut builder = QueryBuilder::new(self.world); + builder.or = true; + builder.first = true; + f(&mut builder); + self.access.extend(builder.access()); + self + } + + /// Returns a reference to the the [`FilteredAccess`] that will be provided to the built [`Query`]. + pub fn access(&self) -> &FilteredAccess { + &self.access + } + + /// Transmute the existing builder adding required accesses. + /// This will maintain all exisiting accesses. + /// + /// If including a filter type see [`Self::transmute_filtered`] + pub fn transmute(&mut self) -> &mut QueryBuilder<'w, NewD> { + self.transmute_filtered::() + } + + /// Transmute the existing builder adding required accesses. + /// This will maintain all existing accesses. + pub fn transmute_filtered( + &mut self, + ) -> &mut QueryBuilder<'w, NewD, NewF> { + let mut fetch_state = NewD::init_state(self.world); + let filter_state = NewF::init_state(self.world); + + NewD::set_access(&mut fetch_state, &self.access); + + let mut access = FilteredAccess::default(); + NewD::update_component_access(&fetch_state, &mut access); + NewF::update_component_access(&filter_state, &mut access); + + self.extend_access(access); + // SAFETY: + // - We have included all required acceses for NewQ and NewF + // - The layout of all QueryBuilder instances is the same + unsafe { std::mem::transmute(self) } + } + + /// Create a [`QueryState`] with the accesses of the builder. + /// + /// Takes `&mut self` to access the innner world reference while initializing + /// state for the new [`QueryState`] + pub fn build(&mut self) -> QueryState { + QueryState::::from_builder(self) + } +} + +#[cfg(test)] +mod tests { + use crate as bevy_ecs; + use crate::prelude::*; + use crate::world::FilteredEntityRef; + + use super::QueryBuilder; + + #[derive(Component, PartialEq, Debug)] + struct A(usize); + + #[derive(Component, PartialEq, Debug)] + struct B(usize); + + #[derive(Component, PartialEq, Debug)] + struct C(usize); + + #[test] + fn builder_with_without_static() { + let mut world = World::new(); + let entity_a = world.spawn((A(0), B(0))).id(); + let entity_b = world.spawn((A(0), C(0))).id(); + + let mut query_a = QueryBuilder::::new(&mut world) + .with::() + .without::() + .build(); + assert_eq!(entity_a, query_a.single(&world)); + + let mut query_b = QueryBuilder::::new(&mut world) + .with::() + .without::() + .build(); + assert_eq!(entity_b, query_b.single(&world)); + } + + #[test] + fn builder_with_without_dynamic() { + let mut world = World::new(); + let entity_a = world.spawn((A(0), B(0))).id(); + let entity_b = world.spawn((A(0), C(0))).id(); + let component_id_a = world.init_component::(); + let component_id_b = world.init_component::(); + let component_id_c = world.init_component::(); + + let mut query_a = QueryBuilder::::new(&mut world) + .with_id(component_id_a) + .without_id(component_id_c) + .build(); + assert_eq!(entity_a, query_a.single(&world)); + + let mut query_b = QueryBuilder::::new(&mut world) + .with_id(component_id_a) + .without_id(component_id_b) + .build(); + assert_eq!(entity_b, query_b.single(&world)); + } + + #[test] + fn builder_or() { + let mut world = World::new(); + world.spawn((A(0), B(0))); + world.spawn(B(0)); + world.spawn(C(0)); + + let mut query_a = QueryBuilder::::new(&mut world) + .or(|builder| { + builder.with::(); + builder.with::(); + }) + .build(); + assert_eq!(2, query_a.iter(&world).count()); + + let mut query_b = QueryBuilder::::new(&mut world) + .or(|builder| { + builder.with::(); + builder.without::(); + }) + .build(); + dbg!(&query_b.component_access); + assert_eq!(2, query_b.iter(&world).count()); + + let mut query_c = QueryBuilder::::new(&mut world) + .or(|builder| { + builder.with::(); + builder.with::(); + builder.with::(); + }) + .build(); + assert_eq!(3, query_c.iter(&world).count()); + } + + #[test] + fn builder_transmute() { + let mut world = World::new(); + world.spawn(A(0)); + world.spawn((A(1), B(0))); + let mut query = QueryBuilder::<()>::new(&mut world) + .with::() + .transmute::<&A>() + .build(); + + query.iter(&world).for_each(|a| assert_eq!(a.0, 1)); + } + + #[test] + fn builder_static_components() { + let mut world = World::new(); + let entity = world.spawn((A(0), B(1))).id(); + + let mut query = QueryBuilder::::new(&mut world) + .data::<&A>() + .data::<&B>() + .build(); + + let entity_ref = query.single(&world); + + assert_eq!(entity, entity_ref.id()); + + let a = entity_ref.get::().unwrap(); + let b = entity_ref.get::().unwrap(); + + assert_eq!(0, a.0); + assert_eq!(1, b.0); + } + + #[test] + fn builder_dynamic_components() { + let mut world = World::new(); + let entity = world.spawn((A(0), B(1))).id(); + let component_id_a = world.init_component::(); + let component_id_b = world.init_component::(); + + let mut query = QueryBuilder::::new(&mut world) + .ref_id(component_id_a) + .ref_id(component_id_b) + .build(); + + let entity_ref = query.single(&world); + + assert_eq!(entity, entity_ref.id()); + + let a = entity_ref.get_by_id(component_id_a).unwrap(); + let b = entity_ref.get_by_id(component_id_b).unwrap(); + + // SAFETY: We set these pointers to point to these components + unsafe { + assert_eq!(0, a.deref::().0); + assert_eq!(1, b.deref::().0); + } + } +} diff --git a/crates/bevy_ecs/src/query/fetch.rs b/crates/bevy_ecs/src/query/fetch.rs index ddf8b3346ecce..5e83637d33019 100644 --- a/crates/bevy_ecs/src/query/fetch.rs +++ b/crates/bevy_ecs/src/query/fetch.rs @@ -1,11 +1,14 @@ use crate::{ - archetype::{Archetype, ArchetypeComponentId}, + archetype::Archetype, change_detection::{Ticks, TicksMut}, component::{Component, ComponentId, ComponentStorage, StorageType, Tick}, entity::Entity, query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery}, storage::{ComponentSparseSet, Table, TableRow}, - world::{unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityRef, Mut, Ref, World}, + world::{ + unsafe_world_cell::UnsafeWorldCell, EntityMut, EntityRef, FilteredEntityMut, + FilteredEntityRef, Mut, Ref, World, + }, }; use bevy_ptr::{ThinSlicePtr, UnsafeCellDeref}; use bevy_utils::all_tuples; @@ -320,15 +323,12 @@ unsafe impl WorldQuery for Entity { fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess) {} - fn update_archetype_component_access( - _state: &Self::State, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(_world: &mut World) {} + fn get_state(_world: &World) -> Option<()> { + Some(()) + } + fn matches_component_set( _state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool, @@ -402,18 +402,12 @@ unsafe impl<'a> WorldQuery for EntityRef<'a> { access.read_all(); } - fn update_archetype_component_access( - _state: &Self::State, - archetype: &Archetype, - access: &mut Access, - ) { - for component_id in archetype.components() { - access.add_read(archetype.get_archetype_component_id(component_id).unwrap()); - } - } - fn init_state(_world: &mut World) {} + fn get_state(_world: &World) -> Option<()> { + Some(()) + } + fn matches_component_set( _state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool, @@ -484,18 +478,12 @@ unsafe impl<'a> WorldQuery for EntityMut<'a> { access.write_all(); } - fn update_archetype_component_access( - _state: &Self::State, - archetype: &Archetype, - access: &mut Access, - ) { - for component_id in archetype.components() { - access.add_write(archetype.get_archetype_component_id(component_id).unwrap()); - } - } - fn init_state(_world: &mut World) {} + fn get_state(_world: &World) -> Option<()> { + Some(()) + } + fn matches_component_set( _state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool, @@ -509,6 +497,218 @@ unsafe impl<'a> QueryData for EntityMut<'a> { type ReadOnly = EntityRef<'a>; } +/// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self` +unsafe impl<'a> WorldQuery for FilteredEntityRef<'a> { + type Fetch<'w> = (UnsafeWorldCell<'w>, Access); + type Item<'w> = FilteredEntityRef<'w>; + type State = FilteredAccess; + + fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> { + item + } + + const IS_DENSE: bool = false; + + unsafe fn init_fetch<'w>( + world: UnsafeWorldCell<'w>, + _state: &Self::State, + _last_run: Tick, + _this_run: Tick, + ) -> Self::Fetch<'w> { + let mut access = Access::default(); + access.read_all(); + (world, access) + } + + #[inline] + unsafe fn set_archetype<'w>( + fetch: &mut Self::Fetch<'w>, + state: &Self::State, + archetype: &'w Archetype, + _table: &Table, + ) { + let mut access = Access::default(); + state.access.reads().for_each(|id| { + if archetype.contains(id) { + access.add_read(id); + } + }); + fetch.1 = access; + } + + #[inline] + unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) { + let mut access = Access::default(); + state.access.reads().for_each(|id| { + if table.has_column(id) { + access.add_read(id); + } + }); + fetch.1 = access; + } + + #[inline] + fn set_access<'w>(state: &mut Self::State, access: &FilteredAccess) { + *state = access.clone(); + state.access_mut().clear_writes(); + } + + #[inline(always)] + unsafe fn fetch<'w>( + (world, access): &mut Self::Fetch<'w>, + entity: Entity, + _table_row: TableRow, + ) -> Self::Item<'w> { + // SAFETY: `fetch` must be called with an entity that exists in the world + let cell = world.get_entity(entity).debug_checked_unwrap(); + // SAFETY: mutable access to every component has been registered. + FilteredEntityRef::new(cell, access.clone()) + } + + fn update_component_access( + state: &Self::State, + filtered_access: &mut FilteredAccess, + ) { + assert!( + filtered_access.access().is_compatible(&state.access), + "FilteredEntityRef conflicts with a previous access in this query. Exclusive access cannot coincide with any other accesses.", + ); + filtered_access.access.extend(&state.access); + } + + fn init_state(_world: &mut World) -> Self::State { + FilteredAccess::default() + } + + fn get_state(_world: &World) -> Option { + Some(FilteredAccess::default()) + } + + fn matches_component_set( + _state: &Self::State, + _set_contains_id: &impl Fn(ComponentId) -> bool, + ) -> bool { + true + } +} + +/// SAFETY: `Self` is the same as `Self::ReadOnly` +unsafe impl<'a> QueryData for FilteredEntityRef<'a> { + type ReadOnly = Self; +} + +/// SAFETY: Access is read-only. +unsafe impl ReadOnlyQueryData for FilteredEntityRef<'_> {} + +/// SAFETY: The accesses of `Self::ReadOnly` are a subset of the accesses of `Self` +unsafe impl<'a> WorldQuery for FilteredEntityMut<'a> { + type Fetch<'w> = (UnsafeWorldCell<'w>, Access); + type Item<'w> = FilteredEntityMut<'w>; + type State = FilteredAccess; + + fn shrink<'wlong: 'wshort, 'wshort>(item: Self::Item<'wlong>) -> Self::Item<'wshort> { + item + } + + const IS_DENSE: bool = false; + + unsafe fn init_fetch<'w>( + world: UnsafeWorldCell<'w>, + _state: &Self::State, + _last_run: Tick, + _this_run: Tick, + ) -> Self::Fetch<'w> { + let mut access = Access::default(); + access.write_all(); + (world, access) + } + + #[inline] + unsafe fn set_archetype<'w>( + fetch: &mut Self::Fetch<'w>, + state: &Self::State, + archetype: &'w Archetype, + _table: &Table, + ) { + let mut access = Access::default(); + state.access.reads().for_each(|id| { + if archetype.contains(id) { + access.add_read(id); + } + }); + state.access.writes().for_each(|id| { + if archetype.contains(id) { + access.add_write(id); + } + }); + fetch.1 = access; + } + + #[inline] + unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table) { + let mut access = Access::default(); + state.access.reads().for_each(|id| { + if table.has_column(id) { + access.add_read(id); + } + }); + state.access.writes().for_each(|id| { + if table.has_column(id) { + access.add_write(id); + } + }); + fetch.1 = access; + } + + #[inline] + fn set_access<'w>(state: &mut Self::State, access: &FilteredAccess) { + *state = access.clone(); + } + + #[inline(always)] + unsafe fn fetch<'w>( + (world, access): &mut Self::Fetch<'w>, + entity: Entity, + _table_row: TableRow, + ) -> Self::Item<'w> { + // SAFETY: `fetch` must be called with an entity that exists in the world + let cell = world.get_entity(entity).debug_checked_unwrap(); + // SAFETY: mutable access to every component has been registered. + FilteredEntityMut::new(cell, access.clone()) + } + + fn update_component_access( + state: &Self::State, + filtered_access: &mut FilteredAccess, + ) { + assert!( + filtered_access.access().is_compatible(&state.access), + "FilteredEntityMut conflicts with a previous access in this query. Exclusive access cannot coincide with any other accesses.", + ); + filtered_access.access.extend(&state.access); + } + + fn init_state(_world: &mut World) -> Self::State { + FilteredAccess::default() + } + + fn get_state(_world: &World) -> Option { + Some(FilteredAccess::default()) + } + + fn matches_component_set( + _state: &Self::State, + _set_contains_id: &impl Fn(ComponentId) -> bool, + ) -> bool { + true + } +} + +/// SAFETY: access of `FilteredEntityRef` is a subset of `FilteredEntityMut` +unsafe impl<'a> QueryData for FilteredEntityMut<'a> { + type ReadOnly = FilteredEntityRef<'a>; +} + #[doc(hidden)] pub struct ReadFetch<'w, T> { // T::Storage = TableStorage @@ -628,20 +828,14 @@ unsafe impl WorldQuery for &T { access.add_read(component_id); } - fn update_archetype_component_access( - &component_id: &ComponentId, - archetype: &Archetype, - access: &mut Access, - ) { - if let Some(archetype_component_id) = archetype.get_archetype_component_id(component_id) { - access.add_read(archetype_component_id); - } - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &state: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -795,20 +989,14 @@ unsafe impl<'__w, T: Component> WorldQuery for Ref<'__w, T> { access.add_read(component_id); } - fn update_archetype_component_access( - &component_id: &ComponentId, - archetype: &Archetype, - access: &mut Access, - ) { - if let Some(archetype_component_id) = archetype.get_archetype_component_id(component_id) { - access.add_read(archetype_component_id); - } - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &state: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -962,20 +1150,14 @@ unsafe impl<'__w, T: Component> WorldQuery for &'__w mut T { access.add_write(component_id); } - fn update_archetype_component_access( - &component_id: &ComponentId, - archetype: &Archetype, - access: &mut Access, - ) { - if let Some(archetype_component_id) = archetype.get_archetype_component_id(component_id) { - access.add_write(archetype_component_id); - } - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &state: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -1079,20 +1261,14 @@ unsafe impl WorldQuery for Option { access.extend_access(&intermediate); } - fn update_archetype_component_access( - state: &T::State, - archetype: &Archetype, - access: &mut Access, - ) { - if T::matches_component_set(state, &|id| archetype.contains(id)) { - T::update_archetype_component_access(state, archetype, access); - } - } - fn init_state(world: &mut World) -> T::State { T::init_state(world) } + fn get_state(world: &World) -> Option { + T::get_state(world) + } + fn matches_component_set( _state: &T::State, _set_contains_id: &impl Fn(ComponentId) -> bool, @@ -1231,17 +1407,14 @@ unsafe impl WorldQuery for Has { // Do nothing as presence of `Has` never affects whether two queries are disjoint } - fn update_archetype_component_access( - _state: &Self::State, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( _state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool, @@ -1368,6 +1541,7 @@ macro_rules! impl_anytuple_fetch { _new_access.extend_access(&intermediate); } else { $name::update_component_access($name, &mut _new_access); + _new_access.required = _access.required.clone(); _not_first = true; } )* @@ -1375,19 +1549,14 @@ macro_rules! impl_anytuple_fetch { *_access = _new_access; } - fn update_archetype_component_access(state: &Self::State, _archetype: &Archetype, _access: &mut Access) { - let ($($name,)*) = state; - $( - if $name::matches_component_set($name, &|id| _archetype.contains(id)) { - $name::update_archetype_component_access($name, _archetype, _access); - } - )* - } - fn init_state(_world: &mut World) -> Self::State { ($($name::init_state(_world),)*) } + fn get_state(_world: &World) -> Option { + Some(($($name::get_state(_world)?,)*)) + } + fn matches_component_set(_state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { let ($($name,)*) = _state; false $(|| $name::matches_component_set($name, _set_contains_id))* @@ -1457,17 +1626,14 @@ unsafe impl WorldQuery for NopWorldQuery { fn update_component_access(_state: &D::State, _access: &mut FilteredAccess) {} - fn update_archetype_component_access( - _state: &D::State, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(world: &mut World) -> Self::State { D::init_state(world) } + fn get_state(world: &World) -> Option { + D::get_state(world) + } + fn matches_component_set( state: &Self::State, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -1527,15 +1693,12 @@ unsafe impl WorldQuery for PhantomData { fn update_component_access(_state: &Self::State, _access: &mut FilteredAccess) {} - fn update_archetype_component_access( - _state: &Self::State, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(_world: &mut World) -> Self::State {} + fn get_state(_world: &World) -> Option { + Some(()) + } + fn matches_component_set( _state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool, diff --git a/crates/bevy_ecs/src/query/filter.rs b/crates/bevy_ecs/src/query/filter.rs index 5853b1fa6bb72..a8d225ab81173 100644 --- a/crates/bevy_ecs/src/query/filter.rs +++ b/crates/bevy_ecs/src/query/filter.rs @@ -1,8 +1,8 @@ use crate::{ - archetype::{Archetype, ArchetypeComponentId}, + archetype::Archetype, component::{Component, ComponentId, ComponentStorage, StorageType, Tick}, entity::Entity, - query::{Access, DebugCheckedUnwrap, FilteredAccess, WorldQuery}, + query::{DebugCheckedUnwrap, FilteredAccess, WorldQuery}, storage::{Column, ComponentSparseSet, Table, TableRow}, world::{unsafe_world_cell::UnsafeWorldCell, World}, }; @@ -175,18 +175,14 @@ unsafe impl WorldQuery for With { access.and_with(id); } - #[inline] - fn update_archetype_component_access( - _state: &ComponentId, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &id: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -287,18 +283,14 @@ unsafe impl WorldQuery for Without { access.and_without(id); } - #[inline] - fn update_archetype_component_access( - _state: &ComponentId, - _archetype: &Archetype, - _access: &mut Access, - ) { - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &id: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -449,6 +441,7 @@ macro_rules! impl_query_filter_tuple { _new_access.extend_access(&intermediate); } else { $filter::update_component_access($filter, &mut _new_access); + _new_access.required = access.required.clone(); _not_first = true; } )* @@ -456,15 +449,14 @@ macro_rules! impl_query_filter_tuple { *access = _new_access; } - fn update_archetype_component_access(state: &Self::State, archetype: &Archetype, access: &mut Access) { - let ($($filter,)*) = state; - $($filter::update_archetype_component_access($filter, archetype, access);)* - } - fn init_state(world: &mut World) -> Self::State { ($($filter::init_state(world),)*) } + fn get_state(world: &World) -> Option { + Some(($($filter::get_state(world)?,)*)) + } + fn matches_component_set(_state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { let ($($filter,)*) = _state; false $(|| $filter::matches_component_set($filter, _set_contains_id))* @@ -677,21 +669,14 @@ unsafe impl WorldQuery for Added { access.add_read(id); } - #[inline] - fn update_archetype_component_access( - &id: &ComponentId, - archetype: &Archetype, - access: &mut Access, - ) { - if let Some(archetype_component_id) = archetype.get_archetype_component_id(id) { - access.add_read(archetype_component_id); - } - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &id: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, @@ -884,21 +869,14 @@ unsafe impl WorldQuery for Changed { access.add_read(id); } - #[inline] - fn update_archetype_component_access( - &id: &ComponentId, - archetype: &Archetype, - access: &mut Access, - ) { - if let Some(archetype_component_id) = archetype.get_archetype_component_id(id) { - access.add_read(archetype_component_id); - } - } - fn init_state(world: &mut World) -> ComponentId { world.init_component::() } + fn get_state(world: &World) -> Option { + world.component_id::() + } + fn matches_component_set( &id: &ComponentId, set_contains_id: &impl Fn(ComponentId) -> bool, diff --git a/crates/bevy_ecs/src/query/mod.rs b/crates/bevy_ecs/src/query/mod.rs index 3beaf664ac33b..618c179f98570 100644 --- a/crates/bevy_ecs/src/query/mod.rs +++ b/crates/bevy_ecs/src/query/mod.rs @@ -1,6 +1,7 @@ //! Contains APIs for retrieving component data from the world. mod access; +mod builder; mod error; mod fetch; mod filter; @@ -11,6 +12,7 @@ mod world_query; pub use access::*; pub use bevy_ecs_macros::{QueryData, QueryFilter}; +pub use builder::*; pub use error::*; pub use fetch::*; pub use filter::*; diff --git a/crates/bevy_ecs/src/query/state.rs b/crates/bevy_ecs/src/query/state.rs index e459baf7a1a75..58f5db9a1c1a8 100644 --- a/crates/bevy_ecs/src/query/state.rs +++ b/crates/bevy_ecs/src/query/state.rs @@ -8,7 +8,7 @@ use crate::{ Access, BatchingStrategy, DebugCheckedUnwrap, FilteredAccess, QueryCombinationIter, QueryIter, QueryParIter, }, - storage::TableId, + storage::{SparseSetIndex, TableId}, world::{unsafe_world_cell::UnsafeWorldCell, World, WorldId}, }; #[cfg(feature = "trace")] @@ -17,8 +17,8 @@ use fixedbitset::FixedBitSet; use std::{any::TypeId, borrow::Borrow, fmt, mem::MaybeUninit}; use super::{ - NopWorldQuery, QueryComponentError, QueryData, QueryEntityError, QueryFilter, QueryManyIter, - QuerySingleError, ROQueryItem, + NopWorldQuery, QueryBuilder, QueryComponentError, QueryData, QueryEntityError, QueryFilter, + QueryManyIter, QuerySingleError, ROQueryItem, }; /// Provides scoped access to a [`World`] state according to a given [`QueryData`] and [`QueryFilter`]. @@ -138,6 +138,34 @@ impl QueryState { state } + /// Creates a new [`QueryState`] from a given [`QueryBuilder`] and inherits it's [`FilteredAccess`]. + pub fn from_builder(builder: &mut QueryBuilder) -> Self { + let mut fetch_state = D::init_state(builder.world_mut()); + let filter_state = F::init_state(builder.world_mut()); + D::set_access(&mut fetch_state, builder.access()); + + let mut state = Self { + world_id: builder.world().id(), + archetype_generation: ArchetypeGeneration::initial(), + matched_table_ids: Vec::new(), + matched_archetype_ids: Vec::new(), + fetch_state, + filter_state, + component_access: builder.access().clone(), + matched_tables: Default::default(), + matched_archetypes: Default::default(), + archetype_component_access: Default::default(), + #[cfg(feature = "trace")] + par_iter_span: bevy_utils::tracing::info_span!( + "par_for_each", + data = std::any::type_name::(), + filter = std::any::type_name::(), + ), + }; + state.update_archetypes(builder.world()); + state + } + /// Checks if the query is empty for the given [`World`], where the last change and current tick are given. /// /// # Panics @@ -250,17 +278,10 @@ impl QueryState { pub fn new_archetype(&mut self, archetype: &Archetype) { if D::matches_component_set(&self.fetch_state, &|id| archetype.contains(id)) && F::matches_component_set(&self.filter_state, &|id| archetype.contains(id)) + && self.matches_component_set(&|id| archetype.contains(id)) { - D::update_archetype_component_access( - &self.fetch_state, - archetype, - &mut self.archetype_component_access, - ); - F::update_archetype_component_access( - &self.filter_state, - archetype, - &mut self.archetype_component_access, - ); + self.update_archetype_component_access(archetype); + let archetype_index = archetype.id().index(); if !self.matched_archetypes.contains(archetype_index) { self.matched_archetypes.grow(archetype_index + 1); @@ -276,6 +297,86 @@ impl QueryState { } } + /// Returns `true` if this query matches a set of components. Otherwise, returns `false`. + pub fn matches_component_set(&self, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { + self.component_access.filter_sets.iter().any(|set| { + set.with + .ones() + .all(|index| set_contains_id(ComponentId::get_sparse_set_index(index))) + && set + .without + .ones() + .all(|index| !set_contains_id(ComponentId::get_sparse_set_index(index))) + }) + } + + /// For the given `archetype`, adds any component accessed used by this query's underlying [`FilteredAccess`] to `access`. + pub fn update_archetype_component_access(&mut self, archetype: &Archetype) { + self.component_access.access.reads().for_each(|id| { + if let Some(id) = archetype.get_archetype_component_id(id) { + self.archetype_component_access.add_read(id); + } + }); + self.component_access.access.writes().for_each(|id| { + if let Some(id) = archetype.get_archetype_component_id(id) { + self.archetype_component_access.add_write(id); + } + }); + } + + /// Use this to transform a [`QueryState`] into a more generic [`QueryState`]. + /// This can be useful for passing to another function that might take the more general form. + /// See [`Query::transmute_lens`](crate::system::Query::transmute_lens) for more details. + /// + /// You should not call [`update_archetypes`](Self::update_archetypes) on the returned [`QueryState`] as the result will be unpredictable. + /// You might end up with a mix of archetypes that only matched the original query + archetypes that only match + /// the new [`QueryState`]. Most of the safe methods on [`QueryState`] call [`QueryState::update_archetypes`] internally, so this + /// best used through a [`Query`](crate::system::Query). + pub fn transmute(&self, world: &World) -> QueryState { + self.transmute_filtered::(world) + } + + /// Creates a new [`QueryState`] with the same underlying [`FilteredAccess`], matched tables and archetypes + /// as self but with a new type signature. + /// + /// Panics if `NewD` or `NewF` require accesses that this query does not have. + pub fn transmute_filtered( + &self, + world: &World, + ) -> QueryState { + let mut component_access = FilteredAccess::default(); + let mut fetch_state = NewD::get_state(world).expect("Could not create fetch_state, Please initialize all referenced components before transmuting."); + let filter_state = NewF::get_state(world).expect("Could not create filter_state, Please initialize all referenced components before transmuting."); + + NewD::set_access(&mut fetch_state, &self.component_access); + NewD::update_component_access(&fetch_state, &mut component_access); + + let mut filter_component_access = FilteredAccess::default(); + NewF::update_component_access(&filter_state, &mut filter_component_access); + + component_access.extend(&filter_component_access); + assert!(component_access.is_subset(&self.component_access), "Transmuted state for {} attempts to access terms that are not allowed by original state {}.", std::any::type_name::<(NewD, NewF)>(), std::any::type_name::<(D, F)>() ); + + QueryState { + world_id: self.world_id, + archetype_generation: self.archetype_generation, + matched_table_ids: self.matched_table_ids.clone(), + matched_archetype_ids: self.matched_archetype_ids.clone(), + fetch_state, + filter_state, + component_access: self.component_access.clone(), + matched_tables: self.matched_tables.clone(), + matched_archetypes: self.matched_archetypes.clone(), + archetype_component_access: self.archetype_component_access.clone(), + #[cfg(feature = "trace")] + par_iter_span: bevy_utils::tracing::info_span!( + "par_for_each", + query = std::any::type_name::(), + filter = std::any::type_name::(), + ), + } + } + /// Gets the query result for the given [`World`] and [`Entity`]. /// /// This can only be called for read-only queries, see [`Self::get_mut`] for write-queries. @@ -1305,9 +1406,17 @@ impl QueryState { } } +impl From> for QueryState { + fn from(mut value: QueryBuilder) -> Self { + QueryState::from_builder(&mut value) + } +} + #[cfg(test)] mod tests { - use crate::{prelude::*, query::QueryEntityError}; + use crate as bevy_ecs; + use crate::world::FilteredEntityRef; + use crate::{component::Component, prelude::*, query::QueryEntityError}; #[test] fn get_many_unchecked_manual_uniqueness() { @@ -1412,4 +1521,201 @@ mod tests { let mut query_state = world_1.query::(); let _panics = query_state.get_many_mut(&mut world_2, []); } + + #[derive(Component, PartialEq, Debug)] + struct A(usize); + + #[derive(Component, PartialEq, Debug)] + struct B(usize); + + #[derive(Component, PartialEq, Debug)] + struct C(usize); + + #[test] + fn can_transmute_to_more_general() { + let mut world = World::new(); + world.spawn((A(1), B(0))); + + let query_state = world.query::<(&A, &B)>(); + let mut new_query_state = query_state.transmute::<&A>(&world); + assert_eq!(new_query_state.iter(&world).len(), 1); + let a = new_query_state.single(&world); + + assert_eq!(a.0, 1); + } + + #[test] + fn cannot_get_data_not_in_original_query() { + let mut world = World::new(); + world.spawn((A(0), B(0))); + world.spawn((A(1), B(0), C(0))); + + let query_state = world.query_filtered::<(&A, &B), Without>(); + let mut new_query_state = query_state.transmute::<&A>(&world); + // even though we change the query to not have Without, we do not get the component with C. + let a = new_query_state.single(&world); + + assert_eq!(a.0, 0); + } + + #[test] + fn can_transmute_empty_tuple() { + let mut world = World::new(); + world.init_component::(); + let entity = world.spawn(A(10)).id(); + + let q = world.query::<()>(); + let mut q = q.transmute::(&world); + assert_eq!(q.single(&world), entity); + } + + #[test] + fn can_transmute_immut_fetch() { + let mut world = World::new(); + world.spawn(A(10)); + + let q = world.query::<&A>(); + let mut new_q = q.transmute::>(&world); + assert!(new_q.single(&world).is_added()); + + let q = world.query::>(); + let _ = q.transmute::<&A>(&world); + } + + #[test] + fn can_transmute_mut_fetch() { + let mut world = World::new(); + world.spawn(A(0)); + + let q = world.query::<&mut A>(); + let _ = q.transmute::>(&world); + let _ = q.transmute::<&A>(&world); + } + + #[test] + fn can_transmute_entity_mut() { + let mut world = World::new(); + world.spawn(A(0)); + + let q: QueryState> = world.query::(); + let _ = q.transmute::(&world); + } + + #[test] + fn can_generalize_with_option() { + let mut world = World::new(); + world.spawn((A(0), B(0))); + + let query_state = world.query::<(Option<&A>, &B)>(); + let _ = query_state.transmute::>(&world); + let _ = query_state.transmute::<&B>(&world); + } + + #[test] + #[should_panic( + expected = "Transmuted state for ((&bevy_ecs::query::state::tests::A, &bevy_ecs::query::state::tests::B), ()) attempts to access terms that are not allowed by original state (&bevy_ecs::query::state::tests::A, ())." + )] + fn cannot_transmute_to_include_data_not_in_original_query() { + let mut world = World::new(); + world.init_component::(); + world.init_component::(); + world.spawn(A(0)); + + let query_state = world.query::<&A>(); + let mut _new_query_state = query_state.transmute::<(&A, &B)>(&world); + } + + #[test] + #[should_panic( + expected = "Transmuted state for (&mut bevy_ecs::query::state::tests::A, ()) attempts to access terms that are not allowed by original state (&bevy_ecs::query::state::tests::A, ())." + )] + fn cannot_transmute_immut_to_mut() { + let mut world = World::new(); + world.spawn(A(0)); + + let query_state = world.query::<&A>(); + let mut _new_query_state = query_state.transmute::<&mut A>(&world); + } + + #[test] + #[should_panic( + expected = "Transmuted state for (&bevy_ecs::query::state::tests::A, ()) attempts to access terms that are not allowed by original state (core::option::Option<&bevy_ecs::query::state::tests::A>, ())." + )] + fn cannot_transmute_option_to_immut() { + let mut world = World::new(); + world.spawn(C(0)); + + let query_state = world.query::>(); + let mut new_query_state = query_state.transmute::<&A>(&world); + let x = new_query_state.single(&world); + assert_eq!(x.0, 1234); + } + + #[test] + #[should_panic( + expected = "Transmuted state for (&bevy_ecs::query::state::tests::A, ()) attempts to access terms that are not allowed by original state (bevy_ecs::world::entity_ref::EntityRef, ())." + )] + fn cannot_transmute_entity_ref() { + let mut world = World::new(); + world.init_component::(); + + let q = world.query::(); + let _ = q.transmute::<&A>(&world); + } + + #[test] + fn can_transmute_filtered_entity() { + let mut world = World::new(); + let entity = world.spawn((A(0), B(1))).id(); + let query = + QueryState::<(Entity, &A, &B)>::new(&mut world).transmute::(&world); + + let mut query = query; + // Our result is completely untyped + let entity_ref = query.single(&world); + + assert_eq!(entity, entity_ref.id()); + assert_eq!(0, entity_ref.get::().unwrap().0); + assert_eq!(1, entity_ref.get::().unwrap().0); + } + + #[test] + fn can_transmute_added() { + let mut world = World::new(); + let entity_a = world.spawn(A(0)).id(); + + let mut query = QueryState::<(Entity, &A, Has)>::new(&mut world) + .transmute_filtered::<(Entity, Has), Added>(&world); + + assert_eq!((entity_a, false), query.single(&world)); + + world.clear_trackers(); + + let entity_b = world.spawn((A(0), B(0))).id(); + assert_eq!((entity_b, true), query.single(&world)); + + world.clear_trackers(); + + assert!(query.get_single(&world).is_err()); + } + + #[test] + fn can_transmute_changed() { + let mut world = World::new(); + let entity_a = world.spawn(A(0)).id(); + + let mut detection_query = QueryState::<(Entity, &A)>::new(&mut world) + .transmute_filtered::>(&world); + + let mut change_query = QueryState::<&mut A>::new(&mut world); + assert_eq!(entity_a, detection_query.single(&world)); + + world.clear_trackers(); + + assert!(detection_query.get_single(&world).is_err()); + + change_query.single_mut(&mut world).0 = 1; + + assert_eq!(entity_a, detection_query.single(&world)); + } } diff --git a/crates/bevy_ecs/src/query/world_query.rs b/crates/bevy_ecs/src/query/world_query.rs index 758a3e56b0fe3..ad09c5fe447bf 100644 --- a/crates/bevy_ecs/src/query/world_query.rs +++ b/crates/bevy_ecs/src/query/world_query.rs @@ -1,8 +1,8 @@ use crate::{ - archetype::{Archetype, ArchetypeComponentId}, + archetype::Archetype, component::{ComponentId, Tick}, entity::Entity, - query::{Access, FilteredAccess}, + query::FilteredAccess, storage::{Table, TableRow}, world::{unsafe_world_cell::UnsafeWorldCell, World}, }; @@ -14,13 +14,11 @@ use bevy_utils::all_tuples; /// # Safety /// /// Implementor must ensure that -/// [`update_component_access`], [`update_archetype_component_access`], [`matches_component_set`], and [`fetch`] +/// [`update_component_access`], [`matches_component_set`], and [`fetch`] /// obey the following: /// /// - For each component mutably accessed by [`fetch`], [`update_component_access`] should add write access unless read or write access has already been added, in which case it should panic. /// - For each component readonly accessed by [`fetch`], [`update_component_access`] should add read access unless write access has already been added, in which case it should panic. -/// - For each component mutably accessed by [`fetch`], [`update_archetype_component_access`] should add write access if that component belongs to the archetype. -/// - For each component readonly accessed by [`fetch`], [`update_archetype_component_access`] should add read access if that component belongs to the archetype. /// - If `fetch` mutably accesses the same component twice, [`update_component_access`] should panic. /// - [`update_component_access`] may not add a `Without` filter for a component unless [`matches_component_set`] always returns `false` when the component set contains that component. /// - [`update_component_access`] may not add a `With` filter for a component unless [`matches_component_set`] always returns `false` when the component set doesn't contain that component. @@ -34,7 +32,6 @@ use bevy_utils::all_tuples; /// [`fetch`]: Self::fetch /// [`matches_component_set`]: Self::matches_component_set /// [`Query`]: crate::system::Query -/// [`update_archetype_component_access`]: Self::update_archetype_component_access /// [`update_component_access`]: Self::update_component_access /// [`QueryData`]: crate::query::QueryData /// [`QueryFilter`]: crate::query::QueryFilter @@ -84,7 +81,6 @@ pub unsafe trait WorldQuery { /// # Safety /// /// - `archetype` and `tables` must be from the same [`World`] that [`WorldQuery::init_state`] was called on. - /// - [`Self::update_archetype_component_access`] must have been previously called with `archetype`. /// - `table` must correspond to `archetype`. /// - `state` must be the [`State`](Self::State) that `fetch` was initialized with. unsafe fn set_archetype<'w>( @@ -100,11 +96,15 @@ pub unsafe trait WorldQuery { /// # Safety /// /// - `table` must be from the same [`World`] that [`WorldQuery::init_state`] was called on. - /// - `table` must belong to an archetype that was previously registered with - /// [`Self::update_archetype_component_access`]. /// - `state` must be the [`State`](Self::State) that `fetch` was initialized with. unsafe fn set_table<'w>(fetch: &mut Self::Fetch<'w>, state: &Self::State, table: &'w Table); + /// Sets available accesses for implementors with dynamic access such as [`FilteredEntityRef`](crate::world::FilteredEntityRef) + /// or [`FilteredEntityMut`](crate::world::FilteredEntityMut). + /// + /// Called when constructing a [`QueryLens`](crate::system::QueryLens) or calling [`QueryState::from_builder`](super::QueryState::from_builder) + fn set_access(_state: &mut Self::State, _access: &FilteredAccess) {} + /// Fetch [`Self::Item`](`WorldQuery::Item`) for either the given `entity` in the current [`Table`], /// or for the given `entity` in the current [`Archetype`]. This must always be called after /// [`WorldQuery::set_table`] with a `table_row` in the range of the current [`Table`] or after @@ -125,18 +125,12 @@ pub unsafe trait WorldQuery { // and forgetting to do so would be unsound. fn update_component_access(state: &Self::State, access: &mut FilteredAccess); - /// For the given `archetype`, adds any component accessed used by this [`WorldQuery`] to `access`. - // This does not have a default body of `{}` because 99% of cases need to add accesses - // and forgetting to do so would be unsound. - fn update_archetype_component_access( - state: &Self::State, - archetype: &Archetype, - access: &mut Access, - ); - /// Creates and initializes a [`State`](WorldQuery::State) for this [`WorldQuery`] type. fn init_state(world: &mut World) -> Self::State; + /// Attempts to initializes a [`State`](WorldQuery::State) for this [`WorldQuery`] type. + fn get_state(world: &World) -> Option; + /// Returns `true` if this query matches a set of components. Otherwise, returns `false`. fn matches_component_set( state: &Self::State, @@ -210,15 +204,14 @@ macro_rules! impl_tuple_world_query { $($name::update_component_access($name, _access);)* } - fn update_archetype_component_access(state: &Self::State, _archetype: &Archetype, _access: &mut Access) { - let ($($name,)*) = state; - $($name::update_archetype_component_access($name, _archetype, _access);)* - } - fn init_state(_world: &mut World) -> Self::State { ($($name::init_state(_world),)*) } + fn get_state(_world: &World) -> Option { + Some(($($name::get_state(_world)?,)*)) + } + fn matches_component_set(state: &Self::State, _set_contains_id: &impl Fn(ComponentId) -> bool) -> bool { let ($($name,)*) = state; true $(&& $name::matches_component_set($name, _set_contains_id))* diff --git a/crates/bevy_ecs/src/system/query.rs b/crates/bevy_ecs/src/system/query.rs index 1d34d8d30a3d5..9f0eaf34d4a50 100644 --- a/crates/bevy_ecs/src/system/query.rs +++ b/crates/bevy_ecs/src/system/query.rs @@ -1431,6 +1431,96 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> { .is_ok() } } + + /// Returns a [`QueryLens`] that can be used to get a query with a more general fetch. + /// + /// For example, this can transform a `Query<(&A, &mut B)>` to a `Query<&B>`. + /// This can be useful for passing the query to another function. Note that since + /// filter terms are dropped, non-archetypal filters like [`Added`](crate::query::Added) and + /// [`Changed`](crate::query::Changed) will not be respected. To maintain or change filter + /// terms see [`Self::transmute_lens_filtered`] + /// + /// ## Panics + /// + /// This will panic if `NewD` is not a subset of the original fetch `Q` + /// + /// ## Example + /// + /// ```rust + /// # use bevy_ecs::prelude::*; + /// # use bevy_ecs::system::QueryLens; + /// # + /// # #[derive(Component)] + /// # struct A(usize); + /// # + /// # #[derive(Component)] + /// # struct B(usize); + /// # + /// # let mut world = World::new(); + /// # + /// # world.spawn((A(10), B(5))); + /// # + /// fn reusable_function(lens: &mut QueryLens<&A>) { + /// assert_eq!(lens.query().single().0, 10); + /// } + /// + /// // We can use the function in a system that takes the exact query. + /// fn system_1(mut query: Query<&A>) { + /// reusable_function(&mut query.as_query_lens()); + /// } + /// + /// // We can also use it with a query that does not match exactly + /// // by transmuting it. + /// fn system_2(mut query: Query<(&mut A, &B)>) { + /// let mut lens = query.transmute_lens::<&A>(); + /// reusable_function(&mut lens); + /// } + /// + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems((system_1, system_2)); + /// # schedule.run(&mut world); + /// ``` + /// + /// ## Allowed Transmutes + /// + /// Besides removing parameters from the query, you can also + /// make limited changes to the types of paramters. + /// + /// * Can always add/remove `Entity` + /// * `Ref` <-> `&T` + /// * `&mut T` -> `&T` + /// * `&mut T` -> `Ref` + /// * [`EntityMut`](crate::world::EntityMut) -> [`EntityRef`](crate::world::EntityRef) + /// + pub fn transmute_lens(&mut self) -> QueryLens<'_, NewD> { + self.transmute_lens_filtered::() + } + + /// Equivalent to [`Self::transmute_lens`] but also includes a [`QueryFilter`] type. + /// + /// Note that the lens will iterate the same tables and archetypes as the original query. This means that + /// additional archetypal query terms like [`With`](crate::query::With) and [`Without`](crate::query::Without) + /// will not necessarily be respected and non-archetypal terms like [`Added`](crate::query::Added) and + /// [`Changed`](crate::query::Changed) will only be respected if they are in the type signature. + pub fn transmute_lens_filtered( + &mut self, + ) -> QueryLens<'_, NewD, NewF> { + // SAFETY: There are no other active borrows of data from world + let world = unsafe { self.world.world() }; + let state = self.state.transmute_filtered::(world); + QueryLens { + world: self.world, + state, + last_run: self.last_run, + this_run: self.this_run, + force_read_only_component_access: self.force_read_only_component_access, + } + } + + /// Gets a [`QueryLens`] with the same accesses as the existing query + pub fn as_query_lens(&mut self) -> QueryLens<'_, D> { + self.transmute_lens() + } } impl<'w, 's, D: QueryData, F: QueryFilter> IntoIterator for &'w Query<'_, 's, D, F> { @@ -1532,3 +1622,43 @@ impl<'w, 's, D: ReadOnlyQueryData, F: QueryFilter> Query<'w, 's, D, F> { } } } + +/// Type returned from [`Query::transmute_lens`] containing the new [`QueryState`]. +/// +/// Call [`query`](QueryLens::query) or [`into`](Into::into) to construct the resulting [`Query`] +pub struct QueryLens<'w, Q: QueryData, F: QueryFilter = ()> { + world: UnsafeWorldCell<'w>, + state: QueryState, + last_run: Tick, + this_run: Tick, + force_read_only_component_access: bool, +} + +impl<'w, Q: QueryData, F: QueryFilter> QueryLens<'w, Q, F> { + /// Create a [`Query`] from the underlying [`QueryState`]. + pub fn query(&mut self) -> Query<'w, '_, Q, F> { + Query { + world: self.world, + state: &self.state, + last_run: self.last_run, + this_run: self.this_run, + force_read_only_component_access: self.force_read_only_component_access, + } + } +} + +impl<'w, 's, Q: QueryData, F: QueryFilter> From<&'s mut QueryLens<'w, Q, F>> + for Query<'w, 's, Q, F> +{ + fn from(value: &'s mut QueryLens<'w, Q, F>) -> Query<'w, 's, Q, F> { + value.query() + } +} + +impl<'w, 'q, Q: QueryData, F: QueryFilter> From<&'q mut Query<'w, '_, Q, F>> + for QueryLens<'q, Q, F> +{ + fn from(value: &'q mut Query<'w, '_, Q, F>) -> QueryLens<'q, Q, F> { + value.transmute_lens_filtered() + } +} diff --git a/crates/bevy_ecs/src/world/entity_ref.rs b/crates/bevy_ecs/src/world/entity_ref.rs index 2641132c651cb..9558ff5c4e73e 100644 --- a/crates/bevy_ecs/src/world/entity_ref.rs +++ b/crates/bevy_ecs/src/world/entity_ref.rs @@ -4,6 +4,7 @@ use crate::{ change_detection::MutUntyped, component::{Component, ComponentId, ComponentTicks, Components, StorageType}, entity::{Entities, Entity, EntityLocation}, + query::{Access, DebugCheckedUnwrap}, removal_detection::RemovedComponentEvents, storage::Storages, world::{Mut, World}, @@ -1495,6 +1496,358 @@ impl<'w, 'a, T: Component> VacantEntry<'w, 'a, T> { } } +/// Provides read-only access to a single entity and some of its components defined by the contained [`Access`]. +#[derive(Clone)] +pub struct FilteredEntityRef<'w> { + entity: UnsafeEntityCell<'w>, + access: Access, +} + +impl<'w> FilteredEntityRef<'w> { + /// # Safety + /// - No `&mut World` can exist from the underlying `UnsafeWorldCell` + /// - If `access` takes read access to a component no mutable reference to that + /// component can exist at the same time as the returned [`FilteredEntityMut`] + /// - If `access` takes any access for a component `entity` must have that component. + pub(crate) unsafe fn new(entity: UnsafeEntityCell<'w>, access: Access) -> Self { + Self { entity, access } + } + + /// Returns the [ID](Entity) of the current entity. + #[inline] + #[must_use = "Omit the .id() call if you do not need to store the `Entity` identifier."] + pub fn id(&self) -> Entity { + self.entity.id() + } + + /// Gets metadata indicating the location where the current entity is stored. + #[inline] + pub fn location(&self) -> EntityLocation { + self.entity.location() + } + + /// Returns the archetype that the current entity belongs to. + #[inline] + pub fn archetype(&self) -> &Archetype { + self.entity.archetype() + } + + /// Returns an iterator over the component ids that are accessed by self. + #[inline] + pub fn components(&self) -> impl Iterator + '_ { + self.access.reads_and_writes() + } + + /// Returns a reference to the underlying [`Access`]. + #[inline] + pub fn access(&self) -> &Access { + &self.access + } + + /// Returns `true` if the current entity has a component of type `T`. + /// Otherwise, this returns `false`. + /// + /// ## Notes + /// + /// If you do not know the concrete type of a component, consider using + /// [`Self::contains_id`] or [`Self::contains_type_id`]. + #[inline] + pub fn contains(&self) -> bool { + self.contains_type_id(TypeId::of::()) + } + + /// Returns `true` if the current entity has a component identified by `component_id`. + /// Otherwise, this returns false. + /// + /// ## Notes + /// + /// - If you know the concrete type of the component, you should prefer [`Self::contains`]. + /// - If you know the component's [`TypeId`] but not its [`ComponentId`], consider using + /// [`Self::contains_type_id`]. + #[inline] + pub fn contains_id(&self, component_id: ComponentId) -> bool { + self.entity.contains_id(component_id) + } + + /// Returns `true` if the current entity has a component with the type identified by `type_id`. + /// Otherwise, this returns false. + /// + /// ## Notes + /// + /// - If you know the concrete type of the component, you should prefer [`Self::contains`]. + /// - If you have a [`ComponentId`] instead of a [`TypeId`], consider using [`Self::contains_id`]. + #[inline] + pub fn contains_type_id(&self, type_id: TypeId) -> bool { + self.entity.contains_type_id(type_id) + } + + /// Gets access to the component of type `T` for the current entity. + /// Returns `None` if the entity does not have a component of type `T`. + #[inline] + pub fn get(&self) -> Option<&'w T> { + let Some(id) = self.entity.world().components().get_id(TypeId::of::()) else { + return None; + }; + self.access + .has_read(id) + // SAFETY: We have read access so we must have the component + .then(|| unsafe { self.entity.get().debug_checked_unwrap() }) + } + + /// Gets access to the component of type `T` for the current entity, + /// including change detection information as a [`Ref`]. + /// + /// Returns `None` if the entity does not have a component of type `T`. + #[inline] + pub fn get_ref(&self) -> Option> { + let Some(id) = self.entity.world().components().get_id(TypeId::of::()) else { + return None; + }; + self.access + .has_read(id) + // SAFETY: We have read access so we must have the component + .then(|| unsafe { self.entity.get_ref().debug_checked_unwrap() }) + } + + /// Retrieves the change ticks for the given component. This can be useful for implementing change + /// detection in custom runtimes. + #[inline] + pub fn get_change_ticks(&self) -> Option { + let Some(id) = self.entity.world().components().get_id(TypeId::of::()) else { + return None; + }; + self.access + .has_read(id) + // SAFETY: We have read access so we must have the component + .then(|| unsafe { self.entity.get_change_ticks::().debug_checked_unwrap() }) + } + + /// Retrieves the change ticks for the given [`ComponentId`]. This can be useful for implementing change + /// detection in custom runtimes. + /// + /// **You should prefer to use the typed API [`Self::get_change_ticks`] where possible and only + /// use this in cases where the actual component types are not known at + /// compile time.** + #[inline] + pub fn get_change_ticks_by_id(&self, component_id: ComponentId) -> Option { + // SAFETY: We have read access so we must have the component + self.access.has_read(component_id).then(|| unsafe { + self.entity + .get_change_ticks_by_id(component_id) + .debug_checked_unwrap() + }) + } + + /// Gets the component of the given [`ComponentId`] from the entity. + /// + /// **You should prefer to use the typed API [`Self::get`] where possible and only + /// use this in cases where the actual component types are not known at + /// compile time.** + /// + /// Unlike [`FilteredEntityRef::get`], this returns a raw pointer to the component, + /// which is only valid while the [`FilteredEntityRef`] is alive. + #[inline] + pub fn get_by_id(&self, component_id: ComponentId) -> Option> { + self.access + .has_read(component_id) + // SAFETY: We have read access so we must have the component + .then(|| unsafe { self.entity.get_by_id(component_id).debug_checked_unwrap() }) + } +} + +impl<'w> From> for FilteredEntityRef<'w> { + fn from(entity_mut: FilteredEntityMut<'w>) -> Self { + // SAFETY: + // - `FilteredEntityMut` guarantees exclusive access to all components in the new `FilteredEntityRef`. + unsafe { FilteredEntityRef::new(entity_mut.entity, entity_mut.access) } + } +} + +impl<'a> From<&'a FilteredEntityMut<'_>> for FilteredEntityRef<'a> { + fn from(entity_mut: &'a FilteredEntityMut<'_>) -> Self { + // SAFETY: + // - `FilteredEntityMut` guarantees exclusive access to all components in the new `FilteredEntityRef`. + unsafe { FilteredEntityRef::new(entity_mut.entity, entity_mut.access.clone()) } + } +} + +/// Provides mutable access to a single entity and some of its components defined by the contained [`Access`]. +pub struct FilteredEntityMut<'w> { + entity: UnsafeEntityCell<'w>, + access: Access, +} + +impl<'w> FilteredEntityMut<'w> { + /// # Safety + /// - No `&mut World` can exist from the underlying `UnsafeWorldCell` + /// - If `access` takes read access to a component no mutable reference to that + /// component can exist at the same time as the returned [`FilteredEntityMut`] + /// - If `access` takes write access to a component, no reference to that component + /// may exist at the same time as the returned [`FilteredEntityMut`] + /// - If `access` takes any access for a component `entity` must have that component. + pub(crate) unsafe fn new(entity: UnsafeEntityCell<'w>, access: Access) -> Self { + Self { entity, access } + } + + /// Returns a new instance with a shorter lifetime. + /// This is useful if you have `&mut FilteredEntityMut`, but you need `FilteredEntityMut`. + pub fn reborrow(&mut self) -> FilteredEntityMut<'_> { + // SAFETY: We have exclusive access to the entire entity and its components. + unsafe { Self::new(self.entity, self.access.clone()) } + } + + /// Gets read-only access to all of the entity's components. + pub fn as_readonly(&self) -> FilteredEntityRef<'_> { + FilteredEntityRef::from(self) + } + + /// Returns the [ID](Entity) of the current entity. + #[inline] + #[must_use = "Omit the .id() call if you do not need to store the `Entity` identifier."] + pub fn id(&self) -> Entity { + self.entity.id() + } + + /// Gets metadata indicating the location where the current entity is stored. + #[inline] + pub fn location(&self) -> EntityLocation { + self.entity.location() + } + + /// Returns the archetype that the current entity belongs to. + #[inline] + pub fn archetype(&self) -> &Archetype { + self.entity.archetype() + } + + /// Returns an iterator over the component ids that are accessed by self. + #[inline] + pub fn components(&self) -> impl Iterator + '_ { + self.access.reads_and_writes() + } + + /// Returns a reference to the underlying [`Access`]. + #[inline] + pub fn access(&self) -> &Access { + &self.access + } + + /// Returns `true` if the current entity has a component of type `T`. + /// Otherwise, this returns `false`. + /// + /// ## Notes + /// + /// If you do not know the concrete type of a component, consider using + /// [`Self::contains_id`] or [`Self::contains_type_id`]. + #[inline] + pub fn contains(&self) -> bool { + self.contains_type_id(TypeId::of::()) + } + + /// Returns `true` if the current entity has a component identified by `component_id`. + /// Otherwise, this returns false. + /// + /// ## Notes + /// + /// - If you know the concrete type of the component, you should prefer [`Self::contains`]. + /// - If you know the component's [`TypeId`] but not its [`ComponentId`], consider using + /// [`Self::contains_type_id`]. + #[inline] + pub fn contains_id(&self, component_id: ComponentId) -> bool { + self.entity.contains_id(component_id) + } + + /// Returns `true` if the current entity has a component with the type identified by `type_id`. + /// Otherwise, this returns false. + /// + /// ## Notes + /// + /// - If you know the concrete type of the component, you should prefer [`Self::contains`]. + /// - If you have a [`ComponentId`] instead of a [`TypeId`], consider using [`Self::contains_id`]. + #[inline] + pub fn contains_type_id(&self, type_id: TypeId) -> bool { + self.entity.contains_type_id(type_id) + } + + /// Gets access to the component of type `T` for the current entity. + /// Returns `None` if the entity does not have a component of type `T`. + #[inline] + pub fn get(&self) -> Option<&'_ T> { + self.as_readonly().get() + } + + /// Gets access to the component of type `T` for the current entity, + /// including change detection information as a [`Ref`]. + /// + /// Returns `None` if the entity does not have a component of type `T`. + #[inline] + pub fn get_ref(&self) -> Option> { + self.as_readonly().get_ref() + } + + /// Gets mutable access to the component of type `T` for the current entity. + /// Returns `None` if the entity does not have a component of type `T`. + #[inline] + pub fn get_mut(&mut self) -> Option> { + let Some(id) = self.entity.world().components().get_id(TypeId::of::()) else { + return None; + }; + self.access + .has_write(id) + // SAFETY: We have write access so we must have the component + .then(|| unsafe { self.entity.get_mut().debug_checked_unwrap() }) + } + + /// Retrieves the change ticks for the given component. This can be useful for implementing change + /// detection in custom runtimes. + #[inline] + pub fn get_change_ticks(&self) -> Option { + self.as_readonly().get_change_ticks::() + } + + /// Retrieves the change ticks for the given [`ComponentId`]. This can be useful for implementing change + /// detection in custom runtimes. + /// + /// **You should prefer to use the typed API [`Self::get_change_ticks`] where possible and only + /// use this in cases where the actual component types are not known at + /// compile time.** + #[inline] + pub fn get_change_ticks_by_id(&self, component_id: ComponentId) -> Option { + self.as_readonly().get_change_ticks_by_id(component_id) + } + + /// Gets the component of the given [`ComponentId`] from the entity. + /// + /// **You should prefer to use the typed API [`Self::get`] where possible and only + /// use this in cases where the actual component types are not known at + /// compile time.** + /// + /// Unlike [`FilteredEntityMut::get`], this returns a raw pointer to the component, + /// which is only valid while the [`FilteredEntityMut`] is alive. + #[inline] + pub fn get_by_id(&self, component_id: ComponentId) -> Option> { + self.as_readonly().get_by_id(component_id) + } + + /// Gets a [`MutUntyped`] of the component of the given [`ComponentId`] from the entity. + /// + /// **You should prefer to use the typed API [`Self::get_mut`] where possible and only + /// use this in cases where the actual component types are not known at + /// compile time.** + /// + /// Unlike [`FilteredEntityMut::get_mut`], this returns a raw pointer to the component, + /// which is only valid while the [`FilteredEntityMut`] is alive. + #[inline] + pub fn get_mut_by_id(&mut self, component_id: ComponentId) -> Option> { + // SAFETY: We have write access so we must have the component + self.access.has_write(component_id).then(|| unsafe { + self.entity + .get_mut_by_id(component_id) + .debug_checked_unwrap() + }) + } +} + /// Inserts a dynamic [`Bundle`] into the entity. /// /// # Safety diff --git a/crates/bevy_ecs/src/world/mod.rs b/crates/bevy_ecs/src/world/mod.rs index fe6533e8310ae..790867764ca84 100644 --- a/crates/bevy_ecs/src/world/mod.rs +++ b/crates/bevy_ecs/src/world/mod.rs @@ -7,7 +7,10 @@ pub mod unsafe_world_cell; mod world_cell; pub use crate::change_detection::{Mut, Ref, CHECK_TICK_THRESHOLD}; -pub use entity_ref::{EntityMut, EntityRef, EntityWorldMut, Entry, OccupiedEntry, VacantEntry}; +pub use entity_ref::{ + EntityMut, EntityRef, EntityWorldMut, Entry, FilteredEntityMut, FilteredEntityRef, + OccupiedEntry, VacantEntry, +}; pub use spawn_batch::*; pub use world_cell::*; diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.rs b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.rs new file mode 100644 index 0000000000000..e7d61d708e16e --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.rs @@ -0,0 +1,39 @@ +use bevy_ecs::prelude::*; +use bevy_ecs::system::SystemState; + +#[derive(Component, Eq, PartialEq, Debug)] +struct Foo(u32); + +#[derive(Component)] +struct Bar; + +fn main() { + let mut world = World::default(); + world.spawn(Foo(10)); + + let mut system_state = SystemState::>::new(&mut world); + let mut query = system_state.get_mut(&mut world); + + { + let mut lens_a = query.transmute_lens::<&mut Foo>(); + let mut lens_b = query.transmute_lens::<&mut Foo>(); + + let mut query_a = lens_a.query(); + let mut query_b = lens_b.query(); + + let a = query_a.single_mut(); + let b = query_b.single_mut(); // oops 2 mutable references to same Foo + assert_eq!(*a, *b); + } + + { + let mut lens = query.transmute_lens::<&mut Foo>(); + + let mut query_a = lens.query(); + let mut query_b = lens.query(); + + let a = query_a.single_mut(); + let b = query_b.single_mut(); // oops 2 mutable references to same Foo + assert_eq!(*a, *b); + } +} diff --git a/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.stderr b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.stderr new file mode 100644 index 0000000000000..4e8e283e39528 --- /dev/null +++ b/crates/bevy_ecs_compile_fail_tests/tests/ui/query_transmute_safety.stderr @@ -0,0 +1,21 @@ +error[E0499]: cannot borrow `query` as mutable more than once at a time + --> tests/ui/query_transmute_safety.rs:19:26 + | +18 | let mut lens_a = query.transmute_lens::<&mut Foo>(); + | ----- first mutable borrow occurs here +19 | let mut lens_b = query.transmute_lens::<&mut Foo>(); + | ^^^^^ second mutable borrow occurs here +20 | +21 | let mut query_a = lens_a.query(); + | ------ first borrow later used here + +error[E0499]: cannot borrow `lens` as mutable more than once at a time + --> tests/ui/query_transmute_safety.rs:33:27 + | +32 | let mut query_a = lens.query(); + | ---- first mutable borrow occurs here +33 | let mut query_b = lens.query(); + | ^^^^ second mutable borrow occurs here +34 | +35 | let a = query_a.single_mut(); + | ------- first borrow later used here diff --git a/examples/README.md b/examples/README.md index 4ce73d2819838..b315256190b57 100644 --- a/examples/README.md +++ b/examples/README.md @@ -224,6 +224,7 @@ Example | Description --- | --- [Component Change Detection](../examples/ecs/component_change_detection.rs) | Change detection on components [Custom Query Parameters](../examples/ecs/custom_query_param.rs) | Groups commonly used compound queries and query filters into a single type +[Dynamic ECS](../examples/ecs/dynamic.rs) | Dynamically create components, spawn entities with those components and query those components [ECS Guide](../examples/ecs/ecs_guide.rs) | Full guide to Bevy's ECS [Event](../examples/ecs/event.rs) | Illustrates event creation, activation, and reception [Fixed Timestep](../examples/ecs/fixed_timestep.rs) | Shows how to create systems that run every fixed timestep, rather than every tick diff --git a/examples/ecs/dynamic.rs b/examples/ecs/dynamic.rs new file mode 100644 index 0000000000000..fa1e84d53f7f2 --- /dev/null +++ b/examples/ecs/dynamic.rs @@ -0,0 +1,243 @@ +//! This example show how you can create components dynamically, spawn entities with those components +//! as well as query for entities with those components. + +use std::{alloc::Layout, io::Write, ptr::NonNull}; + +use bevy::prelude::*; +use bevy::{ + ecs::{ + component::{ComponentDescriptor, ComponentId, ComponentInfo, StorageType}, + query::{QueryBuilder, QueryData}, + world::FilteredEntityMut, + }, + ptr::OwningPtr, + utils::HashMap, +}; + +const PROMPT: &str = " +Commands: + comp, c Create new components + spawn, s Spawn entities + query, q Query for entities +Enter a command with no parameters for usage."; + +const COMPONENT_PROMPT: &str = " +comp, c Create new components + Enter a comma seperated list of type names optionally followed by a size in u64s. + e.g. CompA 3, CompB, CompC 2"; + +const ENTITY_PROMPT: &str = " +spawn, s Spawn entities + Enter a comma seperated list of components optionally followed by values. + e.g. CompA 0 1 0, CompB, CompC 1"; + +const QUERY_PROMPT: &str = " +query, q Query for entities + Enter a query to fetch and update entities + Components with read or write access will be displayed with their values + Components with write access will have their fields incremented by one + + Accesses: 'A' with, '&A' read, '&mut A' write + Operators: '||' or, ',' and, '?' optional + + e.g. &A || &B, &mut C, D, ?E"; + +fn main() { + let mut world = World::new(); + let mut lines = std::io::stdin().lines(); + let mut component_names = HashMap::::new(); + let mut component_info = HashMap::::new(); + + println!("{}", PROMPT); + loop { + print!("\n> "); + let _ = std::io::stdout().flush(); + let Some(Ok(line)) = lines.next() else { + return; + }; + + if line.is_empty() { + return; + }; + + let Some((first, rest)) = line.trim().split_once(|c: char| c.is_whitespace()) else { + match &line.chars().next() { + Some('c') => println!("{}", COMPONENT_PROMPT), + Some('s') => println!("{}", ENTITY_PROMPT), + Some('q') => println!("{}", QUERY_PROMPT), + _ => println!("{}", PROMPT), + } + continue; + }; + + match &first[0..1] { + "c" => { + rest.split(',').for_each(|component| { + let mut component = component.split_whitespace(); + let Some(name) = component.next() else { + return; + }; + let size = match component.next().map(|s| s.parse::()) { + Some(Ok(size)) => size, + _ => 0, + }; + // SAFETY: [u64] is Send + Sync + let id = world.init_component_with_descriptor(unsafe { + ComponentDescriptor::new_with_layout( + name.to_string(), + StorageType::Table, + Layout::array::(size).unwrap(), + None, + ) + }); + let Some(info) = world.components().get_info(id) else { + return; + }; + component_names.insert(name.to_string(), id); + component_info.insert(id, info.clone()); + println!("Component {} created with id: {:?}", name, id.index()); + }); + } + "s" => { + let mut to_insert_ids = Vec::new(); + let mut to_insert_ptr = Vec::new(); + rest.split(',').for_each(|component| { + let mut component = component.split_whitespace(); + let Some(name) = component.next() else { + return; + }; + let Some(&id) = component_names.get(name) else { + println!("Component {} does not exist", name); + return; + }; + let info = world.components().get_info(id).unwrap(); + let len = info.layout().size() / std::mem::size_of::(); + let mut values: Vec = component + .take(len) + .filter_map(|value| value.parse::().ok()) + .collect(); + + // SAFETY: + // - All components will be interpreted as [u64] + // - len and layout are taken directly from the component descriptor + let ptr = unsafe { + let data = std::alloc::alloc_zeroed(info.layout()).cast::(); + data.copy_from(values.as_mut_ptr(), values.len()); + let non_null = NonNull::new_unchecked(data.cast()); + OwningPtr::new(non_null) + }; + + to_insert_ids.push(id); + to_insert_ptr.push(ptr); + }); + + let mut entity = world.spawn_empty(); + // SAFETY: + // - Component ids have been taken from the same world + // - The pointer with the correct layout + unsafe { + entity.insert_by_ids(&to_insert_ids, to_insert_ptr.into_iter()); + } + println!("Entity spawned with id: {:?}", entity.id()); + } + "q" => { + let mut builder = QueryBuilder::::new(&mut world); + parse_query(rest, &mut builder, &component_names); + let mut query = builder.build(); + + query.iter_mut(&mut world).for_each(|filtered_entity| { + let terms = filtered_entity + .components() + .map(|id| { + let ptr = filtered_entity.get_by_id(id).unwrap(); + let info = component_info.get(&id).unwrap(); + let len = info.layout().size() / std::mem::size_of::(); + + // SAFETY: + // - All components are created with layout [u64] + // - len is calculated from the component descriptor + let data = unsafe { + std::slice::from_raw_parts_mut( + ptr.assert_unique().as_ptr().cast::(), + len, + ) + }; + if filtered_entity.access().has_write(id) { + data.iter_mut().for_each(|data| { + *data += 1; + }); + } + + format!("{}: {:?}", info.name(), data[0..len].to_vec()) + }) + .collect::>() + .join(", "); + + println!("{:?}: {}", filtered_entity.id(), terms); + }); + } + _ => continue, + } + } +} + +fn parse_term( + str: &str, + builder: &mut QueryBuilder, + components: &HashMap, +) { + let mut matched = false; + let str = str.trim(); + match str.chars().next() { + Some('?') => { + builder.optional(|b| parse_term(&str[1..], b, components)); + matched = true; + } + Some('&') => { + let mut parts = str.split_whitespace(); + let first = parts.next().unwrap(); + if first == "&mut" { + if let Some(str) = parts.next() { + if let Some(&id) = components.get(str) { + builder.mut_id(id); + matched = true; + } + }; + } else if let Some(&id) = components.get(&first[1..]) { + builder.ref_id(id); + matched = true; + } + } + Some(_) => { + if let Some(&id) = components.get(str) { + builder.with_id(id); + matched = true; + } + } + None => {} + }; + + if !matched { + println!("Unable to find component: {}", str); + } +} + +fn parse_query( + str: &str, + builder: &mut QueryBuilder, + components: &HashMap, +) { + let str = str.split(','); + str.for_each(|term| { + let sub_terms: Vec<_> = term.split("||").collect(); + if sub_terms.len() == 1 { + parse_term(sub_terms[0], builder, components); + } else { + builder.or(|b| { + sub_terms + .iter() + .for_each(|term| parse_term(term, b, components)); + }); + } + }); +}