Preserve extra data in RelationshipTarget with replace_related(_with_difference) (#19588)

# Objective

The methods and commands `replace_related` and
`replace_related_with_difference` may cause data stored at the
`RelationshipTarget` be lost when all original children are removed
before new children are added.

Part of https://github.com/bevyengine/bevy/issues/19589

## Solution

Fix the issue, either by removing the old children _after_ adding the
new ones and not _before_ (`replace_related_with_difference`) or by
taking the whole `RelationshipTarget` to modify it, not only the inner
collection (`replace_related`).

## Testing

I added a new test asserting the data is kept. I also added a general
test of these methods as they had none previously.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
This commit is contained in:
urben1680 2025-06-15 18:41:28 +02:00 committed by GitHub
parent 6ecaf9de53
commit 8bce3e46f2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -139,22 +139,26 @@ impl<'w> EntityWorldMut<'w> {
return self;
}
let Some(mut existing_relations) = self.get_mut::<R::RelationshipTarget>() else {
let Some(existing_relations) = self.get_mut::<R::RelationshipTarget>() else {
return self.add_related::<R>(related);
};
// We take the collection here so we can modify it without taking the component itself (this would create archetype move).
// We replace the component here with a dummy value so we can modify it without taking it (this would create archetype move).
// SAFETY: We eventually return the correctly initialized collection into the target.
let mut existing_relations = mem::replace(
existing_relations.collection_mut_risky(),
Collection::<R>::with_capacity(0),
let mut relations = mem::replace(
existing_relations.into_inner(),
<R as Relationship>::RelationshipTarget::from_collection_risky(
Collection::<R>::with_capacity(0),
),
);
let collection = relations.collection_mut_risky();
let mut potential_relations = EntityHashSet::from_iter(related.iter().copied());
let id = self.id();
self.world_scope(|world| {
for related in existing_relations.iter() {
for related in collection.iter() {
if !potential_relations.remove(related) {
world.entity_mut(related).remove::<R>();
}
@ -169,11 +173,9 @@ impl<'w> EntityWorldMut<'w> {
});
// SAFETY: The entities we're inserting will be the entities that were either already there or entities that we've just inserted.
existing_relations.clear();
existing_relations.extend_from_iter(related.iter().copied());
self.insert(R::RelationshipTarget::from_collection_risky(
existing_relations,
));
collection.clear();
collection.extend_from_iter(related.iter().copied());
self.insert(relations);
self
}
@ -239,11 +241,20 @@ impl<'w> EntityWorldMut<'w> {
assert_eq!(newly_related_entities, entities_to_relate, "`entities_to_relate` ({entities_to_relate:?}) didn't contain all entities that would end up related");
};
if !self.contains::<R::RelationshipTarget>() {
self.add_related::<R>(entities_to_relate);
match self.get_mut::<R::RelationshipTarget>() {
None => {
self.add_related::<R>(entities_to_relate);
return self;
};
return self;
}
Some(mut target) => {
// SAFETY: The invariants expected by this function mean we'll only be inserting entities that are already related.
let collection = target.collection_mut_risky();
collection.clear();
collection.extend_from_iter(entities_to_relate.iter().copied());
}
}
let this = self.id();
self.world_scope(|world| {
@ -252,32 +263,13 @@ impl<'w> EntityWorldMut<'w> {
}
for new_relation in newly_related_entities {
// We're changing the target collection manually so don't run the insert hook
// 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);
}
});
if !entities_to_relate.is_empty() {
if let Some(mut target) = self.get_mut::<R::RelationshipTarget>() {
// SAFETY: The invariants expected by this function mean we'll only be inserting entities that are already related.
let collection = target.collection_mut_risky();
collection.clear();
collection.extend_from_iter(entities_to_relate.iter().copied());
} else {
let mut empty =
<R::RelationshipTarget as RelationshipTarget>::Collection::with_capacity(
entities_to_relate.len(),
);
empty.extend_from_iter(entities_to_relate.iter().copied());
// SAFETY: We've just initialized this collection and we know there's no `RelationshipTarget` on `self`
self.insert(R::RelationshipTarget::from_collection_risky(empty));
}
}
self
}
@ -668,4 +660,61 @@ mod tests {
assert_eq!(world.entity(b).get::<ChildOf>(), None);
assert_eq!(world.entity(c).get::<ChildOf>(), None);
}
#[test]
fn replace_related_works() {
let mut world = World::new();
let child1 = world.spawn_empty().id();
let child2 = world.spawn_empty().id();
let child3 = world.spawn_empty().id();
let mut parent = world.spawn_empty();
parent.add_children(&[child1, child2]);
let child_value = ChildOf(parent.id());
let some_child = Some(&child_value);
parent.replace_children(&[child2, child3]);
let children = parent.get::<Children>().unwrap().collection();
assert_eq!(children, &[child2, child3]);
assert_eq!(parent.world().get::<ChildOf>(child1), None);
assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
assert_eq!(parent.world().get::<ChildOf>(child3), some_child);
parent.replace_children_with_difference(&[child3], &[child1, child2], &[child1]);
let children = parent.get::<Children>().unwrap().collection();
assert_eq!(children, &[child1, child2]);
assert_eq!(parent.world().get::<ChildOf>(child1), some_child);
assert_eq!(parent.world().get::<ChildOf>(child2), some_child);
assert_eq!(parent.world().get::<ChildOf>(child3), None);
}
#[test]
fn replace_related_keeps_data() {
#[derive(Component)]
#[relationship(relationship_target = Parent)]
pub struct Child(Entity);
#[derive(Component)]
#[relationship_target(relationship = Child)]
pub struct Parent {
#[relationship]
children: Vec<Entity>,
pub data: u8,
}
let mut world = World::new();
let child1 = world.spawn_empty().id();
let child2 = world.spawn_empty().id();
let mut parent = world.spawn_empty();
parent.add_related::<Child>(&[child1]);
parent.get_mut::<Parent>().unwrap().data = 42;
parent.replace_related_with_difference::<Child>(&[child1], &[child2], &[child2]);
let data = parent.get::<Parent>().unwrap().data;
assert_eq!(data, 42);
parent.replace_related::<Child>(&[child1]);
let data = parent.get::<Parent>().unwrap().data;
assert_eq!(data, 42);
}
}