Combine Query and QueryLens using a type parameter for state, but don't introduce owned iterators

[chescock]

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.

Split out from #18162. I thought this might be easier to review separately, but if we do merge #18162 then this PR will no longer be necessary.

Solution

Make Query and the various iterator types generic, where normal queries use borrowed &QueryState<D, F> and QueryLens uses owned Box<QueryState<D, F>>.

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.

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)]
pub 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 .LBB1856_2
	xor r11d, r11d
	xor r10d, r10d
	xor esi, esi
	xor eax, eax
	jmp .LBB1856_6
.LBB1856_2:
	xor r11d, r11d
	mov esi, 8
	xor r10d, r10d
	xor edi, edi
	xor eax, eax
	jmp .LBB1856_16
.LBB1856_3:
	xor r11d, r11d
.LBB1856_4:
	xor esi, esi
.LBB1856_5:
	mov edi, esi
	inc esi
	add rax, qword ptr [r11 + 8*rdi]
.LBB1856_6:
	cmp esi, r10d
	jne .LBB1856_5
.LBB1856_7:
	cmp r8, r9
	je .LBB1856_23
	mov r10d, dword ptr [r8]
	add r8, 4
	mov r11, qword ptr [rcx + 416]
	lea rsi, [r10 + 8*r10]
	mov r10, qword ptr [r11 + 8*rsi + 16]
	test r10, r10
	je .LBB1856_7
	lea r11, [r11 + 8*rsi]
	mov rsi, qword ptr [rdx + 272]
	cmp qword ptr [r11 + 64], rsi
	jbe .LBB1856_3
	mov rdi, qword ptr [r11 + 56]
	mov rsi, qword ptr [rdi + 8*rsi]
	test rsi, rsi
	je .LBB1856_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 .LBB1856_4
.LBB1856_13:
	xor r11d, r11d
.LBB1856_14:
	mov rsi, qword ptr [rsi + 80]
	xor edi, edi
.LBB1856_15:
	mov ebx, edi
	shl rbx, 4
	mov ebx, dword ptr [rsi + rbx + 8]
	not ebx
	inc edi
	add rax, qword ptr [r11 + 8*rbx]
.LBB1856_16:
	cmp edi, r10d
	jne .LBB1856_15
.LBB1856_17:
	cmp r8, r9
	je .LBB1856_23
	mov r10d, dword ptr [r8]
	add r8, 4
	mov rsi, qword ptr [rcx + 256]
	lea r11, [r10 + 4*r10]
	shl r11, 5
	mov r10, qword ptr [rsi + r11 + 88]
	test r10, r10
	je .LBB1856_17
	add rsi, r11
	mov r11d, dword ptr [rsi + 148]
	mov rdi, qword ptr [rcx + 416]
	lea rbx, [r11 + 8*r11]
	mov r11, qword ptr [rdx + 272]
	cmp qword ptr [rdi + 8*rbx + 64], r11
	jbe .LBB1856_13
	lea rdi, [rdi + 8*rbx]
	mov rbx, qword ptr [rdi + 56]
	mov r11, qword ptr [rbx + 8*r11]
	test r11, r11
	je .LBB1856_13
	mov rdi, qword ptr [rdi + 24]
	not r11
	lea r11, [r11 + 2*r11]
	shl r11, 4
	mov r11, qword ptr [rdi + r11 + 16]
	jmp .LBB1856_14
.LBB1856_23:
	pop rbx
	pop rdi
	pop rsi
	ret

[ElliottjPierce]

I'm broadly on board with this. I do something similar for my ecs too.

It's really hard to see the changes to all the*_inner functions that return 'w types, but I'm assuming those were just copied over. And the asm didn't change, so the impl is still correct.

The only thought I have is to change S: Deref<Target = QueryState<D, F>> to S: Borrow<QueryState<D, F>>. That would let this work with an owned state directly. In fact, you might even be able to get rid of the Box in the query lens.

[chescock]

It's really hard to see the changes to all the*_inner functions that return 'w types, but I'm assuming those were just copied over. And the asm didn't change, so the impl is still correct.

Yeah, I did put the moves in their own commit in case that helps, but it's still hard to confirm that that commit really is just a move. I wish we had better tooling for reviewing things like that!

The only thought I have is to change S: Deref<Target = QueryState<D, F>> to S: Borrow<QueryState<D, F>>. That would let this work with an owned state directly. In fact, you might even be able to get rid of the Box in the query lens.

Hah, yeah, that's what I did first :). @Victoronz convinced me that Box and Deref are better here: #18162 (comment).

[ElliottjPierce]

Hah, yeah, that's what I did first :). @Victoronz convinced me that Box and Deref are better here: #18162 (comment).

That makes a lot of since. I still hate to loose all the flexibility that Borrow would afford us, but I get the argument for deref.

