Nested Queries

[chescock]

Objective

Support queries that soundly access multiple entities.

This can be used to create queries that follow relations, as in #17647.

This can also be used to create queries that perform resource access. This has been supported since #16843, although that approach may become unsound if we do resources-as-components #19731, such as #21346.

Fixes #20315

Solution

Allow a QueryData that wants to access other entities to store a QueryState<D, F> in its WorldQuery::State, so that it can create a nested Query<D, F> during the outer fetch.

NestedQuery type

Introduce a NestedQuery type that implements QueryData by yielding a Query. It is intended to be used inside other implementations of QueryData, either for manual implementations or #[derive(QueryData)]. It is not normally useful to query directly, since it's equivalent to adding another Query parameter to a system.

In theory, we could directly impl QueryData for Query, but this would be too easy to do accidentally. Having to explicitly import and write NestedQuery will make it clear that it's something unusual, and also allows us to remove the need for passing 'static for the 'w and 's lifetimes.

New WorldQuery methods

For it to be sound to create the Query during fetch, we need to register the FilteredAccess of the nested query and check for conflicts with other parameters. Create a WorldQuery::update_external_component_access method for that purpose. For Query as SystemParam, call this during init_access so the access can be combined with the rest of the system access. For loose QueryStates, call it during QueryState::new.

In order to keep the query cache up-to-date, create a WorldQuery::update_archetypes method where it can call QueryState::update_archetypes_unsafe_world_cell, and call it from there.

New QueryData subtraits

Some operations would not be sound with nested queries! In particular, we want a Parent<D> query that reads data from the parent entity by following the ChildOf relation. But many entities may share a parent, so it's not sound to iterate a Query<Parent<&mut C>>.

It is sound to get_mut, though, so we want the query type to exist, just not be iterable. And following the relation in the other direction for a Query<Children<&mut C>> is sound to iterate, since children are unique to a given parent.

So, introduce two new QueryData subtraits:

• IterQueryData - For anything it's sound to iterate. This is used to bound iter_mut and related methods.
• SingleEntityQueryData - For anything that only accesses data from one entity. This is used to bound EntityRef::get_components (see Unsound to call EntityRef::get_components with a QueryData that performs resource access #20315). It's also used to bound transmute at the moment, although we may be able to relax that later (see below, under Future Work).

Note that SingleEntityQueryData: IterQueryData, since single-entity queries never alias data across entities, and ReadOnlyQueryData: IterQueryData, since it's always sound to alias read-only data.

Here is a summary of the traits implemented by some representative QueryData:

Data	Iter	ReadOnly	SingleEntity
&T	✓	✓	✓
&mut T	✓	x	✓
Parent<&T>	✓	✓	x
Parent<&mut T>	x	x	x
(&mut T, Parent<&U>)	✓	x	x
Children<&mut T>	✓	x	x

Alternatives

We could avoid the need for the IterQueryData trait by making it a requirement for all QueryData. That would reduce the number of traits required, at the cost of making it impossible to support Query<Parent<&mut C>>.

Showcase

Here is an implementation of a Related<R, D, F> query using this PR:

pub struct Related<R: Relationship, D: QueryData + 'static, F: QueryFilter + 'static = ()>(
    RelatedInner<R, D, F>,
);

type RelatedInner<R, D, F> = (
    &'static R,
    NestedQuery<D, (F, With<<R as Relationship>::RelationshipTarget>)>,
);

unsafe impl<R: Relationship, D: QueryData + 'static, F: QueryFilter + 'static> WorldQuery
    for Related<R, D, F>
{
    type Fetch<'w> = <RelatedInner<R, D, F> as WorldQuery>::Fetch<'w>;
    type State = <RelatedInner<R, D, F> as WorldQuery>::State;

    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
        <RelatedInner<R, D, F> as WorldQuery>::shrink_fetch(fetch)
    }

    unsafe fn init_fetch<'w, 's>(
        world: UnsafeWorldCell<'w>,
        state: &'s Self::State,
        last_run: Tick,
        this_run: Tick,
    ) -> Self::Fetch<'w> {
        unsafe {
            <RelatedInner<R, D, F> as WorldQuery>::init_fetch(world, state, last_run, this_run)
        }
    }

    const IS_DENSE: bool = <RelatedInner<R, D, F> as WorldQuery>::IS_DENSE;

    unsafe fn set_archetype<'w, 's>(
        fetch: &mut Self::Fetch<'w>,
        state: &'s Self::State,
        archetype: &'w Archetype,
        table: &'w Table,
    ) {
        unsafe {
            <RelatedInner<R, D, F> as WorldQuery>::set_archetype(fetch, state, archetype, table)
        };
    }

    unsafe fn set_table<'w, 's>(
        fetch: &mut Self::Fetch<'w>,
        state: &'s Self::State,
        table: &'w Table,
    ) {
        unsafe { <RelatedInner<R, D, F> as WorldQuery>::set_table(fetch, state, table) };
    }

    fn update_component_access(state: &Self::State, access: &mut FilteredAccess) {
        <RelatedInner<R, D, F> as WorldQuery>::update_component_access(state, access);
    }

    fn init_nested_access(
        state: &Self::State,
        system_name: Option<&str>,
        component_access_set: &mut FilteredAccessSet,
        world: UnsafeWorldCell,
    ) {
        <RelatedInner<R, D, F> as WorldQuery>::init_nested_access(state, system_name, component_access_set, world);
    }

    fn init_state(world: &mut World) -> Self::State {
        <RelatedInner<R, D, F> as WorldQuery>::init_state(world)
    }

    fn get_state(components: &Components) -> Option<Self::State> {
        <RelatedInner<R, D, F> as WorldQuery>::get_state(components)
    }

    fn matches_component_set(
        state: &Self::State,
        set_contains_id: &impl Fn(ComponentId) -> bool,
    ) -> bool {
        <RelatedInner<R, D, F> as WorldQuery>::matches_component_set(state, set_contains_id)
    }

    fn update_archetypes(state: &mut Self::State, world: UnsafeWorldCell) {
        <RelatedInner<R, D, F> as WorldQuery>::update_archetypes(state, world);
    }
}

