From 7a1972ed3d157168626930aca94552626ef933d1 Mon Sep 17 00:00:00 2001
From: Dmytro Banin <dima_banin@hotmail.com>
Date: Mon, 3 Mar 2025 23:57:35 -0800
Subject: [PATCH] One-to-One Relationships (#18087)

# Objective

Minimal implementation of directed one-to-one relationships via
implementing `RelationshipSourceCollection` for `Entity`.

Now you can do

```rust
#[derive(Component)]
#[relationship(relationship_target = Below)]
pub struct Above(Entity);

#[derive(Component)]
#[relationship_target(relationship = Above)]
pub struct Below(Entity);
```

## Future Work

It would be nice if the relationships could be fully symmetrical in the
future - in the example above, since `Above` is the source of truth you
can't add `Below` to an entity and have `Above` added automatically.

## Testing

Wrote unit tests for new relationship sources and and verified
adding/removing relationships maintains connection as expected.
---
 .../relationship_source_collection.rs         | 83 +++++++++++++++++++
 1 file changed, 83 insertions(+)

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<const N: usize> RelationshipSourceCollection for SmallVec<[Entity; N]> {
     }
 }
 
+impl RelationshipSourceCollection for Entity {
+    type SourceIter<'a> = core::iter::Once<Entity>;
+
+    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::<RelTarget>(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::<Below>(b).unwrap().0);
+
+        // Verify removing target removes relationship
+        world.entity_mut(b).remove::<Below>();
+        assert!(world.get::<Above>(a).is_none());
+
+        // Verify removing relationship removes target
+        world.entity_mut(a).insert(Above(b));
+        world.entity_mut(a).remove::<Above>();
+        assert!(world.get::<Below>(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::<Below>(b).is_none());
+        assert_eq!(a, world.get::<Below>(c).unwrap().0);
+    }
 }