1 files changed, 124 insertions, 12 deletions

diff --git a/ecs/src/relationship.rs b/ecs/src/relationship.rs
index e49bde7..1a8c842 100644
--- a/ecs/src/relationship.rs
+++ b/ecs/src/relationship.rs
@@ -13,10 +13,11 @@ use crate::system::{ComponentRefMut, Input as SystemInput};
 use crate::type_name::TypeName;
 use crate::World;
 
+/// A relationship to one or more targets.
 #[derive(Debug)]
 pub struct Relationship<Kind, ComponentT: Component>
 {
-    entity_uid: EntityUid,
+    entity_uid: SingleOrMultiple<EntityUid>,
     _pd: PhantomData<(Kind, ComponentT)>,
 }
 
@@ -24,9 +25,22 @@ impl<Kind, ComponentT> Relationship<Kind, ComponentT>
 where
     ComponentT: Component,
 {
+    /// Creates a new `Relationship` with a single target.
     pub fn new(entity_uid: EntityUid) -> Self
     {
-        Self { entity_uid, _pd: PhantomData }
+        Self {
+            entity_uid: SingleOrMultiple::Single(entity_uid),
+            _pd: PhantomData,
+        }
+    }
+
+    /// Creates a new `Relationship` with multiple targets.
+    pub fn new_multiple(entity_uids: impl IntoIterator<Item = EntityUid>) -> Self
+    {
+        Self {
+            entity_uid: SingleOrMultiple::Multiple(entity_uids.into_iter().collect()),
+            _pd: PhantomData,
+        }
     }
 }
 
@@ -117,16 +131,16 @@ impl<'rel_comp, Kind, ComponentT> Relation<'rel_comp, Kind, ComponentT>
 where
     ComponentT: Component,
 {
-    /// Retrieves the related-to component.
-    pub fn get(&self) -> Option<ComponentRefMut<'_, ComponentT>>
+    /// Returns the component of the target at the specified index.
+    pub fn get(&self, index: usize) -> Option<ComponentRefMut<'_, ComponentT>>
     {
-        let archetype = self
-            .component_storage_lock
-            .get_entity_archetype(self.relationship_comp.entity_uid)?;
+        let target = self.get_target(index)?;
+
+        let archetype = self.component_storage_lock.get_entity_archetype(*target)?;
 
         let entity = archetype
-            .get_entity(self.relationship_comp.entity_uid)
-            .expect("Entity is gone");
+            .get_entity(*target)
+            .expect("Target entity is gone from archetype");
 
         let component_index =
             archetype.get_index_for_component(&ComponentId::of::<ComponentT>())?;
@@ -147,9 +161,107 @@ where
         Some(component)
     }
 
-    /// Changes the related-to entity to the entity with the given UID:
-    pub fn set_entity(&mut self, entity_uid: EntityUid)
+    /// Returns a reference to the target at the specified index.
+    pub fn get_target(&self, index: usize) -> Option<&EntityUid>
+    {
+        match &self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get(index),
+            _ => None,
+        }
+    }
+
+    /// Returns a mutable reference to the target at the specified index.
+    pub fn get_target_mut(&mut self, index: usize) -> Option<&mut EntityUid>
+    {
+        match &mut self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(entity_uid) if index == 0 => Some(entity_uid),
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.get_mut(index),
+            _ => None,
+        }
+    }
+
+    /// Adds a target to the relationship.
+    pub fn add_target(&mut self, entity_uid: EntityUid)
+    {
+        match &mut self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(prev_entity_uid) => {
+                self.relationship_comp.entity_uid =
+                    SingleOrMultiple::Multiple(vec![*prev_entity_uid, entity_uid]);
+            }
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.push(entity_uid),
+        }
+    }
+
+    /// Removes a target to the relationship, returning it.
+    pub fn remove_target(&mut self, index: usize) -> Option<EntityUid>
+    {
+        match &mut self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(entity_uid) => {
+                let prev_entity_uid = *entity_uid;
+
+                self.relationship_comp.entity_uid =
+                    SingleOrMultiple::Multiple(Vec::new());
+
+                Some(prev_entity_uid)
+            }
+            SingleOrMultiple::Multiple(entity_uids) => {
+                if index >= entity_uids.len() {
+                    return None;
+                }
+
+                Some(entity_uids.remove(index))
+            }
+        }
+    }
+
+    pub fn target_count(&self) -> usize
+    {
+        match &self.relationship_comp.entity_uid {
+            SingleOrMultiple::Single(_) => 1,
+            SingleOrMultiple::Multiple(entity_uids) => entity_uids.len(),
+        }
+    }
+
+    /// Returns a iterator of the components of the targets of this relationship.
+    pub fn iter(&self) -> TargetComponentIter<'_, 'rel_comp, Kind, ComponentT>
+    {
+        TargetComponentIter { relation: self, index: 0 }
+    }
+}
+
+/// Iterator of the components of the targets of a relationship.
+pub struct TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+where
+    Kind: 'static,
+    ComponentT: Component,
+{
+    relation: &'relationship Relation<'rel_comp, Kind, ComponentT>,
+    index: usize,
+}
+
+impl<'relationship, 'rel_comp, Kind, ComponentT> Iterator
+    for TargetComponentIter<'relationship, 'rel_comp, Kind, ComponentT>
+where
+    'relationship: 'rel_comp,
+    Kind: 'static,
+    ComponentT: Component,
+{
+    type Item = ComponentRefMut<'rel_comp, ComponentT>;
+
+    fn next(&mut self) -> Option<Self::Item>
     {
-        self.relationship_comp.entity_uid = entity_uid;
+        let index = self.index;
+
+        self.index += 1;
+
+        self.relation.get(index)
     }
 }
+
+#[derive(Debug)]
+enum SingleOrMultiple<Value>
+{
+    Single(Value),
+    Multiple(Vec<Value>),
+}