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)] 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#5.

Testing

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

(Red is main, yellow is chescock#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(), 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:

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);
// 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:?}");

[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.

[hymm]

This needing to be moved is suspicious. Can you not call get multiple times on an immutable borrow of the world anymore?

[chescock]

Not in this case, but I promise it's not as bad as it looks!

QueryState::get() takes &mut self so that it can call update_archetypes(). And NameOrEntity uses #[derive(QueryData)], which now generates an item struct that includes the 's lifetime. So the item is treated as having a mutable borrow of the QueryState.

This doesn't affect Query::get(), which always takes &self.

And this doesn't affect concrete types where QueryData::Item doesn't include 's. That is, it only affects generic parameters, #[derive(QueryData)] types, and (in the future) FilteredEntityRef|Mut. If we use let mut query = world.query::<(Option<&Name>, Entity)>(); here, it will compile.

And there is a workaround in this case, which is to call query.update_archetypes() manually and then use query.get_manual(), which takes &self and so will allow overlapping borrows.

But, yeah, it is a breaking change here.

[hymm]

After digesting this for a bit, maybe it isn't such a bad change. get takes an &mut self so hopefully it won't be that surprising for new users if you can't get multiple items out of it. And the workaround is more performant, so leading users to that is probably better.

But we should document the workaround on QueryState::get and in the migration guide.

I haven't done a full review yet, but hopefully will get to that soon.

[Victoronz]

Does this not also affect query iteration? As in, for the affected types, any item would have to be dropped before another can be retrieved?
(Like what happens in QueryManyIter::fetch_next(), where item lifetimes are shrunk like this intentionally)

[chescock]

Does this not also affect query iteration?

No, it doesn't. Query has a &'s QueryState, so we can copy that reference and hand it out with the full 's lifetime. This change only affects methods that take &mut QueryState.

(The issue with things like QueryManyIter::fetch_next() and Query::iter_mut() are that the UnsafeWorldCell<'w> is acting like a &'w mut, so we can't copy it and instead have to reborrow it with a shorter 'w lifetime.)

[Victoronz]

Okay, so this does not present an issue within individual iterators.
However, this will be a breaking change to more methods of QueryState than what you've currently listed:

• get
• get_single (no checked manual version)
• get_many (no manual version)
• single (no manual version)
• iter
• iter_many
• iter_combinations (no manual version)
• par_iter (no manual version)
• and naturally, any method that takes at least &'a QueryState while the query items returned by one of the above is live.

The lack of manual versions can be considered API holes for some of the above, however not all of them (f.e. single)

Do I understand correctly?

If so, I think the ability to mix these methods of QueryState contains important use cases, and I would see those as more valuable than this change in its current form :/

[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.

[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 Introduce methods on QueryState to obtain a Query #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 Always collect() when using QueryIterMany::sort methods. #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 Remove the need for collect_inner() on QuerySortedManyIter #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]

How come the extra bound is necessary here?

[chescock]

Traversal is implemented in terms of get_components(), which works by creating D::State in a local, so the returned item can't borrow from the state.

It shouldn't prevent anything useful, since the only QueryData that prohibits is FilteredEntity(Ref|Mut), and those don't have any access at all when default-initialized. Note that it does allow Entity(Ref|Mut), so you can get full access, just not partial!

We could loosen the requirement on Traveral in the future if we need to. The traversal code only needs the query items during traversal, so it would be possible to create D::State in the traversal code to borrow from. But that means making a different version of get_components(), and that's a lot of extra code to support a case that isn't actually useful in the first place.

[cBournhonesque]

Ah I was confused because I was only thinking of traversal in the context of observers, where the Traversal is usually just a component so we could use world.get::<T> instead of get_components.

But I see that in some cases (WindowTraversal) the traversal query data is not a component

[cBournhonesque]

Conceptually it makes sense to me that the Item can borrow from both the UnsafeWorldCell or the State.
I don't have a strong opinion on the sort changes, I actually couldn't find them in the diff

[chescock]

Here is the change to do a collect() as part of the sort methods.

[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