implement the full set of sort methods on QueryIter

[Victoronz]

Objective

Currently, a query iterator can be collected into a Vec and sorted, but this can be quite unwieldy, especially when many Components are involved. The itertools crate helps somewhat, but the need to write a closure over all of QueryData
can sometimes hurt ergonomics, anywhere from slightly to strongly. A key extraction function only partially helps, as sort_by_key does not allow returning non-Copy data. sort_by does not suffer from the Copy restriction, but now the user has to write out a cmp function over two QueryData::Items when it could have just been handled by the Ord impl for the key.
sort requires the entire Iterator Item to be Ord, which is rarely usable without manual helper functionality. If the user wants to hide away unused components with a .. range, they need to track item tuple order across their function. Mutable QueryData can also introduce further complexity.
Additionally, sometimes users solely include Components /Entity to guarantee iteration order.

For a user to write a function to abstract away repeated sorts over various QueryData types they use would require reaching for the all_tuples! macro, and continue tracking tuple order afterwards.

Fixes #1470.

Solution

Custom sort methods on QueryIter, which take a query lens as a generic argument, like transmute_lens in Query.
This allows users to choose what part of their queries they pass to their sort function calls, serving as a kind of "key extraction function" before the sort call. F.e. allowing users to implement Ord for a Component, then call query.iter().sort::<OrdComponent>()

This works independent of mutability in QueryData, QueryData tuple order, or the underlying iter/iter_mut call.
Non-Copy components could also be used this way, an internal Arc<usize> being an example.
If Ord impls on components do not suffice, other sort methods can be used. Notably useful when combined with EntityRef or EntityMut.
Another boon from using underlying transmute functionality, is that with the allowed transmutes, it is possible to sort a Query with Entity even if it wasn't included in the original Query.
The additional generic parameter on the methods other than sort and sort_unstable currently cannot be removed due to Rust limitations, however their types can be inferred.

The new methods do not conflict with the itertools sort methods, as those use the "sorted" prefix.

This is implemented barely touching existing code. That change to existing code being that QueryIter now holds on to the reference to UnsafeWorldCell that is used to initialize it.
A lens query is constructed with Entity attached at the end, sorted, and turned into an iterator. The iterator maps away the lens query, leaving only an iterator of Entity, which is used by QuerySortedIter to retrieve the actual items.
QuerySortedIter resembles a combination of QueryManyIter and QueryIter, but it uses an entity list that is guaranteed to contain unique entities, and implements ExactSizeIterator, DoubleEndedIterator, FusedIterator regardless of mutability or filter kind (archetypal/non-archetypal).

The sort methods are not allowed to be called after next, and will panic otherwise. This is checked using QueryIterationCursor state, which is unique on initialization. Empty queries are an exception to this, as they do not return any item in the first place.
That is because tracking how many iterations have already passed would require regressing either normal query iteration a slight bit, or sorted iteration by a lot. Besides, that would not be the intended use of these methods.

Testing

To ensure that next being called before sort results in a panic, I added some tests. I also test that empty QueryIters do not exhibit this restriction.

The query sorts test checks for equivalence to the underlying sorts.
This change requires that Query<(Entity, Entity)> remains legal, if that is not already guaranteed, which is also ensured by the aforementioned test.

Next Steps

Implement the set of sort methods for QueryManyIter as well.

• This will mostly work the same, other than needing to return a new QuerySortedManyIter to account for iteration
  over lists of entities that are not guaranteed to be unique. This new query iterator will need a bit of internal restructuring
  to allow for double-ended mutable iteration, while not regressing read-only iteration.

The implementations for each pair of

• sort, sort_unstable,
• sort_by, sort_unstable_by,
• sort_by_key, sort_by_cached_key

are the same aside from the panic message and the sort call, so they could be merged with an inner function.
That would require the use of higher-ranked trait bounds on WorldQuery::Item<'1>, and is unclear to me whether it is currently doable.

Iteration in QuerySortedIter might have space for improvement.
When sorting by Entity, an (Entity, Entity) lens QueryData is constructed, is that worth remedying?
When table sorts are implemented, a fast path could be introduced to these sort methods.

Future Possibilities

Implementing Ord for EntityLocation might be useful.
Some papercuts in ergonomics can be improved by future Rust features:

• The additional generic parameter aside from the query lens can be removed once this feature is stable:
  Fn -> impl Trait (impl Trait in Fn trait return position)
