Combine `Query` and `QueryLens` using a type parameter for state #18162

chescock · 2025-03-05T17:25:51Z

Objective

Make QueryLens easier to use by allowing query methods to be called on it directly instead of needing to call the query() method first.

Introduce owned query iterators, so that it's possible for methods to return iterators constructed from Query::transmute or Query::join.

Solution

Make Query and the various iterator types generic, where normal queries use borrowed &QueryState<D, F> and QueryLens uses owned QueryState<D, F>. Introduce a trait QueryStateBorrow: Borrow<QueryState> to abstract over the two types.

Have Query use a default type for the state so that it continues to work without specifying it explicitly.

Note that the 'state lifetime on Query was only used in &'state QueryState, so it is now only used in the default type parameter. It still needs to be part of the Query type in order to be referenced in the default type, so we need a PhantomData so that it's actually used. Another option here would be to make Query a type alias for some new named type, but I think the user experience will be better with a default type parameter than with a type alias.

Testing

I used cargo-show-asm to verify that the assembly of Query::iter did not change.

This rust code:

use crate::prelude::*;
#[derive(Component)]
struct C(usize);

#[unsafe(no_mangle)]
#[inline(never)]
pub fn call_query_iter(query: Query<&C>) -> usize {
    let mut total = 0;
    for c in query.iter() {
        total += c.0;
    }
    total
}

Run with

cargo asm -p bevy_ecs --lib --simplify call_query_iter

Results in the same asm before and after the change

call_query_iter:
	push rsi
	push rdi
	push rbx
	mov rdx, qword ptr [rcx]
	mov rcx, qword ptr [rcx + 8]
	mov r8, qword ptr [rdx + 248]
	mov rax, qword ptr [rdx + 256]
	lea r9, [r8 + 4*rax]
	cmp byte ptr [rdx + 264], 0
	je .LBB1258_2
	xor r11d, r11d
	xor r10d, r10d
	xor esi, esi
	xor eax, eax
	jmp .LBB1258_6
.LBB1258_2:
	mov esi, 8
	xor r11d, r11d
	xor r10d, r10d
	xor edi, edi
	xor eax, eax
	jmp .LBB1258_15
.LBB1258_3:
	xor r11d, r11d
.LBB1258_4:
	xor esi, esi
.LBB1258_5:
	mov edi, esi
	inc rsi
	add rax, qword ptr [r11 + 8*rdi]
.LBB1258_6:
	cmp rsi, r10
	jne .LBB1258_5
.LBB1258_7:
	cmp r8, r9
	je .LBB1258_22
	mov r10d, dword ptr [r8]
	add r8, 4
	mov r11, qword ptr [rcx + 328]
	lea rsi, [r10 + 8*r10]
	mov r10, qword ptr [r11 + 8*rsi + 16]
	test r10, r10
	je .LBB1258_7
	lea r11, [r11 + 8*rsi]
	mov rsi, qword ptr [rdx + 272]
	cmp qword ptr [r11 + 64], rsi
	jbe .LBB1258_3
	mov rdi, qword ptr [r11 + 56]
	mov rsi, qword ptr [rdi + 8*rsi]
	test rsi, rsi
	je .LBB1258_3
	mov r11, qword ptr [r11 + 24]
	not rsi
	lea rsi, [rsi + 2*rsi]
	shl rsi, 4
	mov r11, qword ptr [r11 + rsi + 16]
	jmp .LBB1258_4
.LBB1258_12:
	xor r11d, r11d
.LBB1258_13:
	mov rsi, qword ptr [rsi + 80]
	xor edi, edi
.LBB1258_14:
	mov rbx, rdi
	shl rbx, 4
	mov ebx, dword ptr [rsi + rbx + 8]
	inc rdi
	add rax, qword ptr [r11 + 8*rbx]
.LBB1258_15:
	cmp rdi, r10
	jne .LBB1258_14
