diff --git a/crates/bevy_ecs/src/query/iter.rs b/crates/bevy_ecs/src/query/iter.rs index 05955d5b37854..4c384f3616fc9 100644 --- a/crates/bevy_ecs/src/query/iter.rs +++ b/crates/bevy_ecs/src/query/iter.rs @@ -1,3 +1,5 @@ +use bevy_utils::RandomState; + use crate::{ archetype::{Archetype, ArchetypeEntity, Archetypes}, component::Tick, @@ -10,6 +12,7 @@ use std::{ borrow::Borrow, cmp::Ordering, fmt::{self, Debug, Formatter}, + hash::{BuildHasher, Hash, Hasher}, iter::FusedIterator, mem::MaybeUninit, ops::Range, @@ -167,6 +170,53 @@ impl<'w, 's, D: QueryData, F: QueryFilter> QueryIter<'w, 's, D, F> { accum } + /// Randomizes the order of all query items within a new iterator. + /// + /// This method produces a new order with each call. + /// + /// # Panics + /// + /// This will panic if `next` has been called on `QueryIter` before, unless the underlying `Query` is empty. + /// + /// # Example + /// + /// ``` + /// # use bevy_ecs::prelude::*; + /// # + /// # #[derive(Component)] + /// # pub struct Nothing; + /// # + /// # let mut world = World::new(); + /// # world.spawn_batch((0..100).map(|_| Nothing)); + /// # + /// # fn system(query: Query<()>) { + /// # let n = 7; + /// // Fetch n random query items! + /// let random_items = query.iter().randomize().take(n); + /// # } + /// # + /// # let mut schedule = Schedule::default(); + /// # schedule.add_systems(system); + /// # schedule.run(&mut world); + /// ``` + pub fn randomize( + self, + ) -> QueryRandomIter< + 'w, + 's, + D, + F, + impl ExactSizeIterator + DoubleEndedIterator + FusedIterator + 'w, + > { + // FIXME: This uses AHash, it could be made faster with `EntityHash`. + // However, `EntityHasher` currently does not have a `RandomState` equivalent. + let mut hasher = RandomState::new().build_hasher(); + self.sort_by_cached_key::(move |e| { + e.hash(&mut hasher); + hasher.finish() + }) + } + /// Sorts all query items into a new iterator, using the query lens as a key. /// /// This sort is stable (i.e., does not reorder equal elements). @@ -1071,6 +1121,11 @@ impl<'w, 's, D: QueryData, F: QueryFilter, I: Iterator> Debug } } +/// An [`Iterator`] over randomized query results of a [`Query`](crate::system::Query). +/// +/// This type is created by the [`QueryIter::randomize`] method. +pub type QueryRandomIter<'w, 's, D, F, I> = QuerySortedIter<'w, 's, D, F, I>; + /// An [`Iterator`] over the query items generated from an iterator of [`Entity`]s. /// /// Items are returned in the order of the provided iterator.