• With type parameter defaults, the query lens generic can be defaulted to QueryData::Item, allowing the sort methods
  to look and behave like slice::sort when no query lens is specified.
• With TAIT, the iterator generic on QuerySortedIter and thus the huge visible impl Iterator type in the sort function
  signatures can be removed.
• With specialization, the bound on L could be relaxed to QueryData when the underlying iterator is mutable.

Changelog

Added sort, sort_unstable, sort_by, sort_unstable_by, sort_by_key, sort_by_cached_key to QueryIter.

[alice-i-cecile]

Initial impressions are very positive:

1. This should exist.
2. Good docs.
3. Good tests.
4. Idiomatic and comprehensive API.

Bug me for a full review later if you need a second approval please.

[cBournhonesque]

Overall looks amazing to me! :)
Awesome job.

I'd like to see more tests in the non-panic case, but this is very cool

[cBournhonesque]

The code is basically identical to sort apart from the call to keyed_query.sort_unstable(); ?

I wonder if it would be helpful to put all the duplicate code into a common harness

[Victoronz]

Yes, each pair of sort, sort_by, and sort_by_key methods should be be mergeable (when we eventually have QueryManyIter sorts, they could also be merged in), and I did try that. However, doing so would require writing an inner function over a higher trait bound on WorldQuery::Item<'1>, which I found to be cursed.
Considering that merging isn't necessary for the functionality, I'd like to leave that for follow-up. I'll make a note of it under Next Steps.

[cBournhonesque]

just to be clear, if next() is called before sort() then this method is not valid because we could have duplicate mutable access to a component? (once in next() and once in the query_lens)

Or is it that we would be sorting the entire iterator (in the query_lens) even though part of the iterator has already been 'consumed'?

[Victoronz]

You're correct, mutable reference aliasing would be possible. This is in addition of us having no immediate way of tracking how many iterations have occurred, without regressing the performance of more general iteration/sorted iteration.

[Victoronz]

I changed:

• the bound on L from QueryData to ReadOnlyQueryData
• removed filter from QuerySortedIter iteration, which made the internal fetch_next infallible.
• added a test that checks whether the QueryIter sorts and the std::slice sorts are equal.

[Victoronz]

(forgot cargo fmt)

[Victoronz]

#[allow(clippy::unnecessary_sort_by)] on the new test

[alice-i-cecile]

Lovely API, docs, tests and code. Really nice stuff. I'm very surprised (positive) about how clean and expressive this is.

[Victoronz]

I'm really proud of it! As soon as it clicked how satisfying this extension would be, I just had to implement it.
And it'll become even cleaner with future Rust features! (granted, they'll take time to cook 😄)

[iiYese]

You can get rid of the generics you need to elide at call site here and elsewhere by using compare: impl FnMut(..) instead.

[Victoronz]

Sadly that requires two unstable Rust features, I list those under Future Possibilities above.

[iiYese]

I don't get any compile errors when I clone the fork and change it locally 🤔
Edit: I'm not on nightly btw.

[Victoronz]

hmm, I can try again, what I can think of is that I might only have tried on the sort_by_key methods, which have both the Fn trait and its returned key as a generic.
When attempting this, I did get two separate errors specifically for impl Fn traits and impl Trait returned by Fn traits though. 🤔
Edit: I am on nightly!

[iiYese]

What exactly is cached here?

[Victoronz]

These methods all match the slice::sort methods 1-to-1. What is cached here is the key extraction function.

[Victoronz]

Maybe I wasn't clear enough here:
These do match the slice functions, but they both take an iterator and return an iterator instead,
akin to the Itertools::sorted_* methods.

[iiYese]

The point of providing this functionality as a first party solution isn't just for ergonomics but also performance. My problem with this PR is that it is only marginally more ergonomic than collecting & sorting the query yourself. The energy on optimization is misplaced and would create friction with much bigger optimization items.

The most expensive part of sorting queries are:

1. The sorting operation: This is going to happen every time the system runs regardless of if components were actually changed. It will be very common for components in sorted queries to rarely change. Caching the result of a sort and completely early outing from this eliminates the whole sort operation. A sorted cache also allows you to do very fast suspendable seeking. Like getting N items for rendering a scroll area in UI or any kind of rhythm game (basically very fancy scroll areas).
2. Fragmented iteration: Simply looking up by Entity is going to give you completely random access. We would like to also reorder the underlying tables & then store ranges in the cache from 1. to condense it. This could be done with an explicit accompanying Command initially. It could be followed up by making sorted queries reorder tables by default for better ergonomics. Either by first adding World::parallel_commands or making the world command queue thread local. (Can also add an opt out mechanism for this behavior for systems that request overlapping queries in different orders).

