Allow entities to be added to a relationship using the new spawn api (remade) (#19046)

# Objective

Reopens #18961

## Solution

Copy the code

## Testing

It worked in the previous pull request

## Showcase

See #18961

---------

Co-authored-by: Alice Cecile <[email protected]>

Author: Freyja-moth

crates/bevy_ecs/src/lib.rs

@@ -97,7 +97,7 @@ pub mod prelude {
             common_conditions::*, ApplyDeferred, IntoScheduleConfigs, IntoSystemSet, Schedule,
             Schedules, SystemCondition, SystemSet,
         },
-        spawn::{Spawn, SpawnRelated},
+        spawn::{Spawn, SpawnIter, SpawnRelated, SpawnWith, WithOneRelated, WithRelated},
         system::{
             Command, Commands, Deferred, EntityCommand, EntityCommands, In, InMut, InRef,
             IntoSystem, Local, NonSend, NonSendMut, ParamSet, Populated, Query, ReadOnlySystem,

crates/bevy_ecs/src/spawn.rs

@@ -133,6 +133,80 @@ impl<R: Relationship, F: FnOnce(&mut RelatedSpawner<R>) + Send + Sync + 'static>
     }
 }
 
+/// A [`SpawnableList`] that links already spawned entities to the root entity via relations of type `I`.
+///
+/// This is useful if the entity has already been spawned earlier or if you spawn multiple relationships link to the same entity at the same time.
+/// If you only need to do this for a single entity, consider using [`WithOneRelated`].
+///
+/// ```
+/// # use bevy_ecs::hierarchy::Children;
+/// # use bevy_ecs::spawn::{Spawn, WithRelated, SpawnRelated};
+/// # use bevy_ecs::name::Name;
+/// # use bevy_ecs::world::World;
+/// let mut world = World::new();
+///
+/// let child2 = world.spawn(Name::new("Child2")).id();
+/// let child3 = world.spawn(Name::new("Child3")).id();
+///
+/// world.spawn((
+///     Name::new("Root"),
+///     Children::spawn((
+///         Spawn(Name::new("Child1")),
+///         // This adds the already existing entities as children of Root.
+///         WithRelated([child2, child3].into_iter()),
+///     )),
+/// ));
+/// ```
+pub struct WithRelated<I>(pub I);
+
+impl<R: Relationship, I: Iterator<Item = Entity>> SpawnableList<R> for WithRelated<I> {
+    fn spawn(self, world: &mut World, entity: Entity) {
+        world
+            .entity_mut(entity)
+            .add_related::<R>(&self.0.collect::<Vec<_>>());
+    }
+
+    fn size_hint(&self) -> usize {
+        self.0.size_hint().0
+    }
+}
+
+/// A wrapper over an [`Entity`] indicating that an entity should be added.
+/// This is intended to be used for hierarchical spawning via traits like [`SpawnableList`] and [`SpawnRelated`].
+///
+/// Unlike [`WithRelated`] this only adds one entity.
+///
+/// Also see the [`children`](crate::children) and [`related`](crate::related) macros that abstract over the [`Spawn`] API.
+///
+/// ```
+/// # use bevy_ecs::hierarchy::Children;
+/// # use bevy_ecs::spawn::{Spawn, WithOneRelated, SpawnRelated};
+/// # use bevy_ecs::name::Name;
+/// # use bevy_ecs::world::World;
+/// let mut world = World::new();
+///
+/// let child1 = world.spawn(Name::new("Child1")).id();
+///
+/// world.spawn((
+///     Name::new("Root"),
+///     Children::spawn((
+///         // This adds the already existing entity as a child of Root.
+///         WithOneRelated(child1),
+///     )),
+/// ));
+/// ```
+pub struct WithOneRelated(pub Entity);
+
+impl<R: Relationship> SpawnableList<R> for WithOneRelated {
+    fn spawn(self, world: &mut World, entity: Entity) {
+        world.entity_mut(entity).add_one_related::<R>(self.0);
+    }
+
+    fn size_hint(&self) -> usize {
+        1
+    }
+}
+
 macro_rules! spawnable_list_impl {
     ($($list: ident),*) => {
         #[expect(
@@ -267,7 +341,7 @@ pub trait SpawnRelated: RelationshipTarget {
     /// Returns a [`Bundle`] containing this [`RelationshipTarget`] component. It also spawns a [`SpawnableList`] of entities, each related to the bundle's entity
     /// via [`RelationshipTarget::Relationship`]. The [`RelationshipTarget`] (when possible) will pre-allocate space for the related entities.
     ///
-    /// See [`Spawn`], [`SpawnIter`], and [`SpawnWith`] for usage examples.
+    /// See [`Spawn`], [`SpawnIter`], [`SpawnWith`], [`WithRelated`] and [`WithOneRelated`] for usage examples.
     fn spawn<L: SpawnableList<Self::Relationship>>(
         list: L,
     ) -> SpawnRelatedBundle<Self::Relationship, L>;
@@ -485,3 +559,141 @@ macro_rules! recursive_spawn {
         )
     };
 }
+
+#[cfg(test)]
+mod tests {
+
+    use crate::{
+        name::Name,
+        prelude::{ChildOf, Children, RelationshipTarget},
+        relationship::RelatedSpawner,
+        world::World,
+    };
+
+    use super::{Spawn, SpawnIter, SpawnRelated, SpawnWith, WithOneRelated, WithRelated};
+
+    #[test]
+    fn spawn() {
+        let mut world = World::new();
+
+        let parent = world
+            .spawn((
+                Name::new("Parent"),
+                Children::spawn(Spawn(Name::new("Child1"))),
+            ))
+            .id();
+
+        let children = world
+            .query::<&Children>()
+            .get(&world, parent)
+            .expect("An entity with Children should exist");
+
+        assert_eq!(children.iter().count(), 1);
+
+        for ChildOf(child) in world.query::<&ChildOf>().iter(&world) {
+            assert_eq!(child, &parent);
+        }
+    }
+
+    #[test]
+    fn spawn_iter() {
+        let mut world = World::new();
+
+        let parent = world
+            .spawn((
+                Name::new("Parent"),
+                Children::spawn(SpawnIter(["Child1", "Child2"].into_iter().map(Name::new))),
+            ))
+            .id();
+
+        let children = world
+            .query::<&Children>()
+            .get(&world, parent)
+            .expect("An entity with Children should exist");
+
+        assert_eq!(children.iter().count(), 2);
+
+        for ChildOf(child) in world.query::<&ChildOf>().iter(&world) {
+            assert_eq!(child, &parent);
+        }
+    }
+
+    #[test]
+    fn spawn_with() {
+        let mut world = World::new();
+
+        let parent = world
+            .spawn((
+                Name::new("Parent"),
+                Children::spawn(SpawnWith(|parent: &mut RelatedSpawner<ChildOf>| {
+                    parent.spawn(Name::new("Child1"));
+                })),
+            ))
+            .id();
+
+        let children = world
+            .query::<&Children>()
+            .get(&world, parent)
+            .expect("An entity with Children should exist");
+
+        assert_eq!(children.iter().count(), 1);
+
+        for ChildOf(child) in world.query::<&ChildOf>().iter(&world) {
+            assert_eq!(child, &parent);
+        }
+    }
+
+    #[test]
+    fn with_related() {
+        let mut world = World::new();
+
+        let child1 = world.spawn(Name::new("Child1")).id();
+        let child2 = world.spawn(Name::new("Child2")).id();
+
+        let parent = world
+            .spawn((
+                Name::new("Parent"),
+                Children::spawn(WithRelated([child1, child2].into_iter())),
+            ))
+            .id();
+
+        let children = world
+            .query::<&Children>()
+            .get(&world, parent)
+            .expect("An entity with Children should exist");
+
+        assert_eq!(children.iter().count(), 2);
+
+        assert_eq!(
+            world.entity(child1).get::<ChildOf>(),
+            Some(&ChildOf(parent))
+        );
+        assert_eq!(
+            world.entity(child2).get::<ChildOf>(),
+            Some(&ChildOf(parent))
+        );
+    }
+
+    #[test]
+    fn with_one_related() {
+        let mut world = World::new();
+
+        let child1 = world.spawn(Name::new("Child1")).id();
+
+        let parent = world
+            .spawn((Name::new("Parent"), Children::spawn(WithOneRelated(child1))))
+            .id();
+
+        let children = world
+            .query::<&Children>()
+            .get(&world, parent)
+            .expect("An entity with Children should exist");
+
+        assert_eq!(children.iter().count(), 1);
+
+        assert_eq!(
+            world.entity(child1).get::<ChildOf>(),
+            Some(&ChildOf(parent))
+        );
+    }
+}