[Victoronz]

While there is some follow-up to be done, I think this is a good next step!
Hope realigning the changes with main isn't too tedious.

[alice-i-cecile]

Marked as X-Controversial as it adds another generic to Query. I think that this is probably the right call, but I'm nervous about making query more complex to ease things for QueryLens.

[Victoronz]

To clarify some of the benefits this direction has:
In the past, Query-related code was constrained by the need to work with a &'state QueryState lifetime both because of the lifetime, and the lack of ownership.

This led to various cases where we could not directly create a usable Query even if we wanted to, leading to workarounds like QueryLens, and limitations to transmutes.
The inability to directly work with Query is also part of why a lot of functionality was implemented on QueryState, instead of the more appropriate Query.

Further, this direction allows us to more closely match the 'w and 's lifetimes by letting code create and immediately use owned QueryState, which will alleviate some of our double lifetime pains.
This is not a guarantee, but I suspect that we may be able to remove the get_inner methods (they serve to remove the 'self lifetime constraint), ReleaseStateQueryData (serves to declare query items as independent of 's) , and maybe even 's entirely.

I believe that such improvements surrounding 's will make it easier to mend some of the edge cases surrounding query items that depend on it. The resulting incompatibilities between query features are a headache, and a design risk for upcoming work!

Additionally, this should ease the design space for Queries-as-Entities (QueryState can now live behind various pointer types), Uncached Queries (its functionality could be folded into the generic), Relation Queries (Query is now easier to create on the fly), etc.

Overall, I think the simplifications it will allow are worth the surface complexity it introduces!

The Query implementation lore is quite a scattered puzzle, I hope these are enough (sensible) pieces to develop a feeling for the current situation.

[cBournhonesque]

As Victoronz stated above, some other features like Uncached Queries or MultiSource Queries (Relation Queries) will also need that third generic on Query, so I think it's a step in the right direction!

Currently the third generic is Deref<Target = QueryState<D, F>> = &'state QueryState<D, F>,
Ultimately it would be a trait for QueryState which would have different implementations based on whether or not the state is owned/borrowed, cached/uncached, single-source/multi-source.

It's easier to introduce this third generic for the owned/borrowed case because the code changes are small and the use-case is pretty clear.
The only reason for QueryLens to exist was that to transmute the Query you had to put the new state somewhere, and Query doesn't own the state, only borrows it. It's clearer to indicate directly in the type system that a QueryLens is simply a Query that owns its state.

[chescock]

Marked as X-Controversial

Yup, that label seems appropriate!

as it adds another generic to Query

I'll note that defaulted generic parameters are usually pretty invisible to users. I often forget that Vec<T> is secretly short for Vec<T, A = Global>!

to ease things for QueryLens

One of my goals here is to make world.query::<D>().iter() work. We could almost make that work today by having query() return a QueryLens, but without these changes plus the rest of #18162 it needs to be bound to a local variable to be usable.

So, it's not so much about the current uses of QueryLens, as about making it nicer to use QueryLens or uncached queries in new places.

I suspect that we may be able to remove the get_inner methods (they serve to remove the 'self lifetime constraint), ReleaseStateQueryData (serves to declare query items as independent of 's) , and maybe even 's entirely.

I don't think I see how this change would help with that.

The get_inner methods let you create a Query within a method and return data from it, and that's useful with ordinary state-borrowing queries even if we add new versions.

It might be possible to remove the 's lifetime entirely! We could set it equal to 'w everywhere by having SystemState::get and get_mut take &'w mut self instead of &'s mut self, and most uses would still work. But I don't think this change would make that any easier (or harder).

[alice-i-cecile]

With a bit more docs this LGTM. These arguments have convinced me!

[Victoronz]

I suspect that we may be able to remove the get_inner methods (they serve to remove the 'self lifetime constraint), ReleaseStateQueryData (serves to declare query items as independent of 's) , and maybe even 's entirely.

I don't think I see how this change would help with that.

The get_inner methods let you create a Query within a method and return data from it, and that's useful with ordinary state-borrowing queries even if we add new versions.

It might be possible to remove the 's lifetime entirely! We could set it equal to 'w everywhere by having SystemState::get and get_mut take &'w mut self instead of &'s mut self, and most uses would still work. But I don't think this change would make that any easier (or harder).

How I see it, this change gives us a new lever manipulate 's with:
Every &'state QueryState has to refer to a QueryState that was created and stored somewhere else at some point.
If this original location has more options for how it can store/give out its QueryState, we can prevent some situations where 's is unnaturally shortened.
I'd see query iterators that want to create a transient QueryState for their items as an example of that.

But for now, this is still a suspicion! It might also be that the situations I've seen were hindered by the tech debt that can be addressed as follow-ups to this PR if it goes through, or that there'd be other issues to run into.