I would like to see at least:

• A command to reorder tables
• The sort cache being condensed for cache friendly iteration

These items can be shelved for later because they're blocked on missing features:

• Sorting tables by default (changes to world commands)
• Using change detection to early out of sorts (archetype level change detection)

[Victoronz]

Are these not orthogonal concerns?
I agree that table sorts will be very useful, and they can be used to early out of these sort calls, but I do not see why either feature would need to block the other.
Maybe in addition to these impls being located on QueryIter, it could be helpful to explicitly note that these do not sort the underlying tables, just the iterator.

Cases like under the query-sorts test are the simplest they can get.

I am personally working on a game-like project that has to sync tree-like state with an external binary, in which I use a lot of sorts to guarantee iteration order. This is currently rather painful as they are mostly the same actions, but cannot be abstracted away by a function over query types across systems from the user side.

[Victoronz]

I do agree that it doesn't fully address the concerns of commenters under the issue #1470. Should this issue remain, or a new one be opened?

[iiYese]

I agree that table sorts will be very useful, and they can be used to early out of these sort calls, but I do not see why either feature would need to block the other.

It wouldn't block it, just make a future PR & migration guide more annoying. Currently the sort is done by the iterator not the query. If you want to store the result of a sort to early out the next time a system is called you have to store it somewhere (the QueryState is the best candidate for this). Which means this would break:

fn sys(q: Query<(&A, &B)>) {
    for (a, b) in q.iter().sort::<&A>() {
        // ..
    }
    for (a, b) in q.iter().sort::<&B>() {
        // ..
    }
}

In some surprising ways. Additionally putting the sort on the query rather than the iterator would mean you get mut iteration of sorted queries with less headaches.

[alice-i-cecile]

I do agree that it doesn't fully address the concerns of commenters under the issue #1470. Should this issue remain, or a new one be opened?

IMO this will be clearer as follow-ups.

[Victoronz]

The problem with sort methods on anything other than the query iterators is that they would explode in API surface.

• There are 6 standard sort kinds.
• AFAIK being generic over mutability here is possible exactly because a different type from Query/QueryState is used. This
  might be a Rust limitation or lack of my own knowledge, but I am pretty certain that the "ReadOnly::Item" and
  WorldQuery::Item split prevents unification. This split does not exist at the iterator level.
  This makes for a non_mut/mut method for each sort.
• The sort API is not complete unless iter_many has its own variants, they need their own return type.
  That makes for an additional 6.

That makes for a total of 24 sort methods, which is completely unnecessary IMO. For reference, Query currently has 32 non-trait methods, and QueryState 42.

Could you elaborate on mutable iteration being less of a headache with sort methods not on the iterators? I am not sure what you referring to here, as mutable iteration is completely unrestricted with these sorts.

[iiYese]

The problem with sort methods on anything other than the query iterators is that they would explode in API surface

Isn't the opposite the case? You need to make 6 sorts for each iterator type. If you instead have Query::srot, Query::sort_by_key, etc.. & store the sorted vec of entities/table ranges in the query state that you just check for in each iterator type you only have 6 new functions.

AFAIK being generic over mutability here is possible exactly because a different type

You're going to have to take the Query/QueryState as mut with this method anyway & anything but readonly items in a sort function sounds like a code smell. None of the standard library sort functions allow this afaik.

Could you elaborate on mutable iteration being less of a headache with sort methods not on the iterators

As described. You don't need extra mut iterators multiplied by each iteration type.

[alice-i-cecile]

We could use a trait for sorting (ala iter_tools) to manage some of that complexity. But I'm more than happy to push that to future work. This is useful and pleasant, and gathering feedback from users + performance benchmarks is a more reasonable step than trying to engineer this further in this PR.

I definitely think that some of the ideas above (especially caching!) have value, but I don't see a reason to block on them.

[Victoronz]

The problem with sort methods on anything other than the query iterators is that they would explode in API surface

Isn't the opposite the case? You need to make 6 sorts for each iterator type. If you instead have Query::srot,
Query::sort_by_key, etc.. & store the sorted vec of entities/table ranges in the query state that you just check for in each
iterator type you only have 6 new functions.

