Support non-Vec data structures in relations (#17447)

# Objective

The existing `RelationshipSourceCollection` uses `Vec` as the only
possible backing for our relationships. While a reasonable choice,
benchmarking use cases might reveal that a different data type is better
or faster.

For example:

- Not all relationships require a stable ordering between the
relationship sources (i.e. children). In cases where we a) have many
such relations and b) don't care about the ordering between them, a hash
set is likely a better datastructure than a `Vec`.
- The number of children-like entities may be small on average, and a
`smallvec` may be faster

## Solution

- Implement `RelationshipSourceCollection` for `EntityHashSet`, our
custom entity-optimized `HashSet`.
-~~Implement `DoubleEndedIterator` for `EntityHashSet` to make things
compile.~~
   -  This implementation was cursed and very surprising.
- Instead, by moving the iterator type on `RelationshipSourceCollection`
from an erased RPTIT to an explicit associated type we can add a trait
bound on the offending methods!
- Implement `RelationshipSourceCollection` for `SmallVec`

## Testing

I've added a pair of new tests to make sure this pattern compiles
successfully in practice!

## Migration Guide

`EntityHashSet` and `EntityHashMap` are no longer re-exported in
`bevy_ecs::entity` directly. If you were not using `bevy_ecs` / `bevy`'s
`prelude`, you can access them through their now-public modules,
`hash_set` and `hash_map` instead.

## Notes to reviewers

The `EntityHashSet::Iter` type needs to be public for this impl to be
allowed. I initially renamed it to something that wasn't ambiguous and
re-exported it, but as @Victoronz pointed out, that was somewhat
unidiomatic.

In
1a85648,
I instead made the `entity_hash_set` public (and its `entity_hash_set`)
sister public, and removed the re-export. I prefer this design (give me
module docs please), but it leads to a lot of churn in this PR.

Let me know which you'd prefer, and if you'd like me to split that
change out into its own micro PR.

benches/benches/bevy_ecs/world/entity_hash.rs

@@ -1,4 +1,4 @@
-use bevy_ecs::entity::{Entity, EntityHashSet};
+use bevy_ecs::entity::{hash_set::EntityHashSet, Entity};
 use criterion::{BenchmarkId, Criterion, Throughput};
 use rand::{Rng, SeedableRng};
 use rand_chacha::ChaCha8Rng;

crates/bevy_core_pipeline/src/oit/resolve/mod.rs

@@ -6,7 +6,7 @@ use bevy_app::Plugin;
 use bevy_asset::{load_internal_asset, Handle};
 use bevy_derive::Deref;
 use bevy_ecs::{
-    entity::{EntityHashMap, EntityHashSet},
+    entity::{hash_map::EntityHashMap, hash_set::EntityHashSet},
     prelude::*,
 };
 use bevy_image::BevyDefault as _;

crates/bevy_ecs/Cargo.toml

@@ -125,7 +125,7 @@ derive_more = { version = "1", default-features = false, features = [
 ] }
 nonmax = { version = "0.5", default-features = false }
 arrayvec = { version = "0.7.4", default-features = false, optional = true }
