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.