bevy/examples/ecs/custom_query_param.rs
Chris Russell f7e112a3c9
Let query items borrow from query state to avoid needing to clone (#15396)
# Objective

Improve the performance of `FilteredEntity(Ref|Mut)` and
`Entity(Ref|Mut)Except`.

`FilteredEntityRef` needs an `Access<ComponentId>` to determine what
components it can access. There is one stored in the query state, but
query items cannot borrow from the state, so it has to `clone()` the
access for each row. Cloning the access involves memory allocations and
can be expensive.


## Solution

Let query items borrow from their query state.  

Add an `'s` lifetime to `WorldQuery::Item` and `WorldQuery::Fetch`,
similar to the one in `SystemParam`, and provide `&'s Self::State` to
the fetch so that it can borrow from the state.

Unfortunately, there are a few cases where we currently return query
items from temporary query states: the sorted iteration methods create a
temporary state to query the sort keys, and the
`EntityRef::components<Q>()` methods create a temporary state for their
query.

To allow these to continue to work with most `QueryData`
implementations, introduce a new subtrait `ReleaseStateQueryData` that
converts a `QueryItem<'w, 's>` to `QueryItem<'w, 'static>`, and is
implemented for everything except `FilteredEntity(Ref|Mut)` and
`Entity(Ref|Mut)Except`.

`#[derive(QueryData)]` will generate `ReleaseStateQueryData`
implementations that apply when all of the subqueries implement
`ReleaseStateQueryData`.

This PR does not actually change the implementation of
`FilteredEntity(Ref|Mut)` or `Entity(Ref|Mut)Except`! That will be done
as a follow-up PR so that the changes are easier to review. I have
pushed the changes as chescock/bevy#5.

## Testing

I ran performance traces of many_foxes, both against main and against
chescock/bevy#5, both including #15282. These changes do appear to make
generalized animation a bit faster:

(Red is main, yellow is chescock/bevy#5)

![image](https://github.com/user-attachments/assets/de900117-0c6a-431d-ab62-c013834f97a9)


## Migration Guide

The `WorldQuery::Item` and `WorldQuery::Fetch` associated types and the
`QueryItem` and `ROQueryItem` type aliases now have an additional
lifetime parameter corresponding to the `'s` lifetime in `Query`. Manual
implementations of `WorldQuery` will need to update the method
signatures to include the new lifetimes. Other uses of the types will
need to be updated to include a lifetime parameter, although it can
usually be passed as `'_`. In particular, `ROQueryItem` is used when
implementing `RenderCommand`.

Before: 

```rust
fn render<'w>(
    item: &P,
    view: ROQueryItem<'w, Self::ViewQuery>,
    entity: Option<ROQueryItem<'w, Self::ItemQuery>>,
    param: SystemParamItem<'w, '_, Self::Param>,
    pass: &mut TrackedRenderPass<'w>,
) -> RenderCommandResult;
```

After: 

```rust
fn render<'w>(
    item: &P,
    view: ROQueryItem<'w, '_, Self::ViewQuery>,
    entity: Option<ROQueryItem<'w, '_, Self::ItemQuery>>,
    param: SystemParamItem<'w, '_, Self::Param>,
    pass: &mut TrackedRenderPass<'w>,
) -> RenderCommandResult;
```

---

Methods on `QueryState` that take `&mut self` may now result in
conflicting borrows if the query items capture the lifetime of the
mutable reference. This affects `get()`, `iter()`, and others. To fix
the errors, first call `QueryState::update_archetypes()`, and then
replace a call `state.foo(world, param)` with
`state.query_manual(world).foo_inner(param)`. Alternately, you may be
able to restructure the code to call `state.query(world)` once and then
make multiple calls using the `Query`.

Before:
```rust
let mut state: QueryState<_, _> = ...;
let d1 = state.get(world, e1);
let d2 = state.get(world, e2); // Error: cannot borrow `state` as mutable more than once at a time
println!("{d1:?}");
println!("{d2:?}");
```

After: 
```rust
let mut state: QueryState<_, _> = ...;

state.update_archetypes(world);
let d1 = state.get_manual(world, e1);
let d2 = state.get_manual(world, e2);
// OR
state.update_archetypes(world);
let d1 = state.query(world).get_inner(e1);
let d2 = state.query(world).get_inner(e2);
// OR
let query = state.query(world);
let d1 = query.get_inner(e1);
let d1 = query.get_inner(e2);

println!("{d1:?}");
println!("{d2:?}");
```
2025-06-16 21:05:41 +00:00

196 lines
6.1 KiB
Rust

//! This example illustrates the usage of the [`QueryData`] derive macro, which allows
//! defining custom query and filter types.
//!
//! While regular tuple queries work great in most of simple scenarios, using custom queries
//! declared as named structs can bring the following advantages:
//! - They help to avoid destructuring or using `q.0, q.1, ...` access pattern.
//! - Adding, removing components or changing items order with structs greatly reduces maintenance
//! burden, as you don't need to update statements that destructure tuples, care about order
//! of elements, etc. Instead, you can just add or remove places where a certain element is used.
//! - Named structs enable the composition pattern, that makes query types easier to re-use.
//! - You can bypass the limit of 15 components that exists for query tuples.
//!
//! For more details on the [`QueryData`] derive macro, see the trait documentation.
use bevy::{
ecs::query::{QueryData, QueryFilter},
prelude::*,
};
use std::fmt::Debug;
fn main() {
App::new()
.add_systems(Startup, spawn)
.add_systems(
Update,
(
print_components_read_only,
print_components_iter_mut,
print_components_iter,
print_components_tuple,
)
.chain(),
)
.run();
}
#[derive(Component, Debug)]
struct ComponentA;
#[derive(Component, Debug)]
struct ComponentB;
#[derive(Component, Debug)]
struct ComponentC;
#[derive(Component, Debug)]
struct ComponentD;
#[derive(Component, Debug)]
struct ComponentZ;
#[derive(QueryData)]
#[query_data(derive(Debug))]
struct ReadOnlyCustomQuery<T: Component + Debug, P: Component + Debug> {
entity: Entity,
a: &'static ComponentA,
b: Option<&'static ComponentB>,
nested: NestedQuery,
optional_nested: Option<NestedQuery>,
optional_tuple: Option<(&'static ComponentB, &'static ComponentZ)>,
generic: GenericQuery<T, P>,
empty: EmptyQuery,
}
fn print_components_read_only(
query: Query<
ReadOnlyCustomQuery<ComponentC, ComponentD>,
CustomQueryFilter<ComponentC, ComponentD>,
>,
) {
println!("Print components (read_only):");
for e in &query {
println!("Entity: {}", e.entity);
println!("A: {:?}", e.a);
println!("B: {:?}", e.b);
println!("Nested: {:?}", e.nested);
println!("Optional nested: {:?}", e.optional_nested);
println!("Optional tuple: {:?}", e.optional_tuple);
println!("Generic: {:?}", e.generic);
}
println!();
}
/// If you are going to mutate the data in a query, you must mark it with the `mutable` attribute.
///
/// The [`QueryData`] derive macro will still create a read-only version, which will be have `ReadOnly`
/// suffix.
/// Note: if you want to use derive macros with read-only query variants, you need to pass them with
/// using the `derive` attribute.
#[derive(QueryData)]
#[query_data(mutable, derive(Debug))]
struct CustomQuery<T: Component + Debug, P: Component + Debug> {
entity: Entity,
a: &'static mut ComponentA,
b: Option<&'static mut ComponentB>,
nested: NestedQuery,
optional_nested: Option<NestedQuery>,
optional_tuple: Option<(NestedQuery, &'static mut ComponentZ)>,
generic: GenericQuery<T, P>,
empty: EmptyQuery,
}
// This is a valid query as well, which would iterate over every entity.
#[derive(QueryData)]
#[query_data(derive(Debug))]
struct EmptyQuery {
empty: (),
}
#[derive(QueryData)]
#[query_data(derive(Debug))]
struct NestedQuery {
c: &'static ComponentC,
d: Option<&'static ComponentD>,
}
#[derive(QueryData)]
#[query_data(derive(Debug))]
struct GenericQuery<T: Component, P: Component> {
generic: (&'static T, &'static P),
}
#[derive(QueryFilter)]
struct CustomQueryFilter<T: Component, P: Component> {
_c: With<ComponentC>,
_d: With<ComponentD>,
_or: Or<(Added<ComponentC>, Changed<ComponentD>, Without<ComponentZ>)>,
_generic_tuple: (With<T>, With<P>),
}
fn spawn(mut commands: Commands) {
commands.spawn((ComponentA, ComponentB, ComponentC, ComponentD));
}
fn print_components_iter_mut(
mut query: Query<
CustomQuery<ComponentC, ComponentD>,
CustomQueryFilter<ComponentC, ComponentD>,
>,
) {
println!("Print components (iter_mut):");
for e in &mut query {
// Re-declaring the variable to illustrate the type of the actual iterator item.
let e: CustomQueryItem<'_, '_, _, _> = e;
println!("Entity: {}", e.entity);
println!("A: {:?}", e.a);
println!("B: {:?}", e.b);
println!("Optional nested: {:?}", e.optional_nested);
println!("Optional tuple: {:?}", e.optional_tuple);
println!("Nested: {:?}", e.nested);
println!("Generic: {:?}", e.generic);
}
println!();
}
fn print_components_iter(
query: Query<CustomQuery<ComponentC, ComponentD>, CustomQueryFilter<ComponentC, ComponentD>>,
) {
println!("Print components (iter):");
for e in &query {
// Re-declaring the variable to illustrate the type of the actual iterator item.
let e: CustomQueryReadOnlyItem<'_, '_, _, _> = e;
println!("Entity: {}", e.entity);
println!("A: {:?}", e.a);
println!("B: {:?}", e.b);
println!("Nested: {:?}", e.nested);
println!("Generic: {:?}", e.generic);
}
println!();
}
type NestedTupleQuery<'w> = (&'w ComponentC, &'w ComponentD);
type GenericTupleQuery<'w, T, P> = (&'w T, &'w P);
fn print_components_tuple(
query: Query<
(
Entity,
&ComponentA,
&ComponentB,
NestedTupleQuery,
GenericTupleQuery<ComponentC, ComponentD>,
),
(
With<ComponentC>,
With<ComponentD>,
Or<(Added<ComponentC>, Changed<ComponentD>, Without<ComponentZ>)>,
),
>,
) {
println!("Print components (tuple):");
for (entity, a, b, nested, (generic_c, generic_d)) in &query {
println!("Entity: {entity}");
println!("A: {a:?}");
println!("B: {b:?}");
println!("Nested: {:?} {:?}", nested.0, nested.1);
println!("Generic: {generic_c:?} {generic_d:?}");
}
}