diff --git a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs index 013cdd63aa..633aeb9e66 100644 --- a/crates/bevy_ecs/src/relationship/relationship_source_collection.rs +++ b/crates/bevy_ecs/src/relationship/relationship_source_collection.rs @@ -116,6 +116,35 @@ impl RelationshipSourceCollection for SmallVec<[Entity; N]> { } } +impl RelationshipSourceCollection for Entity { + type SourceIter<'a> = core::iter::Once; + + fn with_capacity(_capacity: usize) -> Self { + Entity::PLACEHOLDER + } + + fn add(&mut self, entity: Entity) { + *self = entity; + } + + fn remove(&mut self, entity: Entity) { + if *self == entity { + *self = Entity::PLACEHOLDER; + } + } + + fn iter(&self) -> Self::SourceIter<'_> { + core::iter::once(*self) + } + + fn len(&self) -> usize { + if *self == Entity::PLACEHOLDER { + return 0; + } + 1 + } +} + #[cfg(test)] mod tests { use super::*; @@ -184,4 +213,58 @@ mod tests { let collection = rel_target.collection(); assert_eq!(collection, &SmallVec::from_buf([a])); } + + #[test] + fn entity_relationship_source_collection() { + #[derive(Component)] + #[relationship(relationship_target = RelTarget)] + struct Rel(Entity); + + #[derive(Component)] + #[relationship_target(relationship = Rel)] + struct RelTarget(Entity); + + let mut world = World::new(); + let a = world.spawn_empty().id(); + let b = world.spawn_empty().id(); + + world.entity_mut(a).insert(Rel(b)); + + let rel_target = world.get::(b).unwrap(); + let collection = rel_target.collection(); + assert_eq!(collection, &a); + } + + #[test] + fn one_to_one_relationships() { + #[derive(Component)] + #[relationship(relationship_target = Below)] + struct Above(Entity); + + #[derive(Component)] + #[relationship_target(relationship = Above)] + struct Below(Entity); + + let mut world = World::new(); + let a = world.spawn_empty().id(); + let b = world.spawn_empty().id(); + + world.entity_mut(a).insert(Above(b)); + assert_eq!(a, world.get::(b).unwrap().0); + + // Verify removing target removes relationship + world.entity_mut(b).remove::(); + assert!(world.get::(a).is_none()); + + // Verify removing relationship removes target + world.entity_mut(a).insert(Above(b)); + world.entity_mut(a).remove::(); + assert!(world.get::(b).is_none()); + + // Actually - a is above c now! Verify relationship was updated correctly + let c = world.spawn_empty().id(); + world.entity_mut(a).insert(Above(c)); + assert!(world.get::(b).is_none()); + assert_eq!(a, world.get::(c).unwrap().0); + } }