use alloc::{
boxed::Box,
collections::{btree_map, btree_set},
rc::Rc,
};
use bevy_platform::collections::HashSet;
use core::{
array,
cmp::Ordering,
fmt::{Debug, Formatter},
hash::{BuildHasher, Hash},
iter::{self, FusedIterator, Product, Sum},
option, ptr, result,
};
use super::{Entity, UniqueEntityEquivalentSlice};
use bevy_platform::sync::Arc;
pub trait ContainsEntity {
fn entity(&self) -> Entity;
}
pub unsafe trait EntityEquivalent: ContainsEntity + Eq {}
impl ContainsEntity for Entity {
fn entity(&self) -> Entity {
*self
}
}
unsafe impl EntityEquivalent for Entity {}
impl<T: ContainsEntity> ContainsEntity for &T {
fn entity(&self) -> Entity {
(**self).entity()
}
}
unsafe impl<T: EntityEquivalent> EntityEquivalent for &T {}
impl<T: ContainsEntity> ContainsEntity for &mut T {
fn entity(&self) -> Entity {
(**self).entity()
}
}
unsafe impl<T: EntityEquivalent> EntityEquivalent for &mut T {}
impl<T: ContainsEntity> ContainsEntity for Box<T> {
fn entity(&self) -> Entity {
(**self).entity()
}
}
unsafe impl<T: EntityEquivalent> EntityEquivalent for Box<T> {}
impl<T: ContainsEntity> ContainsEntity for Rc<T> {
fn entity(&self) -> Entity {
(**self).entity()
}
}
unsafe impl<T: EntityEquivalent> EntityEquivalent for Rc<T> {}
impl<T: ContainsEntity> ContainsEntity for Arc<T> {
fn entity(&self) -> Entity {
(**self).entity()
}
}
unsafe impl<T: EntityEquivalent> EntityEquivalent for Arc<T> {}
pub trait EntitySet: IntoIterator<IntoIter: EntitySetIterator> {}
impl<T: IntoIterator<IntoIter: EntitySetIterator>> EntitySet for T {}
pub unsafe trait EntitySetIterator: Iterator<Item: EntityEquivalent> {
fn collect_set<B: FromEntitySetIterator<Self::Item>>(self) -> B
where
Self: Sized,
{
FromEntitySetIterator::from_entity_set_iter(self)
}
}
unsafe impl<K: EntityEquivalent, V> EntitySetIterator for btree_map::Keys<'_, K, V> {}
unsafe impl<K: EntityEquivalent, V> EntitySetIterator for btree_map::IntoKeys<K, V> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for btree_set::Range<'_, T> {}
unsafe impl<T: EntityEquivalent + Ord> EntitySetIterator for btree_set::Intersection<'_, T> {}
unsafe impl<T: EntityEquivalent + Ord> EntitySetIterator for btree_set::Union<'_, T> {}
unsafe impl<T: EntityEquivalent + Ord> EntitySetIterator for btree_set::Difference<'_, T> {}
unsafe impl<T: EntityEquivalent + Ord> EntitySetIterator for btree_set::SymmetricDifference<'_, T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for btree_set::Iter<'_, T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for btree_set::IntoIter<T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for option::Iter<'_, T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for option::IntoIter<T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for result::Iter<'_, T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for result::IntoIter<T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for array::IntoIter<T, 1> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for array::IntoIter<T, 0> {}
unsafe impl<T: EntityEquivalent, F: FnOnce() -> T> EntitySetIterator for iter::OnceWith<F> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for iter::Once<T> {}
unsafe impl<T: EntityEquivalent> EntitySetIterator for iter::Empty<T> {}
unsafe impl<I: EntitySetIterator + ?Sized> EntitySetIterator for &mut I {}
unsafe impl<I: EntitySetIterator + ?Sized> EntitySetIterator for Box<I> {}
unsafe impl<'a, T: 'a + EntityEquivalent + Copy, I: EntitySetIterator<Item = &'a T>>
EntitySetIterator for iter::Copied<I>
{
}
unsafe impl<'a, T: 'a + EntityEquivalent + Clone, I: EntitySetIterator<Item = &'a T>>
EntitySetIterator for iter::Cloned<I>
{
}
unsafe impl<I: EntitySetIterator, P: FnMut(&<I as Iterator>::Item) -> bool> EntitySetIterator
for iter::Filter<I, P>
{
}
unsafe impl<I: EntitySetIterator> EntitySetIterator for iter::Fuse<I> {}
unsafe impl<I: EntitySetIterator, F: FnMut(&<I as Iterator>::Item)> EntitySetIterator
for iter::Inspect<I, F>
{
}
unsafe impl<I: DoubleEndedIterator + EntitySetIterator> EntitySetIterator for iter::Rev<I> {}
unsafe impl<I: EntitySetIterator> EntitySetIterator for iter::Skip<I> {}
unsafe impl<I: EntitySetIterator, P: FnMut(&<I as Iterator>::Item) -> bool> EntitySetIterator
for iter::SkipWhile<I, P>
{
}
unsafe impl<I: EntitySetIterator> EntitySetIterator for iter::Take<I> {}
unsafe impl<I: EntitySetIterator, P: FnMut(&<I as Iterator>::Item) -> bool> EntitySetIterator
for iter::TakeWhile<I, P>
{
}
unsafe impl<I: EntitySetIterator> EntitySetIterator for iter::StepBy<I> {}
pub trait FromEntitySetIterator<A: EntityEquivalent>: FromIterator<A> {
fn from_entity_set_iter<T: EntitySet<Item = A>>(set_iter: T) -> Self;
}
impl<T: EntityEquivalent + Hash, S: BuildHasher + Default> FromEntitySetIterator<T>
for HashSet<T, S>
{
#[inline]
fn from_entity_set_iter<I: EntitySet<Item = T>>(set_iter: I) -> Self {
let iter = set_iter.into_iter();
let set = HashSet::with_capacity_and_hasher(iter.size_hint().0, S::default());
iter.fold(set, |mut set, e| {
unsafe {
set.insert_unique_unchecked(e);
}
set
})
}
}
#[repr(transparent)]
pub struct UniqueEntityIter<I: Iterator<Item: EntityEquivalent>> {
iter: I,
}
impl<I: EntitySetIterator> UniqueEntityIter<I> {
#[inline]
pub const fn from_entity_set_iter(iter: I) -> Self {
unsafe { Self::from_iter_unchecked(iter) }
}
}
impl<I: Iterator<Item: EntityEquivalent>> UniqueEntityIter<I> {
#[inline]
pub const unsafe fn from_iter_unchecked(iter: I) -> Self {
Self { iter }
}
#[inline]
pub const unsafe fn from_iter_ref_unchecked(iter: &I) -> &Self {
unsafe { &*ptr::from_ref(iter).cast() }
}
#[inline]
pub const unsafe fn from_iter_mut_unchecked(iter: &mut I) -> &mut Self {
unsafe { &mut *ptr::from_mut(iter).cast() }
}
pub fn into_inner(self) -> I {
self.iter
}
pub const fn as_inner(&self) -> &I {
&self.iter
}
pub const unsafe fn as_mut_inner(&mut self) -> &mut I {
&mut self.iter
}
}
impl<I: Iterator<Item: EntityEquivalent>> Iterator for UniqueEntityIter<I> {
type Item = I::Item;
fn next(&mut self) -> Option<Self::Item> {
self.iter.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.iter.size_hint()
}
fn count(self) -> usize {
self.iter.count()
}
fn last(self) -> Option<Self::Item> {
self.iter.last()
}
fn nth(&mut self, n: usize) -> Option<Self::Item> {
self.iter.nth(n)
}
fn for_each<F>(self, f: F)
where
Self: Sized,
F: FnMut(Self::Item),
{
self.iter.for_each(f);
}
fn collect<B: FromIterator<Self::Item>>(self) -> B
where
Self: Sized,
{
self.iter.collect()
}
fn partition<B, F>(self, f: F) -> (B, B)
where
Self: Sized,
B: Default + Extend<Self::Item>,
F: FnMut(&Self::Item) -> bool,
{
self.iter.partition(f)
}
fn fold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.fold(init, f)
}
fn reduce<F>(self, f: F) -> Option<Self::Item>
where
Self: Sized,
F: FnMut(Self::Item, Self::Item) -> Self::Item,
{
self.iter.reduce(f)
}
fn all<F>(&mut self, f: F) -> bool
where
Self: Sized,
F: FnMut(Self::Item) -> bool,
{
self.iter.all(f)
}
fn any<F>(&mut self, f: F) -> bool
where
Self: Sized,
F: FnMut(Self::Item) -> bool,
{
self.iter.any(f)
}
fn find<P>(&mut self, predicate: P) -> Option<Self::Item>
where
Self: Sized,
P: FnMut(&Self::Item) -> bool,
{
self.iter.find(predicate)
}
fn find_map<B, F>(&mut self, f: F) -> Option<B>
where
Self: Sized,
F: FnMut(Self::Item) -> Option<B>,
{
self.iter.find_map(f)
}
fn position<P>(&mut self, predicate: P) -> Option<usize>
where
Self: Sized,
P: FnMut(Self::Item) -> bool,
{
self.iter.position(predicate)
}
fn max(self) -> Option<Self::Item>
where
Self: Sized,
Self::Item: Ord,
{
self.iter.max()
}
fn min(self) -> Option<Self::Item>
where
Self: Sized,
Self::Item: Ord,
{
self.iter.min()
}
fn max_by_key<B: Ord, F>(self, f: F) -> Option<Self::Item>
where
Self: Sized,
F: FnMut(&Self::Item) -> B,
{
self.iter.max_by_key(f)
}
fn max_by<F>(self, compare: F) -> Option<Self::Item>
where
Self: Sized,
F: FnMut(&Self::Item, &Self::Item) -> Ordering,
{
self.iter.max_by(compare)
}
fn min_by_key<B: Ord, F>(self, f: F) -> Option<Self::Item>
where
Self: Sized,
F: FnMut(&Self::Item) -> B,
{
self.iter.min_by_key(f)
}
fn min_by<F>(self, compare: F) -> Option<Self::Item>
where
Self: Sized,
F: FnMut(&Self::Item, &Self::Item) -> Ordering,
{
self.iter.min_by(compare)
}
fn sum<S>(self) -> S
where
Self: Sized,
S: Sum<Self::Item>,
{
self.iter.sum()
}
fn product<P>(self) -> P
where
Self: Sized,
P: Product<Self::Item>,
{
self.iter.product()
}
fn cmp<O>(self, other: O) -> Ordering
where
O: IntoIterator<Item = Self::Item>,
Self::Item: Ord,
Self: Sized,
{
self.iter.cmp(other)
}
fn partial_cmp<O>(self, other: O) -> Option<Ordering>
where
O: IntoIterator,
Self::Item: PartialOrd<O::Item>,
Self: Sized,
{
self.iter.partial_cmp(other)
}
fn eq<O>(self, other: O) -> bool
where
O: IntoIterator,
Self::Item: PartialEq<O::Item>,
Self: Sized,
{
self.iter.eq(other)
}
fn lt<O>(self, other: O) -> bool
where
O: IntoIterator,
Self::Item: PartialOrd<O::Item>,
Self: Sized,
{
self.iter.lt(other)
}
fn le<O>(self, other: O) -> bool
where
O: IntoIterator,
Self::Item: PartialOrd<O::Item>,
Self: Sized,
{
self.iter.le(other)
}
fn gt<O>(self, other: O) -> bool
where
O: IntoIterator,
Self::Item: PartialOrd<O::Item>,
Self: Sized,
{
self.iter.gt(other)
}
fn ge<O>(self, other: O) -> bool
where
O: IntoIterator,
Self::Item: PartialOrd<O::Item>,
Self: Sized,
{
self.iter.ge(other)
}
fn is_sorted(self) -> bool
where
Self: Sized,
Self::Item: PartialOrd,
{
self.iter.is_sorted()
}
fn is_sorted_by<F>(self, compare: F) -> bool
where
Self: Sized,
F: FnMut(&Self::Item, &Self::Item) -> bool,
{
self.iter.is_sorted_by(compare)
}
fn is_sorted_by_key<F, K>(self, f: F) -> bool
where
Self: Sized,
F: FnMut(Self::Item) -> K,
K: PartialOrd,
{
self.iter.is_sorted_by_key(f)
}
}
impl<I: ExactSizeIterator<Item: EntityEquivalent>> ExactSizeIterator for UniqueEntityIter<I> {}
impl<I: DoubleEndedIterator<Item: EntityEquivalent>> DoubleEndedIterator for UniqueEntityIter<I> {
#[inline]
fn next_back(&mut self) -> Option<Self::Item> {
self.iter.next_back()
}
fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
self.iter.nth_back(n)
}
fn rfold<B, F>(self, init: B, f: F) -> B
where
Self: Sized,
F: FnMut(B, Self::Item) -> B,
{
self.iter.rfold(init, f)
}
fn rfind<P>(&mut self, predicate: P) -> Option<Self::Item>
where
Self: Sized,
P: FnMut(&Self::Item) -> bool,
{
self.iter.rfind(predicate)
}
}
impl<I: FusedIterator<Item: EntityEquivalent>> FusedIterator for UniqueEntityIter<I> {}
unsafe impl<I: Iterator<Item: EntityEquivalent>> EntitySetIterator for UniqueEntityIter<I> {}
impl<T, I: Iterator<Item: EntityEquivalent> + AsRef<[T]>> AsRef<[T]> for UniqueEntityIter<I> {
fn as_ref(&self) -> &[T] {
self.iter.as_ref()
}
}
impl<T: EntityEquivalent, I: Iterator<Item: EntityEquivalent> + AsRef<[T]>>
AsRef<UniqueEntityEquivalentSlice<T>> for UniqueEntityIter<I>
{
fn as_ref(&self) -> &UniqueEntityEquivalentSlice<T> {
unsafe { UniqueEntityEquivalentSlice::from_slice_unchecked(self.iter.as_ref()) }
}
}
impl<T: EntityEquivalent, I: Iterator<Item: EntityEquivalent> + AsMut<[T]>>
AsMut<UniqueEntityEquivalentSlice<T>> for UniqueEntityIter<I>
{
fn as_mut(&mut self) -> &mut UniqueEntityEquivalentSlice<T> {
unsafe { UniqueEntityEquivalentSlice::from_slice_unchecked_mut(self.iter.as_mut()) }
}
}
impl<I: EntitySetIterator + Default> Default for UniqueEntityIter<I> {
fn default() -> Self {
Self {
iter: Default::default(),
}
}
}
impl<I: EntitySetIterator + Clone> Clone for UniqueEntityIter<I> {
fn clone(&self) -> Self {
Self {
iter: self.iter.clone(),
}
}
fn clone_from(&mut self, source: &Self) {
self.iter.clone_from(&source.iter);
}
}
impl<I: Iterator<Item: EntityEquivalent> + Debug> Debug for UniqueEntityIter<I> {
fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
f.debug_struct("UniqueEntityIter")
.field("iter", &self.iter)
.finish()
}
}
#[cfg(test)]
mod tests {
use alloc::{vec, vec::Vec};
use crate::{
component::Component,
entity::{Entity, EntityEquivalentHashMap, EntityEquivalentHashSet},
prelude::{Schedule, World},
query::{QueryState, With},
system::Query,
world::Mut,
};
use super::{ContainsEntity, EntityEquivalent, UniqueEntityIter};
#[derive(Component, Clone)]
pub struct Thing;
#[expect(
clippy::iter_skip_zero,
reason = "The `skip(0)` is used to ensure that the `Skip` iterator implements `EntitySet`, which is needed to pass the iterator as the `entities` parameter."
)]
#[test]
fn preserving_uniqueness() {
let mut world = World::new();
let mut query = QueryState::<&mut Thing>::new(&mut world);
let spawn_batch: Vec<Entity> = world.spawn_batch(vec![Thing; 1000]).collect();
let mut unique_entity_iter =
unsafe { UniqueEntityIter::from_iter_unchecked(spawn_batch.iter()) };
let entity_set = unique_entity_iter
.by_ref()
.filter(|_| true)
.fuse()
.inspect(|_| ())
.rev()
.skip(0)
.skip_while(|_| false)
.take(1000)
.take_while(|_| true)
.step_by(2)
.cloned();
let _results: Vec<Mut<Thing>> = query
.iter_many_unique_mut(&mut world, entity_set)
.unwrapped()
.collect();
}
#[test]
fn nesting_queries() {
let mut world = World::new();
world.spawn_batch(vec![Thing; 1000]);
pub fn system(
mut thing_entities: Query<Entity, With<Thing>>,
mut things: Query<&mut Thing>,
) {
things.iter_many_unique(thing_entities.iter());
things.iter_many_unique_mut(thing_entities.iter_mut());
}
let mut schedule = Schedule::default();
schedule.add_systems(system);
schedule.run(&mut world);
}
#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
struct EntityWrapper(Entity);
impl ContainsEntity for EntityWrapper {
fn entity(&self) -> Entity {
self.0
}
}
impl EntityWrapper {
fn new(index: u32) -> EntityWrapper {
EntityWrapper(Entity::from_raw_u32(index).unwrap())
}
}
unsafe impl EntityEquivalent for EntityWrapper {}
#[test]
fn entity_equivalent_map_test() {
type EntityWrapperMap = EntityEquivalentHashMap<EntityWrapper, i32>;
let mut map = EntityWrapperMap::default();
map.insert(EntityWrapper::new(0), 10);
map.insert(EntityWrapper::new(1), 11);
map.insert(EntityWrapper::new(0), 12);
assert_eq!(map.len(), 2);
assert!(map
.get(&EntityWrapper::new(0))
.is_some_and(|val| *val == 12));
map.remove(&EntityWrapper::new(1));
assert_eq!(map.len(), 1);
map.clear();
assert!(map.is_empty());
}
#[test]
fn entity_equivalent_set_test() {
type EntityWrapperSet = EntityEquivalentHashSet<EntityWrapper>;
let mut set = EntityWrapperSet::default();
set.insert(EntityWrapper::new(0));
set.insert(EntityWrapper::new(1));
set.insert(EntityWrapper::new(0));
assert_eq!(set.len(), 2);
assert!(set.get(&EntityWrapper::new(0)).is_some());
set.remove(&EntityWrapper::new(1));
assert_eq!(set.len(), 1);
set.clear();
assert!(set.is_empty());
}
}