From the perspective of the user looking at one type at a time, they would still only be six sorts, not 12 "different" ones.

Hmm, I get the feeling we are imagining different sets of use cases/APIs here.
I think we should separate sorting the underlying data storage vs sorting an iteration over such storage.
One case will affect all future iteration, one is local to one call.
For mutation of storage vs an iteration over it to happen at the same time would conflate purposes.
Part of the intent here is f.e. to be able to call iter().sort_* multiple times in the same system! A user does not always want to touch persisting state.
Your description would also require regressing unsorted iteration, which is a much larger use case.
This change is as minimal as can be to existing code, if there really should be a redesign, then we can go for that later down the line! That would be much less invasive IMO.

AFAIK being generic over mutability here is possible exactly because a different type

You're going to have to take the Query/QueryState as mut with this method anyway & anything but readonly items in a
sort function sounds like a code smell. None of the standard library sort functions allow this afaik.

Keep in mind that the sort function itself does not mutate anything, it works with read-only data. The mutability I was talking about was for the final returned WorldQuery::Items the user actually receives.

[iiYese]

Hmm, I get the feeling we are imagining different sets of use cases/APIs here.
I think we should separate sorting the underlying data storage vs sorting an iteration over such storage.

Yes everything I've described doesn't touch the table storage.

fn sys(mut q: Query<(&mut A, &B)>) {
    // Doesn't change tables
    for (a, b) in q.sort::<&A>().iter() {
    }
    
    // Iters unsorted
    for (a, b) in q.iter_mut() {
    } 
    
    // Iter by different sort
    for (a, b) in q.sort::<&B>().iter() {
    }
}

One case will affect all future iteration, one is local to one call.
Mutation of storage vs an iteration over it to happen at the same time would conflate purposes. Part of the intent here is f.e. to be able to call iter().sort_* multiple times in the same system!

The proposed changes would be no more expensive than what the PR currently has. They would just make the query stateful & you would have less API surface.

Your description would also require regressing unsorted iteration, which is a much larger use case

It is precisely because unsorted is disproportionately more common that this kind of branch is the kind that would likely completely get optimized out by branch predictors.

[Victoronz]

oops, how how do I reopen this? (how did I close this in the first place... I am confused)

[Victoronz]

thanks, is the PR in the correct state now? I tried to rebase on main, then add a new impl Fn version commit

Edit: It seems to be.

[Victoronz]

Because that was written before opaque types were a thing. It is purely a legacy thing. Editions allow for backwards compatible breaking changes at a language level not at a library level. There is no benefit to providing the option to name the generics because

* `Fn` traits cannot be impld by users for types

* Closures are types that are unnameable

Leaving the generics there just creates API noise.

Since a Fn trait generic can be passed explicitly, and impl Trait cannot, I was thinking that it could be useful in cases where inference fails, though that should only happen in esoteric code. (Not relevant for this PR, just wondering)

[NiseVoid]

Sorts all query items into a new iterator [...]

To me this just tells the user they get a new iterator that has sorted query items. Since it uses another query to decide the ordering it could still apply all the weird obscure optimizations end users could think of.

I am also trying not to bloat the docs for these methods too much, I've cut more comments on uses for these than is currently documented, and think the extra docs might better live in a different place, i.e. Query

If info on sorted query performance is the same between methods it makes sense to explain it once in a section on Query, and only link to that section on the methods

The docs can still be adjusted afterwards, and I'd rather base that on user feedback than more speculation in advance.

How is this speculation? I looked at the PR, thought of usecases I've written or seen, then had to look trough a lot of implementation details to see how the performance stacks up the current approaches taken there. The only speculation is that people could assume it's free or way cheaper than it is. I also doubt we'd get meaningful feedback on docs, usually it's just "too little" or "not up-to-date enough" or something vague like that. Most doc improvements are done by observing users running into issues. The more valuable user feedback would probably also come if end users are aware of the performance implications of this feature and can comment on if it's worth it, as well as suggest possible improvements 🤔

[NiseVoid]

is the PR in the correct state now?

Looks like it, but you also don't usually want to rebase once you have gotten reviews, since a lot of reviewers seem to hate that.

[Victoronz]

is the PR in the correct state now?

Looks like it, but you also don't usually want to rebase once you have gotten reviews, since a lot of reviewers seem to hate that.

This is my first proper contribution to anything, so I am still trying to get used to everything there is to manage.
Do you mean rebasing on main, or rebasing to update the PR commits?