.LBB1258_16:
	cmp r8, r9
	je .LBB1258_22
	mov r10d, dword ptr [r8]
	add r8, 4
	mov rsi, qword ptr [rcx + 160]
	lea r11, [r10 + 4*r10]
	shl r11, 5
	mov r10, qword ptr [rsi + r11 + 88]
	test r10, r10
	je .LBB1258_16
	add rsi, r11
	mov r11d, dword ptr [rsi + 148]
	mov rdi, qword ptr [rcx + 328]
	lea rbx, [r11 + 8*r11]
	mov r11, qword ptr [rdx + 272]
	cmp qword ptr [rdi + 8*rbx + 64], r11
	jbe .LBB1258_12
	lea rdi, [rdi + 8*rbx]
	mov rbx, qword ptr [rdi + 56]
	mov r11, qword ptr [rbx + 8*r11]
	test r11, r11
	je .LBB1258_12
	mov rdi, qword ptr [rdi + 24]
	not r11
	lea r11, [r11 + 2*r11]
	shl r11, 4
	mov r11, qword ptr [rdi + r11 + 16]
	jmp .LBB1258_13
.LBB1258_22:
	pop rbx
	pop rdi
	pop rsi
	ret

…borrowed state.

Victoronz · 2025-03-05T19:36:43Z

First of all, thank you for this PR! I've been wanting to do something like this for some time!

Let's get into the details:

I've mentioned this in the discord discussion before, but I'll properly write it out here again:
If we make QueryLens a proper type, then we should not store QueryState on the stack.
QueryState is about a quarter kilobyte big, with a variable size, depending on D and F.
Today, Query is used and stored in various places, and if QueryLens becomes a type that can be used like Query, it is only natural to expect it to be passed around, stored and moved a bunch as well.
Some types even store multiple Query types within them!
Especially since the QueryStateBorrow generic propagates to the iterators, they now also have the same problem!

I think we should only store pointers to QueryState, not QueryState itself. For QueryLens this means storing a Box<QueryState>.
The API of the trait could then equal the API of &QueryState, and can live in just one place (whether that should be the trait or &QueryState I am not yet sure).

The meaning of the QueryStateBorrow trait then becomes close to Deref, instead of Borrow.
With Deref as the base of this trait, we'd have more future possibilities, because more types implement Deref over Borrow.
Examples are: Pin, LazyCell, LazyLock, various mutex guards, RefCell borrow types.
It also spares us from needing to pepper .borrow() everywhere, we can use either *, .deref(), or a direct call when applicable.

Separately, QueryStateBorrow needs to be an unsafe trait, because neither Borrow nor Deref are. An implementor could just conjure up some invalid QueryState when handing out &QueryState, which would lead to UB!

Victoronz · 2025-03-05T19:37:43Z

crates/bevy_ecs/src/query/iter.rs

    where
        D: ReadOnlyQueryData,
+        S: Copy,


We should be using Clone bounds, not Copy, it allows for stuff like Arc<QueryState> in the future.

I don't want to use Clone here because it makes it too easy to accidentally clone the entire QueryState. This is a fairly niche method, so I'm comfortable declaring that it can't be used with unusual iterators.

Note that iter() and iter_mut() will always return a QueryIter with &QueryState, since they borrow from the Query, so this only prevents calls to remaining() on the result of into_iter() with an exotic state.

Victoronz · 2025-03-05T19:51:14Z

crates/bevy_ecs/src/query/iter.rs

-    fetch: D::Fetch<'w>,
-    filter: F::Fetch<'w>,
-    query_state: &'s QueryState<D, F>,
+    fetch: <S::Data as WorldQuery>::Fetch<'w>,


I wonder whether we could simplify some of these/recover the old names with type aliases

crates/bevy_ecs/src/query/iter.rs

chescock · 2025-03-05T20:10:09Z

I think we should only store pointers to QueryState, not QueryState itself. For QueryLens this means storing a Box<QueryState>. The API of the trait could then equal the API of &QueryState, and can live in just one place (whether that should be the trait or &QueryState I am not yet sure).

The meaning of the QueryStateBorrow trait then becomes close to Deref, instead of Borrow. With Deref as the base of this trait, we'd have more future possibilities, because more types implement Deref over Borrow. Examples are: Pin, LazyCell, LazyLock, various mutex guards, RefCell borrow types. It also spares us from needing to pepper .borrow() everywhere, we can use either *, .deref(), or a direct call when applicable.

