diff --git a/crates/bevy_ecs/src/entity/entity_set.rs b/crates/bevy_ecs/src/entity/entity_set.rs index 8b1f47779ef86..6b69897806071 100644 --- a/crates/bevy_ecs/src/entity/entity_set.rs +++ b/crates/bevy_ecs/src/entity/entity_set.rs @@ -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, }; @@ -144,7 +146,21 @@ impl> 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 {} +pub unsafe trait EntitySetIterator: Iterator { + /// 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>(self) -> B + where + Self: Sized, + { + FromEntitySetIterator::from_entity_set_iter(self) + } +} // SAFETY: // A correct `BTreeMap` contains only unique keys. @@ -291,6 +307,36 @@ unsafe impl::Item) -> bool> Enti // SAFETY: Discarding elements maintains uniqueness. unsafe impl EntitySetIterator for iter::StepBy {} +/// 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: FromIterator { + /// Creates a value from an [`EntitySetIterator`]. + fn from_entity_set_iter>(set_iter: T) -> Self; +} + +impl FromEntitySetIterator + for HashSet +{ + fn from_entity_set_iter>(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| { + // 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. diff --git a/crates/bevy_ecs/src/entity/hash_set.rs b/crates/bevy_ecs/src/entity/hash_set.rs index 2f02578e56d0d..41fdb33fa7ca1 100644 --- a/crates/bevy_ecs/src/entity/hash_set.rs +++ b/crates/bevy_ecs/src/entity/hash_set.rs @@ -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))] @@ -195,6 +195,20 @@ impl FromIterator for EntityHashSet { } } +impl FromEntitySetIterator for EntityHashSet { + fn from_entity_set_iter>(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.