How to sort arbitrary queries by component hash? Getting confused by complex types and lifetime around WorldQuery, Item, etc

[bcolloran]

I'm working on a project in which I have a number of systems that I need to run deterministically, and to make sure that that's happening, I'm sorting each system (the need for this is nicely described by @alice-i-cecile at #1470 and #2480). Currently, I'm manually using the workaround described here: https://github.com/bevyengine/bevy/pull/1817/files.

Here's a sketch of what I'm doing -- basically, since everything needs to be deterministic, I have impl'ed a "checksum" method that I use to verify that the tuples of components operated on by the queries requiring determinism are identical between system invocations that start from the same world state (this checksum method can also easily be extended to tuples). This checksum method gives me a u64, and indeed the same collection of u64s every time a system is run from the same starting point, and so it is suitable to be used as a sort key for my use case. The (compiling) code below is a simplified version of the pattern:

trait MyChecksum {
    fn checksum(&self) -> u64;
}

#[derive(Component)]
pub struct A(pub f32);
// need to work on a derive macro...
impl MyChecksum for A {
    fn checksum(&self) -> u64 {
        0
    }
}

#[derive(Component)]
pub struct B(pub f32);
impl MyChecksum for B {
    fn checksum(&self) -> u64 {
        0
    }
}

#[derive(Component)]
pub struct C(pub f32);
impl MyChecksum for C {
    fn checksum(&self) -> u64 {
        0
    }
}

// in the real implementation I provide impls for tuples of other sizes, 
// and for tuples including `Mut<_>` wrapped components.
impl<P, Q, R> MyChecksum for (&P, &Q, &R)
where
    P: MyChecksum,
    Q: MyChecksum,
    R: MyChecksum,
{
    fn checksum(&self) -> u64 {
        self.0.checksum() + self.1.checksum() + self.2.checksum()
    }
}

pub fn test_system_that_i_have(mut query: Query<(&A, &B, &C)>) {
    let mut components_vec = query.iter_mut().collect::<Vec<_>>();
    components_vec.sort_by(|a, b| (a.checksum()).partial_cmp(&b.checksum()).unwrap());

    for (a, b, c) in components_vec {
        // do stuff
    }
}

This workaround is really not so bad, but what I'd really like to do is just be able to write one generic function that I can use to sort any query, like this:

pub fn test_system_that_i_want(mut query: Query<(&A, &B, &C)>) {
    for (a, b, c) in sort_query_by_checksum(query) {
        // do stuff
    }
}

(and additionally, I'd like it be wrapped in one function so that I can add additional logging and verification logic in the sorting function).

Getting that function sort_query_by_checksum to work has really got me stymied, though. When I manually sort the component tuples, as in the test_system_that_i_have example, it works fine -- I can figure out what I need to do because I am given concrete tuples by the query.iter_mut() call -- intellisense tells me that components_vec has type Vec<(&A, &B, &C)>. But I don't get anything nearly so concrete when I try to do this in a function where I'm passing in a query.

Here is my thrashing attempt to get something working:

fn sort_query_by_checksum<'w, 'a, Q, F>(
    mut query: Query<'w, 'a, Q, F>,
) -> Vec<<Q as WorldQuery>::Item<'a>>
where
    Q: WorldQuery,
    F: ReadOnlyWorldQuery,
    <Q as WorldQuery>::Item<'a>: MyChecksum,
    'w: 'a,
{
    let mut components_vec = query.iter_mut().collect::<Vec<_>>();
    components_vec.sort_by(|a, b| (a.checksum()).partial_cmp(&b.checksum()).unwrap());
    components_vec
}

This gives the following compiler warning...

error[E0515]: cannot return value referencing function parameter `query`
   --> src/player_systems.rs:445:5
    |
443 |     let mut components_vec = query.iter().collect::<Vec<_>>();
    |                              ------------ `query` is borrowed here
444 |     components_vec.sort_by(|a, b| (a.checksum()).partial_cmp(&b.checksum()).unwrap());
445 |     components_vec
    |     ^^^^^^^^^^^^^^ returns a value referencing data owned by the current function

As a newbie to Rust I'm still struggling with lifetimes and borrowing. I think I need to be passing in a mutable reference to the query, and then passing out a vec of mutable references to the items from the query, but as I try to add this in, the type signatures and lifetimes get totally nuts. I'm sure there is an easier way, and that someone who understands (a) rust in general, and (b) the lifetimes of Query, WorldQuery, etc, could make much more sense of this than I can.

Looking forward to learning from you all, thanks!!

[SkiFire13]

Looking at the definition of Query::iter_mut in the docs you can see that it returns a QueryIter<'_, 's, Q, F>. '_ means an "inferred" lifetime, and since the only generic lifetime in the input is the one of the &mut self, it means that they're the same. Then, looking at the docs for QueryIter<'w, 's, Q, F> you can see that its Iterator implementation yields <Q as WorldQuery>::Item<'w>, thus the lifetime of the elements is the first one in the QueryIter type, and tracing back to the .iter_mut call the first lifetime was the lifetime of &mut self. Thus that's the only lifetime you need to care about and use. Here's the result (with the bonus of using sort_by_key instead of sort_by):

fn sort_query_by_checksum<'a, Q, F>(
    query: &'a mut Query<'_, '_, Q, F>,
) -> Vec<<Q as WorldQuery>::Item<'a>>
where
    Q: WorldQuery,
    F: ReadOnlyWorldQuery,
    <Q as WorldQuery>::Item<'a>: MyChecksum,
{
    let mut components_vec = query.iter_mut().collect::<Vec<_>>();
    components_vec.sort_by_key(|a| a.checksum());
    components_vec
}

You can also replace <Q as WorldQuery>::Item<'a> with QueryItem<'a, Q> to make it look better and remove the '_, '_, since the meaning is the same.

One thing to note though is that this is't really efficient, especially if your query start to return lot of elements. I'm not sure if there are valid alternatives though. Maybe with the future you could have a sorted index?

[bcolloran]

Thanks so much @SkiFire13 , this is tremendously helpful! :-D