Yup, that's a good idea! I had written this before those discussions, and was trying for the smallest possible change to QueryLens, but I agree that boxing it makes sense. It costs an extra allocation, but QueryState is already full of allocations, so one more shouldn't hurt. And the diff will get smaller without the borrow() calls!

Separately, QueryStateBorrow needs to be an unsafe trait, because neither Borrow nor Deref are. An implementor could just conjure up some invalid QueryState when handing out &QueryState, which would lead to UB!

Yup, that makes sense, too. I had been thinking that we could just roll that into the safety requirements of the existing unsafe fn new(). But they'll be easier to express more clearly on the trait, and there aren't going to be so many implementations of this that we need to worry about a little more unsafe.

Use Deref instead of Borrow.

…pl EntitySetIterator for QueryIter`.

crates/bevy_ecs/src/query/fetch.rs

Victoronz · 2025-03-06T05:51:50Z

crates/bevy_ecs/src/query/state.rs

+/// Abstracts over an owned or borrowed [`QueryState`].
+///
+/// # Safety
+///
+/// This must `deref` to a `QueryState` that does not change.
+pub unsafe trait QueryStateDeref:
+    Deref<Target = QueryState<Self::Data, Self::Filter>>
+{
+    /// The [`QueryData`] for this `QueryState`.
+    type Data: QueryData;
+
+    /// The [`QueryFilter`] for this `QueryState`.
+    type Filter: QueryFilter;
+
+    /// The type returned by [`Self::storage_ids`].
+    type StorageIter: Iterator<Item = StorageId> + Clone + Default;
+
+    /// A read-only version of the state.
+    type ReadOnly: QueryStateDeref<
+        Data = <Self::Data as QueryData>::ReadOnly,
+        Filter = Self::Filter,
+    >;
+
+    /// Iterates the storage ids that this [`QueryState`] matches.
+    fn storage_ids(&self) -> Self::StorageIter;
+
+    /// Borrows the remainder of the [`Self::StorageIter`] as an iterator
+    /// usable with `&QueryState<D, F>`.
+    fn reborrow_storage_ids(
+        storage_iter: &Self::StorageIter,
+    ) -> iter::Copied<slice::Iter<StorageId>>;
+
+    /// Converts this state to a read-only version.
+    fn into_readonly(self) -> Self::ReadOnly;
+}


Hmm, I said earlier I wasn't sure where the API should live, but I am now certain that it should be on &QueryState.
By putting it on the trait, we would need to describe the correct implementation for each method in the safety contract, or we cannot use these methods ourselves.

This trait should instead have no API for now, and would then only have to describe its supertrait and QueryData/QueryFilter.

I delegated everything I could, but these methods are here because the types vary between owned and borrowed versions. The S::StorageIter is stored at the same time as the S, so it can't be slice::Iter for owned QueryData. And into_readonly needs to return an owned QueryState, since there's nothing left to borrow from.

Right, it seems QueryIterationCursor code doesn't compile when attempting to borrow the storage iterator from a deref.
However, this means that we are introducing some storage iterator clones into hot iteration code.
We should try to avoid that if possible, these clones shouldn't be necessary!
That does mean changing up the code to allow for this, which I hope isn't too difficult...
The iteration code of QueryIterationCursor is important enough that changing it involves benchmarks/proving it isn't being regressed.

With into_readonly, we run into the problem of needing to cast the original S instead of just the deref.
If we need to put the cast after the deref, then we can just wrap S with a struct that retains the deref, but appends the cast in its own impl!

For that, we don't need a dedicated method on QueryStateDeref, it can be done in Query::into_readonly directly.
If we later want it on QueryStateDeref, then we can put it there as a provided method! Though the more private solution is preferable at first I think.
This would change the returned S type of Query::into_readonly into ReadOnly<S>.

This solution would also transfer to owned transmutes, which into_readonly is essentially a simple case of.

The technique of wrapping an S in a ReadOnly<S>/Transmuted<S> might actually be a nice speedup for QueryState::transmute/Query::transmute_lens in general, because we can then avoid creating a new QueryState like we do now!
We can just keep the "access is a subset" check, then cast S this way!

However, this means that we are introducing some storage iterator clones into hot iteration code.
We should try to avoid that if possible, these clones shouldn't be necessary!
That does mean changing up the code to allow for this, which I hope isn't too difficult...
The iteration code of QueryIterationCursor is important enough that changing it involves benchmarks/proving it isn't being regressed.

Yup, the clones are unfortunate! They only occur for QueryLens::into_iter(), though. Calling into_iter() on an ordinary Query, or calling iter() or iter_mut() on anything, will use &QueryState. That still uses a slice iterator, so this won't change the behavior of any existing code. And if you're consuming a QueryLens then you're already spending allocations on the QueryState, so one more shouldn't be too bad.

I really don't want to make any big changes to QueryIterationCursor in this PR, exactly because it's important for performance! If performance of QueryLens::into_iter() becomes a problem, we can do a more targeted follow-up to rework QueryIterationCursor.

For that, we don't need a dedicated method on QueryStateDeref, it can be done in Query::into_readonly directly. ... This would change the returned S type of Query::into_readonly into ReadOnly<S>.

Oh, wrapping the type is a clever idea! ... Hmm, it would mean that Query<&mut T>::as_readonly() is no longer the same type as Query<&T>, though, since the state type would differ. And I really don't want to make a breaking change like that as part of this PR!

The technique of wrapping an S in a ReadOnly<S>/Transmuted<S> might actually be a nice speedup for QueryState::transmute/Query::transmute_lens in general, because we can then avoid creating a new QueryState like we do now!

Yeah, I think there are some good opportunities in that space! It's a little harder than that because you also need to store a new D::State and F::State for the new D and F. Maybe Transmuted<S> could hold new versions of those and delegate the rest to a wrapped state? Although then you're definitely not a Deref<Target=QueryState>.

However, this means that we are introducing some storage iterator clones into hot iteration code.
We should try to avoid that if possible, these clones shouldn't be necessary!
That does mean changing up the code to allow for this, which I hope isn't too difficult...
The iteration code of QueryIterationCursor is important enough that changing it involves benchmarks/proving it isn't being regressed.

Yup, the clones are unfortunate! They only occur for QueryLens::into_iter(), though. Calling into_iter() on an ordinary Query, or calling iter() or iter_mut() on anything, will use &QueryState. That still uses a slice iterator, so this won't change the behavior of any existing code. And if you're consuming a QueryLens then you're already spending allocations on the QueryState, so one more shouldn't be too bad.

I really don't want to make any big changes to QueryIterationCursor in this PR, exactly because it's important for performance! If performance of QueryLens::into_iter() becomes a problem, we can do a more targeted follow-up to rework QueryIterationCursor.

True, let us leave it for a follow-up then.

For that, we don't need a dedicated method on QueryStateDeref, it can be done in Query::into_readonly directly. ... This would change the returned S type of Query::into_readonly into ReadOnly<S>.

Oh, wrapping the type is a clever idea! ... Hmm, it would mean that Query<&mut T>::as_readonly() is no longer the same type as Query<&T>, though, since the state type would differ. And I really don't want to make a breaking change like that as part of this PR!

Hmm, I don't quite follow, how would those differ?
But for now, what could be done in this PR is add an unsafe cast method to QueryStateDeref (mirroring ptr::cast), which would serve the purpose we'd need it for. The impl for Box<QueryState> should then also be changed into a into_raw -> cast -> from_raw roundtrip.

The technique of wrapping an S in a ReadOnly<S>/Transmuted<S> might actually be a nice speedup for QueryState::transmute/Query::transmute_lens in general, because we can then avoid creating a new QueryState like we do now!

Yeah, I think there are some good opportunities in that space! It's a little harder than that because you also need to store a new D::State and F::State for the new D and F. Maybe Transmuted<S> could hold new versions of those and delegate the rest to a wrapped state? Although then you're definitely not a Deref<Target=QueryState>.

Right! I forgot that ReadOnly has the equal state restriction. Maybe something can be done to add a similar restriction here.

Hmm, it would mean that Query<&mut T>::as_readonly() is no longer the same type as Query<&T>, though, since the state type would differ.

Hmm, I don't quite follow, how would those differ?

If I understood your suggestion, then Query<&mut T, (), &QS>::as_readonly() would return Query<&T, (), ReadOnly<&QS>>. But Query<&T> is Query<&T, (), &QS>, and &QS != ReadOnly<&QS>. So a user couldn't pass the result of as_readonly() to a method that took Query<&T>, which I believe is the main use case for that method today.

But for now, what could be done in this PR is add an unsafe cast method to QueryStateDeref (mirroring ptr::cast), which would serve the purpose we'd need it for. The impl for Box<QueryState> should then also be changed into a into_raw -> cast -> from_raw roundtrip.

I don't think I see how to implement that. Do you mean exposing an equivalent of QueryState::as_transmuted_state instead of QueryState::as_readonly? That would still wind up needing an associated type, but it would need to be generic over NewD, NewF. Since the only uses are as_readonly and as_nop, and as_nop never needs an owned state, it seems simpler to only expose as_readonly.

I can change the Box<QueryState> implementation to do a pointer cast. I like that the current impl doesn't need unsafe, but the safety there isn't any worse than the &QueryState cast, so we might as well avoid reallocating!

Victoronz · 2025-03-10T21:45:34Z

Quick note:
Github pings in commit messages should be avoided, because it means Github will notify that person each time the repo is cloned!
(Which can be a lot...)
Last time I heard about this happening, Github didn't have a fine-grained way of turning this kind of ping off either.
This might only apply to commit titles and not the full message, I'm not sure.

chescock · 2025-03-10T22:57:19Z

Github pings in commit messages should be avoided, because it means Github will notify that person each time the repo is cloned!

Oh, sorry! I was just trying to give you proper credit :). I didn't actually know that would turn into a ping!

Do you want me to force-push a new message, or is it okay here because nobody is cloning my personal fork and the commit will be squashed if it gets merged?

Victoronz · 2025-03-10T23:17:30Z

Github pings in commit messages should be avoided, because it means Github will notify that person each time the repo is cloned!

Oh, sorry! I was just trying to give you proper credit :). I didn't actually know that would turn into a ping!

Do you want me to force-push a new message, or is it okay here because nobody is cloning my personal fork and the commit will be squashed if it gets merged?

I think its fine here because of the squash, just noting for the future; some repos don't do that!
Thanks for the credit :)

…ding to allocate a new box.

…tate # Conflicts: # crates/bevy_ecs/src/query/iter.rs # crates/bevy_ecs/src/query/par_iter.rs # crates/bevy_ecs/src/system/query.rs

`cargo fmt`.

…tate # Conflicts: # crates/bevy_ecs/src/query/par_iter.rs

github-actions · 2025-05-05T17:25:20Z

It looks like your PR is a breaking change, but you didn't provide a migration guide.

Please review the instructions for writing migration guides, then expand or revise the content in the migration guides directory to reflect your changes.

chescock · 2025-05-07T15:30:19Z

I added a migration guide for the deprecated QueryLens::query() method. I think that was the only breaking change, but let me know if I missed something!

chescock added 3 commits March 5, 2025 09:55

Make Query generic over the State type so it can use either owned or …

e943852

…borrowed state.

Make QueryLens a type alias.

f7ff2b3

Update doc link now that StorageId is pub.

0a68186

chescock mentioned this pull request Mar 5, 2025

Store QueryState in Query as a Cow and remove distinct QueryLens type #15848

Closed

chescock added A-ECS Entities, components, systems, and events C-Usability A targeted quality-of-life change that makes Bevy easier to use S-Needs-Review Needs reviewer attention (from anyone!) to move forward labels Mar 5, 2025

Victoronz reviewed Mar 5, 2025

View reviewed changes

crates/bevy_ecs/src/query/iter.rs Show resolved Hide resolved

chescock added 5 commits March 5, 2025 15:55

Box QueryState.

49ae0f3

Use Deref instead of Borrow.

Rename QueryStateBorrow to QueryStateDeref.

2d2d78d

Make QueryStateDeref an unsafe trait.

4c367db

Introduce a TrustedEntityQueryData trait so we can make a blanket `im…

