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:

Details

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?

[cBournhonesque]

Could we remove this bound?

[chescock]

Could we remove this bound?

Hmm, I think we could, but I don't think it's worth it for this PR.

The goal is to be able to cheaply convert &QueryState<D, F> to &QueryState<D::ReadOnly, F> during Query::as_readonly() for things like Query::iter(). That's done by making QueryState be #[repr(transparent)] and only containing D::State and F::State, which means QueryState<D, F> and QueryState<D::ReadOnly, F> are guaranteed to be transmutable.

We could technically drop the constraint and have a safety requirement that D::State be transmutable to F::State. But that would make it easier to get wrong, and the transmutes are already tricky to justify.

Another approach would be to rework it like

#[repr(transparent)]
struct QueryState<D, F> {
    inner: QueryStateInner<D::State, F::State>,
}

So that we have a type parameterized by the state types instead of the query types themselves. Then we could do type State = QueryStateInner<D::State, F::State> and it would work. That would also let us justify the transmutes without the need for #[repr(c)], letting the compiler optimize the layout slightly. But that would be a pretty extensive refactor!

I'd be inclined to see what happens to QueryState with #21607 and #19787 before trying to clean this up more. If QueryState winds up behind a trait bound, we might have more room to maneuver here.

[cBournhonesque]

Is get_state only used for World::try_query where we try to create a query using an immutable reference to world?
Could we change the API to take &World as input, and then call QueryState::try_new()?

Although apparently allowing &World allows UB: #13343

[chescock]

Is get_state only used for World::try_query where we try to create a query using an immutable reference to world? Could we change the API to take &World as input, and then call QueryState::try_new()?

Although apparently allowing &World allows UB: #13343

It's also used by query transmutes and joins, like transmute_lens. That's where the UB comes in: transmute_lens is called with the access of the source query, which wouldn't have access to the DefaultQueryFilters resource that's needed to create a new QueryState.

I think the ideal behavior would be to build the nested query state by transmuting the nested query state from the source query. Like, transmuting Query<Parent<(&T, &U)>> to Query<Parent<&T>> would still only match entities whose parent has both T and U. Then we wouldn't need to re-evaluate DefaultQueryFilters. But I don't see any good way to match up the nested queries in the source with those in the target.

[cBournhonesque]

I want to pursue multi-source queries using a different approach (dynamic query plans similar to what flecs is doing), but I think that the traits introduced are broadly useful.
It will also be nice to have a shorter path towards relation queries and to gain more experience in usage/API design related to them.

Awesome work!

[alice-i-cecile]

I want to pursue multi-source queries using a different approach (dynamic query plans similar to what flecs is doing)

FWIW, my opinion is that both approaches will be useful and can happily co-exist :) I do think that this is a key missing feature for bevy_ecs, regardless of what route we take to get there.

[Trashtalk217]

My main concern with this PR is ballooning complexity. Especially SingleEntityQueryData is very FactoryManagerAdaptor-coded and in my opinion signals that we're reaching the limit of our type-based query language.

That being said, the change by itself is not complicated, and it has clear and immediate advantages. I would like to see the documentation and safety comments improved a bit. Also, if I were given the final say I would probably go with the alternative where IterQueryData would be required for all QueryData and reduce the number of traits, but this is not super important. I will approve either version as long as the documentation is improved.

[Trashtalk217]

What I understand is that some QueryData (probably the nested types) might go over the same entity twice, generating two Items, which can conflict if they both touch the same component and at least one has write access. Is this correct?

[Trashtalk217]

This briefly confused me in the sense that I thought this QueryData was only ever used for Single<> queries, but apparently it allows you to access data on more than one entity as long as you don't do it at the same time.

[Trashtalk217]

I don't understand this. Could you not just do a world.query<Entity>().iter(&world).collect()? I understand that this would still be fine because Entity is read-only. Still, you can access more entities at the same time.

[Trashtalk217]

I have the same question for most other implementors of SingleEntityQueryData.

[Trashtalk217]

What is the 'current' entity? In what context?

[Trashtalk217]

Shouldn't this safety-warning be removed for this function? If I'm reading this right, you do not have to call Self::init_access before into_readonly. Why is it even unsafe then?

[Trashtalk217]

Move this to the safety section.

[Trashtalk217]

The first thing I think when reading this is 'but we've always been able to do that, I can do for (compA, compB) in query.iter() { ... } already right? That accesses multiple entities, right?'

I know that this is a migration guide, not a release note, but can you word it more clearly? I don't really have any alternative because the word 'Access' is just very overloaded here.