EntityWorldMut methods do not automatically overwrite Relationship components

crates/bevy_ecs/macros/src/component.rs

@@ -792,6 +792,11 @@ fn derive_relationship(
                     #relationship_member: entity
                 }
             }
+
+            #[inline]
+            fn set_risky(&mut self, entity: Entity) {
+                self.#relationship_member = entity;
+            }
         }
     }))
 }

crates/bevy_ecs/src/relationship/mod.rs

@@ -82,6 +82,20 @@ pub trait Relationship: Component + Sized {
     /// Creates this [`Relationship`] from the given `entity`.
     fn from(entity: Entity) -> Self;
 
+    /// Changes the current [`Entity`] ID of the entity containing the [`RelationshipTarget`] to another one.
+    ///
+    /// This is useful for updating the relationship without overwriting other fields stored in `Self`.
+    ///
+    /// # Warning
+    ///
+    /// This should generally not be called by user code, as modifying the related entity could invalidate the
+    /// relationship. If this method is used, then the hooks [`on_replace`](Relationship::on_replace) have to
+    /// run before and [`on_insert`](Relationship::on_insert) after it.
+    /// This happens automatically when this method is called with [`EntityWorldMut::modify_component`].
+    ///
+    /// Prefer to use regular means of insertions when possible.
+    fn set_risky(&mut self, entity: Entity);
+
     /// The `on_insert` component hook that maintains the [`Relationship`] / [`RelationshipTarget`] connection.
     fn on_insert(
         mut world: DeferredWorld,

crates/bevy_ecs/src/relationship/related_methods.rs

@@ -6,7 +6,7 @@ use crate::{
         Relationship, RelationshipHookMode, RelationshipSourceCollection, RelationshipTarget,
     },
     system::{Commands, EntityCommands},
-    world::{EntityWorldMut, World},
+    world::{DeferredWorld, EntityWorldMut, World},
 };
 use bevy_platform::prelude::{Box, Vec};
 use core::{marker::PhantomData, mem};
@@ -42,7 +42,12 @@ impl<'w> EntityWorldMut<'w> {
         let id = self.id();
         self.world_scope(|world| {
             for related in related {
-                world.entity_mut(*related).insert(R::from(id));
+                world
+                    .entity_mut(*related)
+                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
+                        id,
+                        RelationshipHookMode::Run,
+                    );
             }
         });
         self
@@ -98,7 +103,12 @@ impl<'w> EntityWorldMut<'w> {
                         .collection_mut_risky()
                         .place(*related, index);
                 } else {
-                    world.entity_mut(*related).insert(R::from(id));
+                    world
+                        .entity_mut(*related)
+                        .modify_or_insert_relation_with_relationship_hook_mode::<R>(
+                            id,
+                            RelationshipHookMode::Run,
+                        );
                     world
                         .get_mut::<R::RelationshipTarget>(id)
                         .expect("hooks should have added relationship target")
@@ -165,10 +175,13 @@ impl<'w> EntityWorldMut<'w> {
             }
 
             for related in potential_relations {
-                // SAFETY: We'll manually be adjusting the contents of the parent to fit the final state.
+                // SAFETY: We'll manually be adjusting the contents of the `RelationshipTarget` to fit the final state.
                 world
                     .entity_mut(related)
-                    .insert_with_relationship_hook_mode(R::from(id), RelationshipHookMode::Skip);
+                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
+                        id,
+                        RelationshipHookMode::Skip,
+                    );
             }
         });
 
@@ -266,7 +279,10 @@ impl<'w> EntityWorldMut<'w> {
                 // We changed the target collection manually so don't run the insert hook
                 world
                     .entity_mut(*new_relation)
-                    .insert_with_relationship_hook_mode(R::from(this), RelationshipHookMode::Skip);
+                    .modify_or_insert_relation_with_relationship_hook_mode::<R>(
+                        this,
+                        RelationshipHookMode::Skip,
+                    );
             }
         });
 
@@ -352,6 +368,40 @@ impl<'w> EntityWorldMut<'w> {
 
         self
     }