unsafe impl<R: Relationship, D: QueryData + 'static, F: QueryFilter + 'static> QueryData
    for Related<R, D, F>
{
    const IS_READ_ONLY: bool = D::IS_READ_ONLY;
    type ReadOnly = Related<R, D::ReadOnly, F>;
    type Item<'w, 's> = Option<D::Item<'w, 's>>;

    fn shrink<'wlong: 'wshort, 'wshort, 's>(
        item: Self::Item<'wlong, 's>,
    ) -> Self::Item<'wshort, 's> {
        item.map(D::shrink)
    }

    unsafe fn fetch<'w, 's>(
        state: &'s Self::State,
        fetch: &mut Self::Fetch<'w>,
        entity: Entity,
        table_row: TableRow,
    ) -> Self::Item<'w, 's> {
        let (relationship, query) =
            unsafe { <RelatedInner<R, D, F> as QueryData>::fetch(state, fetch, entity, table_row) };
        query.get_inner(relationship.get()).ok()
    }
}

unsafe impl<R: Relationship, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> ReadOnlyQueryData for Related<R, D, F> { }

// Note that we require `D: ReadOnlyQueryData` for `Related: IterQueryData`
unsafe impl<R: Relationship, D: ReadOnlyQueryData + 'static, F: QueryFilter + 'static> IterQueryData for Related<R, D, F> { }

I'd like to leave that to a follow-up PR to allow bikeshedding the API, and to take advantage of #21581 to remove the Option, but I think it works!

Future Work

There is more to do here, but this PR is already pretty big. Future work includes:

• WorldQuery types for working with relationships #17647
• Following Store resources as components on singleton entities (v2) #21346, update AssetChanged to use nested queries for resource access, and stop tracking resource access separately in Access
• Implement get_state for NestedQuery. This is difficult because constructing a QueryState requires reading the DefaultQueryFilters resource, but get_state can be called from transmute with no access.
• Relax the SingleEntityQueryData bound on transmutes and joins. This will require checking that the nested query access is also a subset of the original access. Although unless we also solve the problem of implementing get_state, transmuting to a query with nested queries won't work anyway.
• Support streaming iteration for QueryIter by offering a fn fetch_next(&self) -> D::Item<'_> method and relaxing the IterQueryData bound on Query::into_iter and Query::iter_mut. This would work similar to iter_many_mut and iter_many_inner.
• Relax the IterQueryData bound on Query::single_inner, Query::single_mut, and Single<D, F>. This seems like it should be straightforward, because the method only returns a single item. But the way it checks that there is only one item is by fetching the second one!