556eced

…pl EntitySetIterator for QueryIter`.

Doc fixes.

c81671d

Victoronz reviewed Mar 6, 2025

View reviewed changes

crates/bevy_ecs/src/query/fetch.rs Outdated Show resolved Hide resolved

Victoronz reviewed Mar 6, 2025

View reviewed changes

Better docs for TrustedEntityQueryData from @Victoronz

e94970c

chescock added 5 commits March 11, 2025 13:09

Use pointer casts for Box<QueryState>::into_readonly() to avoid nee…

39c8c16

…ding to allocate a new box.

Merge remote-tracking branch 'remotes/origin/main' into query-owned-s…

3d2bba9

…tate # Conflicts: # crates/bevy_ecs/src/query/iter.rs # crates/bevy_ecs/src/query/par_iter.rs # crates/bevy_ecs/src/system/query.rs

Fix other references to TrustedEntityBorrow.

ddd30a5

`cargo fmt`.

Rename TrustedEntityQueryData to EntityEquivalentQueryData.

0e61022

Merge remote-tracking branch 'remotes/origin/main' into query-owned-s…

3e83269

…tate # Conflicts: # crates/bevy_ecs/src/query/par_iter.rs

alice-i-cecile added M-Needs-Migration-Guide A breaking change to Bevy's public API that needs to be noted in a migration guide X-Contentious There are nontrivial implications that should be thought through labels May 5, 2025

alice-i-cecile added C-Code-Quality A section of code that is hard to understand or change D-Complex Quite challenging from either a design or technical perspective. Ask for help! labels May 5, 2025

chescock added 3 commits May 7, 2025 10:09

Add migration guide.

de74769

Merge branch 'main' into query-owned-state

1751fa0

bevy_platform_support was renamed bevy_platform.

fccca4d

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Combine `Query` and `QueryLens` using a type parameter for state #18162

Combine `Query` and `QueryLens` using a type parameter for state #18162

chescock commented Mar 5, 2025

Victoronz commented Mar 5, 2025 •

edited

Loading

Victoronz Mar 5, 2025

chescock Mar 5, 2025

Victoronz Mar 5, 2025

chescock commented Mar 5, 2025

Victoronz Mar 6, 2025

chescock Mar 6, 2025

Victoronz Mar 10, 2025

Victoronz Mar 11, 2025 •

edited

Loading

chescock Mar 11, 2025

Victoronz Mar 11, 2025 •

edited

Loading

chescock Mar 11, 2025

Victoronz commented Mar 10, 2025 •

edited

Loading

chescock commented Mar 10, 2025

Victoronz commented Mar 10, 2025 •

edited

Loading

github-actions bot commented May 5, 2025

chescock commented May 7, 2025

Combine Query and QueryLens using a type parameter for state #18162

Are you sure you want to change the base?

Combine Query and QueryLens using a type parameter for state #18162

Conversation

chescock commented Mar 5, 2025

Objective

Solution

Testing

Victoronz commented Mar 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chescock commented Mar 5, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Victoronz Mar 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Victoronz Mar 11, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Victoronz commented Mar 10, 2025 • edited Loading

chescock commented Mar 10, 2025

Victoronz commented Mar 10, 2025 • edited Loading

github-actions bot commented May 5, 2025

chescock commented May 7, 2025

Combine `Query` and `QueryLens` using a type parameter for state #18162

Combine `Query` and `QueryLens` using a type parameter for state #18162

Victoronz commented Mar 5, 2025 •

edited

Loading

Victoronz Mar 11, 2025 •

edited

Loading

Victoronz Mar 11, 2025 •

edited

Loading

Victoronz commented Mar 10, 2025 •

edited

Loading

Victoronz commented Mar 10, 2025 •

edited

Loading