[NiseVoid]

Do you mean rebasing on main, to rebasing to update the PR commits?

Rebasing your PR onto main and force pushing it. Iirc it breaks this feature:

[Victoronz]

Sorts all query items into a new iterator [...]

To me this just tells the user they get a new iterator that has sorted query items. Since it uses another query to decide the ordering it could still apply all the weird obscure optimizations end users could think of.

I am also trying not to bloat the docs for these methods too much, I've cut more comments on uses for these than is currently documented, and think the extra docs might better live in a different place, i.e. Query

If info on sorted query performance is the same between methods it makes sense to explain it once in a section on Query, and only link to that section on the methods

The docs can still be adjusted afterwards, and I'd rather base that on user feedback than more speculation in advance.

How is this speculation? I looked at the PR, thought of usecases I've written or seen, then had to look trough a lot of implementation details to see how the performance stacks up the current approaches taken there. The only speculation is that people could assume it's free or way cheaper than it is. I also doubt we'd get meaningful feedback on docs, usually it's just "too little" or "not up-to-date enough" or something vague like that. Most doc improvements are done by observing users running into issues. The more valuable user feedback would probably also come if end users are aware of the performance implications of this feature and can comment on if it's worth it, as well as suggest possible improvements 🤔

Oh, I didn't mean to refer your performance comment suggestion as speculation, I agree we should add one in some form!

What I meant by speculation, is assuming what the users will interpret as behavior once they see these methods based on discussions that are present among people already familiar with engine concepts, like here:

They won't think it's free but might expect it to actually sort the underlying tables which is what is usually desired when this functionality is brought up in discussion

(Though speculation might be too harsh a word)

[Victoronz]

In my mind, a user unfamiliar with bevy, would travel from Query > iter/iter_mut > QueryIter > sort_*
The guarantees of what the structs and methods can do, stack on top of what the previous one does.
My assumption here is that the audience, someone at least somewhat familiar with Rust, would be most likely to approach it this way.
If enough people ask about things they'd find unclear, f.e. whether we can sort underlying table storage with it, I would think that clarifying comments could be added.

[Victoronz]

Would "Sorts all query results into a new iterator [...]" be preferable wording?

[NiseVoid]

What I meant by speculation, is assuming what the users will interpret as behavior once they see these methods based on discussions that are present among people already familiar with engine concepts, like here:

Ah yes, I don't think people will assume it does anything beyond the power of the QueryIter they passed in. But it's fairly easy to assume sorting something like Has<X>, Has<Y>, Has<Z> is cheap because in theory it could be, even if in practice that would be very niche and a pain to implement. Kind of similar to how people currently assume Changed<T> is very cheap when nothing changed, because they can imagine how it could be done cheaper than it currently is.

Would "Sorts all query results into a new iterator [...]" be preferable wording?

That definitely seems clearer 🤔

[Victoronz]

What I meant by speculation, is assuming what the users will interpret as behavior once they see these methods based on discussions that are present among people already familiar with engine concepts, like here:

Ah yes, I don't think people will assume it does anything beyond the power of the QueryIter they passed in. But it's fairly easy to assume sorting something like Has<X>, Has<Y>, Has<Z> is cheap because in theory it could be, even if in practice that would be very niche and a pain to implement. Kind of similar to how people currently assume Changed<T> is very cheap when nothing changed, because they can imagine how it could be done cheaper than it currently is.

I'm not sure what you mean, here. Has<X>, Has<Y>, Has<Z> results in (bool, bool, bool) which is can be sorted easily, and "cheaply" compared to more complex sort keys.

Do you mean cheap because they could imagine caching across system runs?

[Victoronz]

I opened #13443, if someone wants to see what the companion for QueryManyIter would look like.

On the docs:
How about an additional "The sort is not cached across system runs." sentence on each method?
I feel like much of the suggested confusion comes from "ECS is magic, how does this fit in?" from the user perspective.

This would partially address the concerns from both the performance and table sort side for now, and we could add a section on Query that talks about the performance in more detail, and table sort situation in more detail.

(Also my bad on calling your concern speculation @iiYese, I didn't phrase that well)

[alice-i-cecile]

I think that sentence would be helpful and appropriately communicate the existing limitations.

[Victoronz]

Added.

[alice-i-cecile]

Thank you to everyone involved with the authoring or reviewing of this PR! This work is relatively important and needs release notes! Head over to bevyengine/bevy-website#1322 if you'd like to help out.