+
+    fn modify_or_insert_relation_with_relationship_hook_mode<R: Relationship>(
+        &mut self,
+        entity: Entity,
+        relationship_hook_mode: RelationshipHookMode,
+    ) {
+        // Check if the relation edge holds additional data
+        if size_of::<R>() > size_of::<Entity>() {
+            self.assert_not_despawned();
+
+            let this = self.id();
+
+            let modified = self.world_scope(|world| {
+                let modified = DeferredWorld::from(&mut *world)
+                    .modify_component_with_relationship_hook_mode::<R, _>(
+                        this,
+                        relationship_hook_mode,
+                        |r| r.set_risky(entity),
+                    )
+                    .expect("entity access must be valid")
+                    .is_some();
+
+                world.flush();
+
+                modified
+            });
+
+            if modified {
+                return;
+            }
+        }
+
+        self.insert_with_relationship_hook_mode(R::from(entity), relationship_hook_mode);
+    }
 }
 
 impl<'a> EntityCommands<'a> {
@@ -689,17 +739,132 @@ mod tests {
     }
 
     #[test]
-    fn replace_related_keeps_data() {
+    fn add_related_keeps_relationship_data() {
+        #[derive(Component, PartialEq, Debug)]
+        #[relationship(relationship_target = Parent)]
+        struct Child {
+            #[relationship]
+            parent: Entity,
+            data: u8,
+        }
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Child)]
+        struct Parent(Vec<Entity>);
+
+        let mut world = World::new();
+        let parent1 = world.spawn_empty().id();
+        let parent2 = world.spawn_empty().id();
+        let child = world
+            .spawn(Child {
+                parent: parent1,
+                data: 42,
+            })
+            .id();
+
+        world.entity_mut(parent2).add_related::<Child>(&[child]);
+        assert_eq!(
+            world.get::<Child>(child),
+            Some(&Child {
+                parent: parent2,
+                data: 42
+            })
+        );
+    }
+
+    #[test]
+    fn insert_related_keeps_relationship_data() {
+        #[derive(Component, PartialEq, Debug)]
+        #[relationship(relationship_target = Parent)]
+        struct Child {
+            #[relationship]
+            parent: Entity,
+            data: u8,
+        }
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Child)]
+        struct Parent(Vec<Entity>);
+
+        let mut world = World::new();
+        let parent1 = world.spawn_empty().id();
+        let parent2 = world.spawn_empty().id();
+        let child = world
+            .spawn(Child {
+                parent: parent1,
+                data: 42,
+            })
+            .id();
+
+        world
+            .entity_mut(parent2)
+            .insert_related::<Child>(0, &[child]);
+        assert_eq!(
+            world.get::<Child>(child),
+            Some(&Child {
+                parent: parent2,
+                data: 42
+            })
+        );
+    }
+
+    #[test]
+    fn replace_related_keeps_relationship_data() {
+        #[derive(Component, PartialEq, Debug)]
+        #[relationship(relationship_target = Parent)]
+        struct Child {
+            #[relationship]
+            parent: Entity,
+            data: u8,
+        }
+
+        #[derive(Component)]
+        #[relationship_target(relationship = Child)]
+        struct Parent(Vec<Entity>);
+
+        let mut world = World::new();
+        let parent1 = world.spawn_empty().id();
+        let parent2 = world.spawn_empty().id();
+        let child = world
+            .spawn(Child {
+                parent: parent1,
+                data: 42,
+            })
+            .id();
+
+        world
+            .entity_mut(parent2)
+            .replace_related_with_difference::<Child>(&[], &[child], &[child]);
+        assert_eq!(
+            world.get::<Child>(child),
+            Some(&Child {
+                parent: parent2,
+                data: 42
+            })
+        );
+
+        world.entity_mut(parent1).replace_related::<Child>(&[child]);
+        assert_eq!(
+            world.get::<Child>(child),
+            Some(&Child {
+                parent: parent1,
+                data: 42
+            })
+        );
+    }
+
+    #[test]
+    fn replace_related_keeps_relationship_target_data() {
         #[derive(Component)]
         #[relationship(relationship_target = Parent)]
-        pub struct Child(Entity);
+        struct Child(Entity);
 
         #[derive(Component)]
         #[relationship_target(relationship = Child)]
-        pub struct Parent {
+        struct Parent {
             #[relationship]
             children: Vec<Entity>,
-            pub data: u8,
+            data: u8,
         }
 
         let mut world = World::new();

crates/bevy_ecs/src/world/deferred_world.rs

@@ -97,22 +97,28 @@ impl<'w> DeferredWorld<'w> {
     /// is mutable, prefer using [`get_mut`](DeferredWorld::get_mut).
     #[inline]
     #[track_caller]
-    pub(crate) fn modify_component<T: Component, R>(
+    pub(crate) fn modify_component_with_relationship_hook_mode<T: Component, R>(
         &mut self,
         entity: Entity,
+        relationship_hook_mode: RelationshipHookMode,
         f: impl FnOnce(&mut T) -> R,
     ) -> Result<Option<R>, EntityMutableFetchError> {
         // If the component is not registered, then it doesn't exist on this entity, so no action required.
         let Some(component_id) = self.component_id::<T>() else {
             return Ok(None);
         };
 
-        self.modify_component_by_id(entity, component_id, move |component| {
-            // SAFETY: component matches the component_id collected in the above line
-            let mut component = unsafe { component.with_type::<T>() };
+        self.modify_component_by_id_with_relationship_hook_mode(
+            entity,
+            component_id,
+            relationship_hook_mode,
+            move |component| {
+                // SAFETY: component matches the component_id collected in the above line
+                let mut component = unsafe { component.with_type::<T>() };
 
-            f(&mut component)
-        })
+                f(&mut component)
+            },
+        )
     }
 
     /// Temporarily removes a [`Component`] identified by the provided
@@ -127,14 +133,15 @@ impl<'w> DeferredWorld<'w> {
     /// If you do not need to ensure the above hooks are triggered, and your component
     /// is mutable, prefer using [`get_mut_by_id`](DeferredWorld::get_mut_by_id).
     ///
-    /// You should prefer the typed [`modify_component`](DeferredWorld::modify_component)
+    /// You should prefer the typed [`modify_component_with_relationship_hook_mode`](DeferredWorld::modify_component_with_relationship_hook_mode)
     /// whenever possible.
     #[inline]
     #[track_caller]
-    pub(crate) fn modify_component_by_id<R>(
+    pub(crate) fn modify_component_by_id_with_relationship_hook_mode<R>(
         &mut self,
         entity: Entity,
         component_id: ComponentId,
+        relationship_hook_mode: RelationshipHookMode,
         f: impl for<'a> FnOnce(MutUntyped<'a>) -> R,
     ) -> Result<Option<R>, EntityMutableFetchError> {
         let entity_cell = self.get_entity_mut(entity)?;
@@ -157,7 +164,7 @@ impl<'w> DeferredWorld<'w> {
                 entity,
                 [component_id].into_iter(),
                 MaybeLocation::caller(),
-                RelationshipHookMode::Run,
+                relationship_hook_mode,
             );
             if archetype.has_replace_observer() {
                 self.trigger_observers(
@@ -197,7 +204,7 @@ impl<'w> DeferredWorld<'w> {
                 entity,
                 [component_id].into_iter(),
                 MaybeLocation::caller(),
-                RelationshipHookMode::Run,
+                relationship_hook_mode,
             );
             if archetype.has_insert_observer() {
                 self.trigger_observers(

crates/bevy_ecs/src/world/mod.rs

@@ -1319,7 +1319,11 @@ impl World {
     ) -> Result<Option<R>, EntityMutableFetchError> {
         let mut world = DeferredWorld::from(&mut *self);
 
-        let result = world.modify_component(entity, f)?;
+        let result = world.modify_component_with_relationship_hook_mode(
+            entity,
+            RelationshipHookMode::Run,
+            f,
+        )?;
 
         self.flush();
         Ok(result)
@@ -1349,7 +1353,12 @@ impl World {
     ) -> Result<Option<R>, EntityMutableFetchError> {
         let mut world = DeferredWorld::from(&mut *self);
 
-        let result = world.modify_component_by_id(entity, component_id, f)?;
+        let result = world.modify_component_by_id_with_relationship_hook_mode(
+            entity,
+            component_id,
+            RelationshipHookMode::Run,
+            f,
+        )?;
 
         self.flush();
         Ok(result)