Let query items borrow from query state to avoid needing to clone

[chescock]

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)] should generate ReleaseStateQueryData implementations, but that is not currently implemented in this PR.

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)

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:

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:

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(), get_many(), iter(), and iter_many(). To fix the errors, first call QueryState::update_archetypes(), and then replace the calls with the corresponding _manual methods, which take &self.

Before:

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:

let mut state: QueryState<_, _> = ...;
state.update_archetypes(world);
let d1 = state.get_manual(world, e1);
let d2 = state.get_manual(world, e2);
println!("{d1:?}");
println!("{d2:?}");

[Victoronz]

IIUC, the lack of ReleaseStateQueryData impls for FilteredEntityRef|Mut means they can no longer be used in the sort methods. This would mean removing support for sorting with untyped QueryData, which is important for some people. Is this change possible without removing that feature?

[chescock]

IIUC, the lack of ReleaseStateQueryData impls for FilteredEntityRef|Mut means they can no longer be used in the sort methods. This would mean removing support for sorting with untyped QueryData, which is important for some people. Is this change possible without removing that feature?

It might be. The issue is that the L::Items borrow from the state, and they get included in the returned Iterator<Item=Entity>. But they aren't actually used after the sort method returns. The simplest fix would be to collect() entity_iter into a Vec to drop the L::Items, but I saw you had tried to avoid that in #13443.

For what it's worth, I don't think FilteredEntityRef|Mut will actually work with the sort() methods right now, which is why I hadn't been worried about it. They do transmute_filtered::<(L, Entity), F>(), and calling set_access() on (FilteredEntityRef, Entity) doesn't pass it down to the FilteredEntityRef. Also, transmute_filtered() calls set_access() before update_component_access(), so the access is empty anyway.

[Victoronz]

IIUC, the lack of ReleaseStateQueryData impls for FilteredEntityRef|Mut means they can no longer be used in the sort methods. This would mean removing support for sorting with untyped QueryData, which is important for some people. Is this change possible without removing that feature?

It might be. The issue is that the L::Items borrow from the state, and they get included in the returned Iterator<Item=Entity>. But they aren't actually used after the sort method returns. The simplest fix would be to collect() entity_iter into a Vec to drop the L::Items, but I saw you had tried to avoid that in #13443.

Yeah, an additional collect would regress sort performance. That would make this change a perf trade-off, not a pure gain. By how much, would have to be benchmarked.

For what it's worth, I don't think FilteredEntityRef|Mut will actually work with the sort() methods right now, which is why I hadn't been worried about it. They do transmute_filtered::<(L, Entity), F>(), and calling set_access() on (FilteredEntityRef, Entity) doesn't pass it down to the FilteredEntityRef. Also, transmute_filtered() calls set_access() before update_component_access(), so the access is empty anyway.

The tuple behavior is a bug: #14349. Entity, EntityLocation and &Archetype should always be available QueryData, regardless of access. That access is not propagated to a transmuted FilteredEntityRef|Mut sounds like another bug, does FilteredEntityRef|Mut even work with transmutes in the first place?

[chescock]

Yeah, an additional collect would regress sort performance. That would make this change a perf trade-off, not a pure gain. By how much, would have to be benchmarked.

Yeah, there are definitely tradeoffs here! Given that sort doesn't actually work with FilteredEntityRef right now, I'd like to leave the ReleaseStateQueryData constraint in place on those methods. If we decide to merge this and we fix #14349, then we can revisit the constraint.

(It feels like there must be some way to get the best of both options by dropping the L::Items in place before the method returns without allocating a new Vec, but I don't yet see how to make it work. ManuallyDrop<T> still captures lifetimes.)

That access is not propagated to a transmuted FilteredEntityRef|Mut sounds like another bug, does FilteredEntityRef|Mut even work with transmutes in the first place?

Sorry, I was wrong; it works fine. I confused the local component_access, which was empty, with self.component_access, which gets passed to set_access.

[Victoronz]

To note, I think super let could also solve the kind of problem this PR tackles, once we get some form of it in Rust.

[SkiFire13]

note, I think super let could also solve the kind of problem this PR tackles, once we get some form of it in Rust.

I'm not sure how this would help avoiding to clone the Access in FilteredEntityRef. Do you have some examples?

[chescock]

Okay, I think I've finally fixed the problems that were brought up! I've updated this PR, and it should be reviewable now. The issues had been:

1. We had been missing _manual versions of some methods on QueryState that made the workaround for Query::get() not always available. Following #15858, it's now possible to rewrite all of them as state.foo(world, params) => state.query_manual(world).foo_inner(params). And in many cases, you can even call state.query(world) once and re-use the Query!

2. The sort methods required ReleaseStateQueryData, so were not usable with FilteredEntityRef. In #16844, we started doing a collect() in the sort methods on QueryManyIter, so I've added a collect() to the sort methods on Query to match. (If we ever need the performance back, we can try something unsafe like #16650.)

3. ReleaseStateQueryData was not automatically implemented for #[derive(QueryData)] types. That wasn't a major issue when I first opened this PR, but the Traversal trait requires ReleaseStateQueryData and is now being used with #[derive(QueryData)] types. I've updated the derive macro to also derive ReleaseStateQueryData when possible.

[cBournhonesque]

I also like how the code for FilteredEntityRef because more similar to the one for FilteredResources, where you need to provide &'s Access to create the SystemParam. The FilteredResource is a SystemParam so the Item resource already has both 'w and 's lifetimes, I don't see a reason why it shouldn't be the same for QueryItems

[Victoronz]

IIUC, this change will mean that the consuming methods will not work for what is not ReleaseStateQueryData. Given that the current default way of iterating over Query uses a consuming method to iterate, users will effectively not be able to use the default for loop over FilteredEntityRef and other non-ReleaseStateQueryData types.

Example of what fails:

pub fn foo(query: Query<FilteredEntityRef>) {
    for entity_ref in query {
        // ...
    }
}

Note: I say "default" as in, what we intend for users to reach for first. Using &query in the for loop should work, but that is not obvious to an unsuspecting user, and becomes inconsistent with for loops elsewhere.

[chescock]

IIUC, this change will mean that the consuming methods will not work for what is not ReleaseStateQueryData.

That will still work fine! Query has an &'s QueryState, and the items will borrow from that.

The issue is if you create a QueryState locally, such as in a QueryLens, the items will borrow from it. So you'll no longer be able to return FilteredEntityRefs created from a local lens, like:

fn foo(query: Query<&Transform>) -> FilteredEntityRef {
    let lens = query.transmute_lens_inner::<FilteredEntityRef>();
    let query = lens.query_inner();
    query.single_inner().unwrap()
}

If we also do #18162, then the new consuming methods on QueryLens won't be able to work with non-ReleaseStateQueryData. But those methods don't exist at all today, so that' not a breaking change :).

[Victoronz]

IIUC, this change will mean that the consuming methods will not work for what is not ReleaseStateQueryData.

That will still work fine! Query has an &'s QueryState, and the items will borrow from that.

The issue is if you create a QueryState locally, such as in a QueryLens, the items will borrow from it. So you'll no longer be able to return FilteredEntityRefs created from a local lens, like:

You're right, I was missing that context.

If we also do #18162, then the new consuming methods on QueryLens won't be able to work with non-ReleaseStateQueryData. But those methods don't exist at all today, so that' not a breaking change :).

Even if its not a breaking change, it is still inconsistent!

[alice-i-cecile]

Really excellent work. This is extremely complex, but you've done a great job navigating that in your code, docs, review comments and PR description. Thanks!