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:
parent
6ecaf9de53
commit
8bce3e46f2
@ -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);
|
||||
}
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user