[Freyja-moth]

so it's not sound to iterate a Query<Parent<&mut C>>.

I'm not certain it's the best way of going about it but couldn't we implement this by storing the entity of the parent instead of the data and then resolving the data from the entity when it is needed?

[chescock]

couldn't we implement this by storing the entity of the parent instead of the data and then resolving the data from the entity when it is needed?

I'm not sure what you mean. Like, instead of yielding a D::Item<'w, 's>, we could yield some type Foo with a fn get_mut(&mut self) -> D::Item<'_, '_>? That wouldn't help, since it would still be possible to collect the Foo values and then call get_mut on several of them concurrently, like:

let mut items = query.iter_mut().collect::<Vec<_>>();
let mapped = items.iter_mut().map(|item| item.get_mut()).collect::<Vec<_>>();

[cBournhonesque]

I think these changes are a great idea; I guess I would like to know how this contrasts to how flecs queries for relation data. Maybe @james-j-obrien knows?

[james-j-obrien]

I think these changes are a great idea; I guess I would like to know how this contrasts to how flecs queries for relation data. Maybe @james-j-obrien knows?

That's a quite involved question to answer (although a very interesting one).

The big difference between the bevy and flecs queries are tied to one being defined in the type system and the other being defined dynamically. Due to this bevy queries are fundamentally based on nesting, you have tuples of query terms that each store their own state and generate their own code for managing that state. In flecs all the query terms are just stored in a flat array.

For example in this PR we express querying our parent as creating a query term that traverses the relationship and then nested in that is the set of components we want to access on the target, whereas in flecs you would have a set of instructions that said: "get me any entity A with relationship of the form (ChildOf, B) and store B as a variable", "get me component Y on entity B", "get me component Z on entity B".

This structure allows flecs to optimize/batch/reorder terms since they can be evaluated in the full context of the rest of the query, but for simple queries it's mostly a different path to the same goal.

Since flecs also has fragmenting relations they can do stuff like cache the tables for B since you know that entities in A's table will always have parent B.

All that being said, with bevy's queries as they exist today this PR seems like the shortest path to querying on multiple entities so seems like a positive step.

[cBournhonesque]

I think these changes are a great idea; I guess I would like to know how this contrasts to how flecs queries for relation data. Maybe @james-j-obrien knows?

That's a quite involved question to answer (although a very interesting one).

The big difference between the bevy and flecs queries are tied to one being defined in the type system and the other being defined dynamically. Due to this bevy queries are fundamentally based on nesting, you have tuples of query terms that each store their own state and generate their own code for managing that state. In flecs all the query terms are just stored in a flat array.

For example in this PR we express querying our parent as creating a query term that traverses the relationship and then nested in that is the set of components we want to access on the target, whereas in flecs you would have a set of instructions that said: "get me any entity A with relationship of the form (ChildOf, B) and store B as a variable", "get me component Y on entity B", "get me component Z on entity B".

This structure allows flecs to optimize/batch/reorder terms since they can be evaluated in the full context of the rest of the query, but for simple queries it's mostly a different path to the same goal.

Since flecs also has fragmenting relations they can do stuff like cache the tables for B since you know that entities in A's table will always have parent B.

All that being said, with bevy's queries as they exist today this PR seems like the shortest path to querying on multiple entities so seems like a positive step.

Thanks for the answer!
I guess bevy can also use the dynamic QueryBuilder, but the entire design is heavily influenced by primarily using types.
After reading https://ajmmertens.medium.com/building-games-in-ecs-with-entity-relationships-657275ba2c6c, it sounds like flecs creates some kind of data structure (node graph) that allow it to efficiently match entities.

I guess we could do something similar: for tuple queries, build such a node graph and use it to match entities. I guess we do already create a data structure that helps us find matching entities; that data structure is the QueryState.
The main difference that our state simply has:

• matched archetypes (from QueryData::matches_component_set)
• uses types to filter out entities (F::filter)

