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.

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 that almost works:

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

struct RelatedState<R: Relationship, D: QueryData + 'static, F: QueryFilter + 'static> {
    relationship: <&'static R as WorldQuery>::State,
    nested: QueryState<D, (F, With<R::RelationshipTarget>)>,
}

struct RelatedFetch<'w, R: Relationship> {
    relationship: <&'static R as WorldQuery>::Fetch<'w>,
    world: UnsafeWorldCell<'w>,
    last_run: Tick,
    this_run: Tick,
}

impl<R: Relationship> Clone for RelatedFetch<'_, R> {
    fn clone(&self) -> Self {
        Self { relationship: self.relationship.clone(), world: self.world, last_run: self.last_run, this_run: self.this_run }
    }
}

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

    fn shrink_fetch<'wlong: 'wshort, 'wshort>(fetch: Self::Fetch<'wlong>) -> Self::Fetch<'wshort> {
        let RelatedFetch { relationship, world, last_run, this_run } = fetch;
        let relationship = <&R>::shrink_fetch(relationship);
        RelatedFetch { relationship, world, last_run, this_run }
    }

    unsafe fn init_fetch<'w, 's>(world: UnsafeWorldCell<'w>, state: &'s Self::State, last_run: Tick, this_run: Tick) -> Self::Fetch<'w> {
        let relationship = unsafe { <&R>::init_fetch(world, &state.relationship, last_run, this_run) };
        RelatedFetch { relationship, world, this_run, last_run }
    }

    const IS_DENSE: bool = <&R>::IS_DENSE;

    unsafe fn set_archetype<'w, 's>(fetch: &mut Self::Fetch<'w>, state: &'s Self::State, archetype: &'w Archetype, table: &'w Table) {
        unsafe { <&R>::set_archetype(&mut fetch.relationship, &state.relationship, archetype, table) };
    }

    unsafe fn set_table<'w, 's>(fetch: &mut Self::Fetch<'w>, state: &'s Self::State, table: &'w Table) {
        unsafe { <&R>::set_table(&mut fetch.relationship, &state.relationship, table) };
    }

    fn update_component_access(state: &Self::State, access: &mut FilteredAccess) {
        <&R>::update_component_access(&state.relationship, access);
    }

    fn update_external_component_access(state: &Self::State, system_name: Option<&str>, component_access_set: &mut FilteredAccessSet, world: UnsafeWorldCell) {
        state.nested.update_external_component_access(system_name, component_access_set, world);
    }

    fn init_state(world: &mut World) -> Self::State {
        RelatedState {
            relationship: <&R>::init_state(world),
            nested: unsafe { QueryState::new_unchecked(world) },
        }
    }

    fn get_state(_components: &Components) -> Option<Self::State> {
        // :(
        None
    }

    fn matches_component_set(state: &Self::State, set_contains_id: &impl Fn(ComponentId) -> bool) -> bool {
        <&R>::matches_component_set(&state.relationship, set_contains_id)
    }

    fn update_archetypes(state: &mut Self::State, world: UnsafeWorldCell) {
        state.nested.update_archetypes_unsafe_world_cell(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 = <&R>::fetch(&state.relationship, &mut fetch.relationship, entity, table_row);
        let query = unsafe { state.nested.query_unchecked_manual_with_ticks(fetch.world, fetch.last_run, fetch.this_run ) };
        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> { }

That has a few flaws, notably that it fails with

error[E0271]: type mismatch resolving `<Related<R, <D as QueryData>::ReadOnly, F> as WorldQuery>::State == RelatedState<R, D, F>`

because QueryData requires that State = ReadOnly::State, but QueryState<D, F> != QueryState<D::ReadOnly, F>.

It's also impossible to implement get_state, because constructing a QueryState requires reading the DefaultQueryFilters resource, but get_state can be called from transmute with no access.

I believe it's possible to resolve those issues, but I don't think those solutions belong in this PR.

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
• 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 get_state being impossible to implement, 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?

[cBournhonesque]

The PR looks good, but I don't fully understand IterQueryData.

The trait says

A [`QueryData`] for which instances may be alive for different entities concurrently.

But when doing iter_mut, we never have two 2 items concurrently right? since we would iterate through them one by one. So even if we have a nested query which accesses multiple entities, there would still not be a collision?

[chescock]

But when doing iter_mut, we never have two 2 items concurrently right? since we would iterate through them one by one.

Nope, you can keep the old items around! The lifetime in fn next(&mut self) isn't connected to the Item type, so later calls don't invalidate earlier items. That's how things like collect() work.

[cBournhonesque]

How come we can only transmute to a SingleEntityQueryData?

[cBournhonesque]

Ah i see in the PR description that this could be relaxed in the future

[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<_>>();