Skip to content

Fix unsoundness in QueryIter::sort_by #17826

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 26, 2025
Merged

Fix unsoundness in QueryIter::sort_by #17826

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
299d904
Shorten lifetime of lens items passed to sort_by callbacks.
chescock Feb 12, 2025
72585c1
Use closures instead of passing Ord::cmp directly to get a higher-ran…
chescock Feb 12, 2025
bbe214f
The key in `sort_by_key` can no longer borrow from the lens item.
chescock Feb 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions crates/bevy_ecs/compile_fail/tests/ui/system_query_iter_sort_lifetime_safety.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
use bevy_ecs::prelude::*;
use std::cmp::Ordering;

#[derive(Component)]
struct A(usize);

fn system(mut query: Query<&mut A>) {
let iter = query.iter_mut();
let mut stored: Option<&A> = None;
let mut sorted = iter.sort_by::<&A>(|left, _right| {
// Try to smuggle the lens item out of the closure.
stored = Some(left);
//~^ E0521
Ordering::Equal
});
let r: &A = stored.unwrap();
let m: &mut A = &mut sorted.next().unwrap();
assert!(std::ptr::eq(r, m));
}
10 changes: 10 additions & 0 deletions crates/bevy_ecs/compile_fail/tests/ui/system_query_iter_sort_lifetime_safety.stderr
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
error[E0521]: borrowed data escapes outside of closure
--> tests/ui\system_query_iter_sort_lifetime_safety.rs:12:9
|
9 | let mut stored: Option<&A> = None;
| ---------- `stored` declared here, outside of the closure body
10 | let mut sorted = iter.sort_by::<&A>(|left, _right| {
| ---- `left` is a reference that is only valid in the closure body
11 | // Try to smuggle the lens item out of the closure.
12 | stored = Some(left);
| ^^^^^^^^^^^^^^^^^^^ `left` escapes the closure body here
54 changes: 30 additions & 24 deletions crates/bevy_ecs/src/query/iter.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,15 +487,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// # schedule.add_systems((system_1, system_2, system_3));
/// # schedule.run(&mut world);
/// ```
pub fn sort<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
pub fn sort<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
>
where
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}

Expand Down Expand Up @@ -541,15 +544,18 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// # schedule.add_systems((system_1));
/// # schedule.run(&mut world);
/// ```
pub fn sort_unstable<L: ReadOnlyQueryData<Item<'w>: Ord> + 'w>(
pub fn sort_unstable<L: ReadOnlyQueryData + 'w>(
self,
) -> QuerySortedIter<
'w,
's,
D,
F,
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
> {
>
where
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}

Expand Down Expand Up @@ -604,7 +610,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// ```
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -636,7 +642,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -688,7 +694,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// #[derive(Component)]
/// struct AvailableMarker;
///
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
/// enum Rarity {
/// Common,
/// Rare,
Expand Down Expand Up @@ -716,7 +722,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
/// (
/// entity_ref.contains::<AvailableMarker>(),
/// entity_ref.get::<Rarity>()
/// entity_ref.get::<Rarity>().copied()
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A little bit of collateral damage here: sort_by_key can no longer borrow from the lens item. It's theoretically possible to support that, but it's hard to express with Fn traits.

It's still possible to express those sorts using sort_by, so that doesn't make anything impossible, just inconvenient.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think this is likely just the original Rust limitation of sort_by_key/sort_by_cached_key, the closure can only return Copy or unusually long-lived data.

/// )
/// })
/// .rev()
Expand All @@ -728,7 +734,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// ```
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -761,7 +767,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -796,7 +802,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -826,7 +832,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> {
/// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
) -> QuerySortedIter<
'w,
's,
Expand Down Expand Up @@ -1333,7 +1339,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
>
where
L::Item<'w>: Ord,
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort())
}
Expand Down Expand Up @@ -1391,7 +1397,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
impl ExactSizeIterator<Item = Entity> + DoubleEndedIterator + FusedIterator + 'w,
>
where
L::Item<'w>: Ord,
for<'lw> L::Item<'lw>: Ord,
{
self.sort_impl::<L>(|keyed_query| keyed_query.sort_unstable())
}
Expand Down Expand Up @@ -1448,7 +1454,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// ```
pub fn sort_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -1479,7 +1485,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_unstable_by<L: ReadOnlyQueryData + 'w>(
self,
mut compare: impl FnMut(&L::Item<'w>, &L::Item<'w>) -> Ordering,
mut compare: impl FnMut(&L::Item<'_>, &L::Item<'_>) -> Ordering,
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -1531,7 +1537,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// #[derive(Component)]
/// struct AvailableMarker;
///
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord)]
/// #[derive(Component, PartialEq, Eq, PartialOrd, Ord, Copy, Clone)]
/// enum Rarity {
/// Common,
/// Rare,
Expand Down Expand Up @@ -1561,7 +1567,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// .sort_by_key::<EntityRef, _>(|entity_ref| {
/// (
/// entity_ref.contains::<AvailableMarker>(),
/// entity_ref.get::<Rarity>()
// entity_ref.get::<Rarity>().copied()
/// )
/// })
/// .rev()
Expand All @@ -1573,7 +1579,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// ```
pub fn sort_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -1605,7 +1611,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_unstable_by_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -1639,7 +1645,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
pub fn sort_by_cached_key<L: ReadOnlyQueryData + 'w, K>(
self,
mut f: impl FnMut(&L::Item<'w>) -> K,
mut f: impl FnMut(&L::Item<'_>) -> K,
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -1668,7 +1674,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator<Item: EntityBorrow>>
/// called on [`QueryManyIter`] before.
fn sort_impl<L: ReadOnlyQueryData + 'w>(
self,
f: impl FnOnce(&mut Vec<(L::Item<'w>, NeutralOrd<Entity>)>),
f: impl FnOnce(&mut Vec<(L::Item<'_>, NeutralOrd<Entity>)>),
) -> QuerySortedManyIter<
'w,
's,
Expand Down Expand Up @@ -2898,13 +2904,13 @@ mod tests {
{
let mut query = query_state
.iter_many_mut(&mut world, [id, id])
.sort_by::<&C>(Ord::cmp);
.sort_by::<&C>(|l, r| Ord::cmp(l, r));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How does this fail if just the function pointer is passed?
Unlike with sort_by_key, slice::sort_by is able to accept these without the need to wrap it in a closure.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It failed CI with error: implementation of `FnMut` is not general enough:
https://github.com/bevyengine/bevy/actions/runs/13291587519/job/37113490763#step:6:119

I admit I didn't investigate it too carefully; the rules for when rust infers higher-ranked lifetimes versus specific ones seem very finnicky to me.

Copy link
Contributor

@hymm hymm Feb 14, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

found a description of the problem here https://users.rust-lang.org/t/problem-with-lifetimes-in-closure-return-types/79026

TLDR: the lifetime of Ord::cmp has a specific Item<'a> lifetime, while sort_by expects 'a to be hrtb. Doesn't seem fixable without tradeoffs.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, I bet it'll infer HRTB over lifetime parameters to the function but not to lifetime parameters to the trait or type. That's reasonable, even if I wish it were looser. I remember hitting some of the odd cases mentioned in https://rust-lang.github.io/rfcs/3216-closure-lifetime-binder.html, though, which left me with the impression that higher-ranked closure inference was always a little arbitrary.

while query.fetch_next().is_some() {}
}
{
let mut query = query_state
.iter_many_mut(&mut world, [id, id])
.sort_unstable_by::<&C>(Ord::cmp);
.sort_unstable_by::<&C>(|l, r| Ord::cmp(l, r));
while query.fetch_next().is_some() {}
}
{
Expand Down