And the main difference is that the flecs "QueryState" is more elaborate since it can contain sources, relationships, etc.
So this PR's NestedQueries is one way to add more complexity to our QueryState. But we still have a simple 'combined query' since our Tuple QueryState combines the inner WorldQueries' QueryState in a very simple manner. In flecs it would combine them by adding them into a dynamic graph that can then be optimized. Our equivalent would be to add a WorldQuery::add_to_query_plan method so that we would also be able to optimize a query that contains multiple terms

[alice-i-cecile]

@eugineerd, I liked your work over in #21601; can I get your review here in turn?

[eugineerd]

Mostly skimmed through this as I see there are other changes planned for the trait methods part, just a bit confused about some safety comments

[SanderMertens]

@james-j-obrien @cBournhonesque Another difference with Flecs relationship queries, and a limitation (I think) of the solution proposed in this PR is that you cannot have constraints that span multiple terms. For example:

// find all entities that like the same food that they're eating
(Likes, $food), (Eats, $food)

or

// Match all spaceships that are docked to a planet of the same faction that they belong to
SpaceShip, (Faction, $f), (DockedTo, $obj), Faction($obj, $f)

[cBournhonesque]

Yes, I think we would need a more complex QueryBuilder like james suggested that can handle term 'variables'. Personally I'm happy with the PR; I think it's a great effort and @chescock always show a great attention detail. However I wonder if we shouldn't instead focus on providing extra capabilities to a more dynamic QueryBuilder.

I think the idea was floated to have a dynamic QueryBuilder which would be transmuted to Dynamic, which is a WorldQuery that is similar to FilteredEntityRef except that it would give you even more freedom on how to fetch data from the ECS. For example, FilteredEntityRef lets you fetch arbitrary data from the current term ($this variable in flecs) but Dynamic could allow you to fetch even data for other query terms.

Maybe it's independent from this PR?
Basically our matching logic is currently

• archetype matches if <D as WorldQuery>::matches() && <F as WorldQuery>::matches()
  where the WorldQuery is usually a tuple of inner WorldQueries, and we call matches() for each inner term one by one.
  If we use a dynamic query, we would have <Dynamic as WorldQuery>::matches() which would be using the underlying QueryBuilder (that can span multiple terms, etc.) to do the matching.

I guess some questions I have are:

• would we want to reimplement all possible typed-queries inside the dynamic QueryBuilder? For example could we specify 'Changed' inside the dynamic query builder?
• I think @chescock mentioned that the static type-based matching logic might be more efficient because the compiler can optimize a number of things, since it knows the exact shape of the WorldQuery (for example the tuple size)
• @chescock also mentioned that we already accomplish a lot of the more complex queries (such as multi-term queries) by simply using multiple queries in the System. For instance we can do SpaceShip, (Faction, $f), (DockedTo, $obj), Faction($obj, $f) using multiple Queries. I guess the 2 main reasons why we would want to build a dynamic querybuilder is:
  1. performance. Currently each term of the query has a special QueryState which does through new archetype. In a dynamic context we could probably optimize these dynamic queries, and we would store a single QueryState
  2. ergonomics. Sometimes it's easier and less cumbersome to express things in a single query. I guess here we need to have some clever transmute logic to check that the dynamic query can be transmuted into something &C1, ChildOf<(), With<C2>>

[eugineerd]

couldn't we implement this by storing the entity of the parent instead of the data and then resolving the data from the entity when it is needed?

I'm not sure what you mean. Like, instead of yielding a D::Item<'w, 's>, we could yield some type Foo with a fn get_mut(&mut self) -> D::Item<'_, '_>? That wouldn't help, since it would still be possible to collect the Foo values and then call get_mut on several of them concurrently, like:

let mut items = query.iter_mut().collect::<Vec<_>>();
let mapped = items.iter_mut().map(|item| item.get_mut()).collect::<Vec<_>>();

Would it be possible if Query implemented a lending iterator interface? It will be less ergonomic to use (while let instead of for), but at least such query will be expressible and usable where appropriate. Not saying we need it in this PR, just if it would be possible in general.

[chescock]

Would it be possible if Query implemented a lending iterator interface?

Yeah, I want to leave that for a follow-up, but I do think we should do it!

My plan is to give QueryIter a fetch_next method like QueryManyIter has, and make the Iterator impl conditional. So you can always call query.iter_mut(), but whether you can use it as an iterator or only call fetch_next depends on the QueryData.

I don't think it will be hard to do, but I think it would make the diff harder to review if it were combined with this PR.