From 94a238b0ef6db1c866be309f739637c42538b9bd Mon Sep 17 00:00:00 2001 From: Vic <59878206+Victoronz@users.noreply.github.com> Date: Fri, 24 Jan 2025 06:39:35 +0100 Subject: [PATCH] implement FromEntitySetIterator (#17513) # Objective Some collections are more efficient to construct when we know that every element is unique in advance. We have `EntitySetIterator`s from #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 https://github.com/rust-lang/rust/issues/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, With)>) { dippers.iter().collect_set::(); // or EntityHashSet::from_entity_set_iter(dippers); } --------- Co-authored-by: SpecificProtagonist --- crates/bevy_ecs/src/entity/entity_set.rs | 48 +++++++++++++++++++++++- crates/bevy_ecs/src/entity/hash_set.rs | 16 +++++++- 2 files changed, 62 insertions(+), 2 deletions(-) diff --git a/crates/bevy_ecs/src/entity/entity_set.rs b/crates/bevy_ecs/src/entity/entity_set.rs index 8b1f47779e..6b69897806 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 2f02578e56..41fdb33fa7 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.