-smallvec = { version = "1", features = ["union"] }
+smallvec = { version = "1", features = ["union", "const_generics"] }
 indexmap = { version = "2.5.0", default-features = false }
 variadics_please = { version = "1.1", default-features = false }
 spin = { version = "0.9.8", default-features = false, features = [

crates/bevy_ecs/src/entity/hash_map.rs

@@ -1,3 +1,7 @@
+//! Contains the [`EntityHashMap`] type, a [`HashMap`] pre-configured to use [`EntityHash`] hashing.
+//!
+//! This module is a lightweight wrapper around [`hashbrown`](bevy_utils::hashbrown)'s [`HashMap`] that is more performant for [`Entity`] keys.
+
 use core::{
     fmt::{self, Debug, Formatter},
     iter::FusedIterator,

crates/bevy_ecs/src/entity/hash_set.rs

@@ -1,3 +1,7 @@
+//! Contains the [`EntityHashSet`] type, a [`HashSet`] pre-configured to use [`EntityHash`] hashing.
+//!
+//! This module is a lightweight wrapper around [`hashbrown`](bevy_utils::hashbrown)'s [`HashSet`] that is more performant for [`Entity`] keys.
+
 use core::{
     fmt::{self, Debug, Formatter},
     iter::FusedIterator,
@@ -38,6 +42,16 @@ impl EntityHashSet {
         Self(HashSet::with_capacity_and_hasher(n, EntityHash))
     }
 
+    /// Returns the number of elements in the set.
+    pub fn len(&self) -> usize {
+        self.0.len()
+    }
+
+    /// Returns `true` if the set contains no elements.
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
     /// Returns the inner [`HashSet`].
     pub fn into_inner(self) -> HashSet<Entity, EntityHash> {
         self.0

crates/bevy_ecs/src/entity/map_entities.rs

@@ -4,7 +4,7 @@ use crate::{
     world::World,
 };
 
-use super::{EntityHashMap, VisitEntitiesMut};
+use super::{hash_map::EntityHashMap, VisitEntitiesMut};
 
 /// Operation to map all contained [`Entity`] fields in a type to new values.
 ///
@@ -71,7 +71,7 @@ impl<T: VisitEntitiesMut> MapEntities for T {
 ///
 /// ```
 /// # use bevy_ecs::entity::{Entity, EntityMapper};
-/// # use bevy_ecs::entity::EntityHashMap;
+/// # use bevy_ecs::entity::hash_map::EntityHashMap;
 /// #
 /// pub struct SimpleEntityMapper {
 ///   map: EntityHashMap<Entity>,
@@ -194,7 +194,7 @@ impl<'m> SceneEntityMapper<'m> {
 #[cfg(test)]
 mod tests {
     use crate::{
-        entity::{Entity, EntityHashMap, EntityMapper, SceneEntityMapper},
+        entity::{hash_map::EntityHashMap, Entity, EntityMapper, SceneEntityMapper},
         world::World,
     };
 

crates/bevy_ecs/src/entity/mod.rs

@@ -53,11 +53,8 @@ pub use visit_entities::*;
 mod hash;
 pub use hash::*;
 
-mod hash_map;
-mod hash_set;
-
-pub use hash_map::EntityHashMap;
-pub use hash_set::EntityHashSet;
+pub mod hash_map;
+pub mod hash_set;
 
 use crate::{
     archetype::{ArchetypeId, ArchetypeRow},

crates/bevy_ecs/src/entity/visit_entities.rs

@@ -58,7 +58,7 @@ impl VisitEntitiesMut for Entity {
 mod tests {
     use crate::{
         self as bevy_ecs,
-        entity::{EntityHashMap, MapEntities, SceneEntityMapper},
+        entity::{hash_map::EntityHashMap, MapEntities, SceneEntityMapper},
         world::World,
     };
     use alloc::{string::String, vec, vec::Vec};

crates/bevy_ecs/src/observer/mod.rs

@@ -9,7 +9,7 @@ pub use runner::*;
 use crate::{
     archetype::ArchetypeFlags,
     component::ComponentId,
-    entity::EntityHashMap,
+    entity::hash_map::EntityHashMap,
     prelude::*,
     system::IntoObserverSystem,
     world::{DeferredWorld, *},

crates/bevy_ecs/src/relationship/mod.rs

@@ -130,12 +130,22 @@ pub trait Relationship: Component + Sized {
     }
 }
 
+/// The iterator type for the source entities in a [`RelationshipTarget`] collection,
+/// as defined in the [`RelationshipSourceCollection`] trait.
+pub type SourceIter<'w, R> =
+    <<R as RelationshipTarget>::Collection as RelationshipSourceCollection>::SourceIter<'w>;
+
 /// A [`Component`] containing the collection of entities that relate to this [`Entity`] via the associated `Relationship` type.
 /// See the [`Relationship`] documentation for more information.
 pub trait RelationshipTarget: Component<Mutability = Mutable> + Sized {
     /// The [`Relationship`] that populates this [`RelationshipTarget`] collection.
     type Relationship: Relationship<RelationshipTarget = Self>;
     /// The collection type that stores the "source" entities for this [`RelationshipTarget`] component.
+    ///
+    /// Check the list of types which implement [`RelationshipSourceCollection`] for the data structures that can be used inside of your component.
+    /// If you need a new collection type, you can implement the [`RelationshipSourceCollection`] trait
+    /// for a type you own which wraps the collection you want to use (to avoid the orphan rule),
+    /// or open an issue on the Bevy repository to request first-party support for your collection type.
     type Collection: RelationshipSourceCollection;
 
     /// Returns a reference to the stored [`RelationshipTarget::Collection`].
@@ -210,7 +220,7 @@ pub trait RelationshipTarget: Component<Mutability = Mutable> + Sized {
 
     /// Iterates the entities stored in this collection.
     #[inline]
-    fn iter(&self) -> impl DoubleEndedIterator<Item = Entity> {
+    fn iter(&self) -> SourceIter<'_, Self> {
         self.collection().iter()
     }
 

crates/bevy_ecs/src/relationship/relationship_query.rs

@@ -7,6 +7,8 @@ use crate::{
 use alloc::collections::VecDeque;
 use smallvec::SmallVec;
 
+use super::SourceIter;
+
 impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     /// If the given `entity` contains the `R` [`Relationship`] component, returns the
     /// target entity of that relationship.
@@ -59,6 +61,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     ) -> impl Iterator<Item = Entity> + 'w
     where
         <D as QueryData>::ReadOnly: WorldQuery<Item<'w> = &'w S>,
+        SourceIter<'w, S>: DoubleEndedIterator,
     {
         self.iter_descendants_depth_first(entity).filter(|entity| {
             self.get(*entity)
@@ -114,6 +117,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter> Query<'w, 's, D, F> {
     ) -> DescendantDepthFirstIter<'w, 's, D, F, S>
     where
         D::ReadOnly: WorldQuery<Item<'w> = &'w S>,
+        SourceIter<'w, S>: DoubleEndedIterator,
     {
         DescendantDepthFirstIter::new(self, entity)
     }
@@ -195,6 +199,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget>
     DescendantDepthFirstIter<'w, 's, D, F, S>
 where
     D::ReadOnly: WorldQuery<Item<'w> = &'w S>,
+    SourceIter<'w, S>: DoubleEndedIterator,
 {
     /// Returns a new [`DescendantDepthFirstIter`].
     pub fn new(children_query: &'w Query<'w, 's, D, F>, entity: Entity) -> Self {
@@ -211,6 +216,7 @@ impl<'w, 's, D: QueryData, F: QueryFilter, S: RelationshipTarget> Iterator
     for DescendantDepthFirstIter<'w, 's, D, F, S>
 where
     D::ReadOnly: WorldQuery<Item<'w> = &'w S>,
+    SourceIter<'w, S>: DoubleEndedIterator,
 {
     type Item = Entity;
 

crates/bevy_ecs/src/relationship/relationship_source_collection.rs

@@ -1,9 +1,21 @@
-use crate::entity::Entity;
+use crate::entity::{hash_set::EntityHashSet, Entity};
 use alloc::vec::Vec;
+use smallvec::SmallVec;
 
 /// The internal [`Entity`] collection used by a [`RelationshipTarget`](crate::relationship::RelationshipTarget) component.
 /// This is not intended to be modified directly by users, as it could invalidate the correctness of relationships.
 pub trait RelationshipSourceCollection {
+    /// The type of iterator returned by the `iter` method.
+    ///
+    /// This is an associated type (rather than using a method that returns an opaque return-position impl trait)
+    /// to ensure that all methods and traits (like [`DoubleEndedIterator`]) of the underlying collection's iterator
+    /// are available to the user when implemented without unduly restricting the possible collections.
+    ///
+    /// The [`SourceIter`](super::SourceIter) type alias can be helpful to reduce confusion when working with this associated type.
+    type SourceIter<'a>: Iterator<Item = Entity>
+    where
+        Self: 'a;
+
     /// Returns an instance with the given pre-allocated entity `capacity`.
     fn with_capacity(capacity: usize) -> Self;
 
@@ -14,7 +26,7 @@ pub trait RelationshipSourceCollection {
     fn remove(&mut self, entity: Entity);
 
     /// Iterates all entities in the collection.
-    fn iter(&self) -> impl DoubleEndedIterator<Item = Entity>;
+    fn iter(&self) -> Self::SourceIter<'_>;
 
     /// Returns the current length of the collection.
     fn len(&self) -> usize;
@@ -27,6 +39,8 @@ pub trait RelationshipSourceCollection {
 }
 
 impl RelationshipSourceCollection for Vec<Entity> {
+    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
+
     fn with_capacity(capacity: usize) -> Self {
         Vec::with_capacity(capacity)
     }
@@ -41,11 +55,134 @@ impl RelationshipSourceCollection for Vec<Entity> {
         }
     }
 
-    fn iter(&self) -> impl DoubleEndedIterator<Item = Entity> {
+    fn iter(&self) -> Self::SourceIter<'_> {
         <[Entity]>::iter(self).copied()
     }
 
     fn len(&self) -> usize {
         Vec::len(self)
     }
 }
+
+impl RelationshipSourceCollection for EntityHashSet {
+    type SourceIter<'a> = core::iter::Copied<crate::entity::hash_set::Iter<'a>>;
+
+    fn with_capacity(capacity: usize) -> Self {
+        EntityHashSet::with_capacity(capacity)
+    }
+
+    fn add(&mut self, entity: Entity) {
+        self.insert(entity);
+    }
+
+    fn remove(&mut self, entity: Entity) {
+        // We need to call the remove method on the underlying hash set,
+        // which takes its argument by reference
+        self.0.remove(&entity);
+    }
+
+    fn iter(&self) -> Self::SourceIter<'_> {
+        self.iter().copied()
+    }
+
+    fn len(&self) -> usize {
+        self.len()
+    }
+}
+
+impl<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
+    type SourceIter<'a> = core::iter::Copied<core::slice::Iter<'a, Entity>>;
+
+    fn with_capacity(capacity: usize) -> Self {
+        SmallVec::with_capacity(capacity)
+    }
+
+    fn add(&mut self, entity: Entity) {
+        SmallVec::push(self, entity);
+    }
+
+    fn remove(&mut self, entity: Entity) {
+        if let Some(index) = <[Entity]>::iter(self).position(|e| *e == entity) {
+            SmallVec::remove(self, index);
+        }
+    }
+
+    fn iter(&self) -> Self::SourceIter<'_> {
+        <[Entity]>::iter(self).copied()
+    }
+
+    fn len(&self) -> usize {
+        SmallVec::len(self)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate as bevy_ecs;
+    use crate::prelude::{Component, World};
+    use crate::relationship::RelationshipTarget;
+
+    #[test]
+    fn vec_relationship_source_collection() {
+        #[derive(Component)]
+        #[relationship(relationship_target = RelTarget)]
+        struct Rel(Entity);
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Rel, despawn_descendants)]
+        struct RelTarget(Vec<Entity>);
+
+        let mut world = World::new();
+        let a = world.spawn_empty().id();
+        let b = world.spawn_empty().id();
+
+        world.entity_mut(a).insert(Rel(b));
+
+        let rel_target = world.get::<RelTarget>(b).unwrap();
+        let collection = rel_target.collection();
+        assert_eq!(collection, &alloc::vec!(a));
+    }
+
+    #[test]
+    fn entity_hash_set_relationship_source_collection() {
+        #[derive(Component)]
+        #[relationship(relationship_target = RelTarget)]
+        struct Rel(Entity);
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Rel, despawn_descendants)]
+        struct RelTarget(EntityHashSet);
+
+        let mut world = World::new();
+        let a = world.spawn_empty().id();
+        let b = world.spawn_empty().id();
+
+        world.entity_mut(a).insert(Rel(b));
+
+        let rel_target = world.get::<RelTarget>(b).unwrap();
+        let collection = rel_target.collection();
+        assert_eq!(collection, &EntityHashSet::from([a]));
+    }
+
+    #[test]
+    fn smallvec_relationship_source_collection() {
+        #[derive(Component)]
+        #[relationship(relationship_target = RelTarget)]
+        struct Rel(Entity);
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Rel, despawn_descendants)]
+        struct RelTarget(SmallVec<[Entity; 4]>);
+
+        let mut world = World::new();
+        let a = world.spawn_empty().id();
+        let b = world.spawn_empty().id();
+
+        world.entity_mut(a).insert(Rel(b));
+
+        let rel_target = world.get::<RelTarget>(b).unwrap();
+        let collection = rel_target.collection();
+        assert_eq!(collection, &SmallVec::from_buf([a]));
+    }
+}

crates/bevy_ecs/src/world/deferred_world.rs

@@ -189,8 +189,8 @@ impl<'w> DeferredWorld<'w> {
     /// For examples, see [`DeferredWorld::entity_mut`].
     ///
     /// [`EntityMut`]: crate::world::EntityMut
-    /// [`&EntityHashSet`]: crate::entity::EntityHashSet
-    /// [`EntityHashMap<EntityMut>`]: crate::entity::EntityHashMap
+    /// [`&EntityHashSet`]: crate::entity::hash_set::EntityHashSet
+    /// [`EntityHashMap<EntityMut>`]: crate::entity::hash_map::EntityHashMap
     /// [`Vec<EntityMut>`]: alloc::vec::Vec
     #[inline]
     pub fn get_entity_mut<F: WorldEntityFetch>(
@@ -298,7 +298,7 @@ impl<'w> DeferredWorld<'w> {
     /// ## [`&EntityHashSet`]
     ///
     /// ```
-    /// # use bevy_ecs::{prelude::*, entity::EntityHashSet, world::DeferredWorld};
+    /// # use bevy_ecs::{prelude::*, entity::hash_set::EntityHashSet, world::DeferredWorld};
     /// #[derive(Component)]
     /// struct Position {
     ///   x: f32,
@@ -321,8 +321,8 @@ impl<'w> DeferredWorld<'w> {
     /// ```
     ///
     /// [`EntityMut`]: crate::world::EntityMut
-    /// [`&EntityHashSet`]: crate::entity::EntityHashSet
-    /// [`EntityHashMap<EntityMut>`]: crate::entity::EntityHashMap
+    /// [`&EntityHashSet`]: crate::entity::hash_set::EntityHashSet
+    /// [`EntityHashMap<EntityMut>`]: crate::entity::hash_map::EntityHashMap
     /// [`Vec<EntityMut>`]: alloc::vec::Vec
     #[inline]
     pub fn entity_mut<F: WorldEntityFetch>(&mut self, entities: F) -> F::DeferredMut<'_> {