Skip to content

Commit

Permalink
implement FromEntitySetIterator (bevyengine#17513)
Browse files Browse the repository at this point in the history
# Objective

Some collections are more efficient to construct when we know that every
element is unique in advance.
We have `EntitySetIterator`s from bevyengine#16547, but currently no API to safely
make use of them this way.

## Solution

Add `FromEntitySetIterator` as a subtrait to `FromIterator`, and
implement it for the `EntityHashSet`/`hashbrown::HashSet` types.
To match the normal `FromIterator`, we also add a
`EntitySetIterator::collect_set` method.
It'd be better if these methods could shadow `from_iter` and `collect`
completely, but rust-lang/rust#89151 is needed
for that.

While currently only `HashSet`s implement this trait, future
`UniqueEntityVec`/`UniqueEntitySlice` functionality comes with more
implementors.

Because `HashMap`s are collected from tuples instead of singular types,
implementing this same optimization for them is more complex, and has to
be done separately.

## Showcase

This is basically a free speedup for collecting `EntityHashSet`s!

```rust
pub fn collect_milk_dippers(dippers: Query<Entity, (With<Milk>, With<Cookies>)>) {
    dippers.iter().collect_set::<EntityHashSet>();
    // or
    EntityHashSet::from_entity_set_iter(dippers);
}

---------

Co-authored-by: SpecificProtagonist <[email protected]>
  • Loading branch information
Victoronz and SpecificProtagonist authored Jan 24, 2025
1 parent 14ad252 commit 94a238b
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 2 deletions.
48 changes: 47 additions & 1 deletion crates/bevy_ecs/src/entity/entity_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,10 +3,12 @@ use alloc::{
collections::{btree_map, btree_set},
rc::Rc,
};
use bevy_platform_support::collections::HashSet;

use core::{
array,
fmt::{Debug, Formatter},
hash::{BuildHasher, Hash},
iter::{self, FusedIterator},
option, result,
};
Expand Down Expand Up @@ -144,7 +146,21 @@ impl<T: IntoIterator<IntoIter: EntitySetIterator>> EntitySet for T {}
///
/// `x != y` must hold for any 2 elements returned by the iterator.
/// This is always true for iterators that cannot return more than one element.
pub unsafe trait EntitySetIterator: Iterator<Item: TrustedEntityBorrow> {}
pub unsafe trait EntitySetIterator: Iterator<Item: TrustedEntityBorrow> {
/// Transforms an `EntitySetIterator` into a collection.
///
/// This is a specialized form of [`collect`], for collections which benefit from the uniqueness guarantee.
/// When present, this should always be preferred over [`collect`].
///
/// [`collect`]: Iterator::collect
// FIXME: When subtrait item shadowing stabilizes, this should be renamed and shadow `Iterator::collect`
fn collect_set<B: FromEntitySetIterator<Self::Item>>(self) -> B
where
Self: Sized,
{
FromEntitySetIterator::from_entity_set_iter(self)
}
}

// SAFETY:
// A correct `BTreeMap` contains only unique keys.
Expand Down Expand Up @@ -291,6 +307,36 @@ unsafe impl<I: EntitySetIterator, P: FnMut(&<I as Iterator>::Item) -> bool> Enti
// SAFETY: Discarding elements maintains uniqueness.
unsafe impl<I: EntitySetIterator> EntitySetIterator for iter::StepBy<I> {}

/// Conversion from an `EntitySetIterator`.
///
/// Some collections, while they can be constructed from plain iterators,
/// benefit strongly from the additional uniqeness guarantee [`EntitySetIterator`] offers.
/// Mirroring [`Iterator::collect`]/[`FromIterator::from_iter`], [`EntitySetIterator::collect_set`] and
/// `FromEntitySetIterator::from_entity_set_iter` can be used for construction.
///
/// See also: [`EntitySet`].
// FIXME: When subtrait item shadowing stabilizes, this should be renamed and shadow `FromIterator::from_iter`
pub trait FromEntitySetIterator<A: TrustedEntityBorrow>: FromIterator<A> {
/// Creates a value from an [`EntitySetIterator`].
fn from_entity_set_iter<T: EntitySet<Item = A>>(set_iter: T) -> Self;
}

impl<T: TrustedEntityBorrow + Hash, S: BuildHasher + Default> FromEntitySetIterator<T>
for HashSet<T, S>
{
fn from_entity_set_iter<I: EntitySet<Item = T>>(set_iter: I) -> Self {
let iter = set_iter.into_iter();
let set = HashSet::<T, S>::with_capacity_and_hasher(iter.size_hint().0, S::default());
iter.fold(set, |mut set, e| {
// SAFETY: Every element in self is unique.
unsafe {
set.insert_unique_unchecked(e);
}
set
})
}
}

/// An iterator that yields unique entities.
///
/// This wrapper can provide an [`EntitySetIterator`] implementation when an instance of `I` is known to uphold uniqueness.
Expand Down
16 changes: 15 additions & 1 deletion crates/bevy_ecs/src/entity/hash_set.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use bevy_platform_support::collections::hash_set::{self, HashSet};
#[cfg(feature = "bevy_reflect")]
use bevy_reflect::Reflect;

use super::{Entity, EntityHash, EntitySetIterator};
use super::{Entity, EntityHash, EntitySet, EntitySetIterator, FromEntitySetIterator};

/// A [`HashSet`] pre-configured to use [`EntityHash`] hashing.
#[cfg_attr(feature = "bevy_reflect", derive(Reflect))]
Expand Down Expand Up @@ -195,6 +195,20 @@ impl FromIterator<Entity> for EntityHashSet {
}
}

impl FromEntitySetIterator<Entity> for EntityHashSet {
fn from_entity_set_iter<I: EntitySet<Item = Entity>>(set_iter: I) -> Self {
let iter = set_iter.into_iter();
let set = EntityHashSet::with_capacity(iter.size_hint().0);
iter.fold(set, |mut set, e| {
// SAFETY: Every element in self is unique.
unsafe {
set.insert_unique_unchecked(e);
}
set
})
}
}

/// An iterator over the items of an [`EntityHashSet`].
///
/// This struct is created by the [`iter`] method on [`EntityHashSet`]. See its documentation for more.
Expand Down

0 comments on commit 94a238b

Please sign in to comment.