implement FromEntitySetIterator (bevyengine#17513)

# 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]>

Author: 2 people

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<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.
@@ -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.

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<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.