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#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(), 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.

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

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