From 6eb6afeb2dfea943dd24f9d3d2a7cd994b83ed5c Mon Sep 17 00:00:00 2001
From: Guillaume Wafo-Tapa <43912646+gwafotapa@users.noreply.github.com>
Date: Tue, 1 Jul 2025 00:13:38 +0200
Subject: [PATCH] Spawn batch with relationship (#19519)

# Objective

Fixes #19356
Issue: Spawning a batch of entities in relationship with the same target
adds the relationship between the target and only the last entity of the
batch. `spawn_batch` flushes only after having spawned all entities.
This means each spawned entity will have run the `on_insert` hook of its
`Relationship` component. Here is the relevant part of that hook:
```Rust
            if let Some(mut relationship_target) =
                target_entity_mut.get_mut::<Self::RelationshipTarget>()
            {
                relationship_target.collection_mut_risky().add(entity);
            } else {
                let mut target = <Self::RelationshipTarget as RelationshipTarget>::with_capacity(1);
                target.collection_mut_risky().add(entity);
                world.commands().entity(target_entity).insert(target);
            }
```
Given the above snippet and since there's no flush between spawns, each
entity finds the target without a `RelationshipTarget` component and
defers the insertion of that component with the entity's id as the sole
member of its collection. When the commands are finally flushed, each
insertion after the first replaces the one before and in the process
triggers the `on_replace` hook of `RelationshipTarget` which removes the
`Relationship` component from the corresponding entity. That's how we
end up in the invalid state.

## Solution

I see two possible solutions
1. Flush after every spawn
2. Defer the whole code snippet above

I don't know enough about bevy as a whole but 2. seems much more
efficient to me. This is what I'm proposing here. I have a doubt though
because I've started to look at #19348 that 1. would fix as well.

## Testing

I added a test for the issue. I've put it in `relationship/mod.rs` but I
could see it in `world/spawn_batch.rs` or `lib.rs` because the test is
as much about `spawn_batch` as it is about relationships.
---
 crates/bevy_ecs/src/relationship/mod.rs | 55 +++++++++++++++++++------
 1 file changed, 43 insertions(+), 12 deletions(-)

diff --git a/crates/bevy_ecs/src/relationship/mod.rs b/crates/bevy_ecs/src/relationship/mod.rs
index 12def083bd..f95214262b 100644
--- a/crates/bevy_ecs/src/relationship/mod.rs
+++ b/crates/bevy_ecs/src/relationship/mod.rs
@@ -126,16 +126,18 @@ pub trait Relationship: Component + Sized {
             world.commands().entity(entity).remove::<Self>();
             return;
         }
-        if let Ok(mut target_entity_mut) = world.get_entity_mut(target_entity) {
-            if let Some(mut relationship_target) =
-                target_entity_mut.get_mut::<Self::RelationshipTarget>()
-            {
-                relationship_target.collection_mut_risky().add(entity);
-            } else {
-                let mut target = <Self::RelationshipTarget as RelationshipTarget>::with_capacity(1);
-                target.collection_mut_risky().add(entity);
-                world.commands().entity(target_entity).insert(target);
-            }
+        if let Ok(mut entity_commands) = world.commands().get_entity(target_entity) {
+            // Deferring is necessary for batch mode
+            entity_commands
+                .entry::<Self::RelationshipTarget>()
+                .and_modify(move |mut relationship_target| {
+                    relationship_target.collection_mut_risky().add(entity);
+                })
+                .or_insert_with(|| {
+                    let mut target = Self::RelationshipTarget::with_capacity(1);
+                    target.collection_mut_risky().add(entity);
+                    target
+                });
         } else {
             warn!(
                 "{}The {}({target_entity:?}) relationship on entity {entity:?} relates to an entity that does not exist. The invalid {} relationship has been removed.",
@@ -323,6 +325,7 @@ pub enum RelationshipHookMode {
 
 #[cfg(test)]
 mod tests {
+    use crate::prelude::{ChildOf, Children};
     use crate::world::World;
     use crate::{component::Component, entity::Entity};
     use alloc::vec::Vec;
@@ -418,8 +421,6 @@ mod tests {
 
     #[test]
     fn parent_child_relationship_with_custom_relationship() {
-        use crate::prelude::ChildOf;
-
         #[derive(Component)]
         #[relationship(relationship_target = RelTarget)]
         struct Rel(Entity);
@@ -474,4 +475,34 @@ mod tests {
         assert!(world.get_entity(child).is_err());
         assert!(!world.entity(parent).contains::<RelTarget>());
     }
+
+    #[test]
+    fn spawn_batch_with_relationship() {
+        let mut world = World::new();
+        let parent = world.spawn_empty().id();
+        let children = world
+            .spawn_batch((0..10).map(|_| ChildOf(parent)))
+            .collect::<Vec<_>>();
+
+        for &child in &children {
+            assert!(world
+                .get::<ChildOf>(child)
+                .is_some_and(|child_of| child_of.parent() == parent));
+        }
+        assert!(world
+            .get::<Children>(parent)
+            .is_some_and(|children| children.len() == 10));
+    }
+
+    #[test]
+    fn insert_batch_with_relationship() {
+        let mut world = World::new();
+        let parent = world.spawn_empty().id();
+        let child = world.spawn_empty().id();
+        world.insert_batch([(child, ChildOf(parent))]);
+        world.flush();
+
+        assert!(world.get::<ChildOf>(child).is_some());
+        assert!(world.get::<Children>(parent).is_some());
+    }
 }