feat(ecs): replace Relationship component with pair UID support

3 files changed, 168 insertions, 24 deletions

diff --git a/ecs/src/component/storage.rs b/ecs/src/component/storage.rs
index 53f51f2..9cf4433 100644
--- a/ecs/src/component/storage.rs
+++ b/ecs/src/component/storage.rs
@@ -33,14 +33,43 @@ pub struct ArchetypeSearchTerms<'a>
 
 impl ArchetypeSearchTerms<'_>
 {
-    fn excluded_contains(&self, uid: Uid) -> bool
+    fn excluded_contains(&self, comp_id: Uid) -> bool
     {
-        self.excluded_components.binary_search(&uid).is_ok()
+        let comp_id_kind = comp_id.kind();
+
+        debug_assert!(
+            comp_id_kind == UidKind::Component
+                || (comp_id_kind == UidKind::Pair
+                    && comp_id.target_component() != Uid::wildcard())
+        );
+
+        let is_found = self.excluded_components.binary_search(&comp_id).is_ok();
+
+        if !is_found && comp_id_kind == UidKind::Pair {
+            return self.excluded_components.iter().any(|excluded_comp_id| {
+                excluded_comp_id.kind() == UidKind::Pair
+                    && excluded_comp_id.has_same_relation_as(comp_id)
+                    && excluded_comp_id.target_component() == Uid::wildcard()
+            });
+        }
+
+        is_found
     }
 
-    fn required_contains(&self, uid: Uid) -> bool
+    fn contains_conflicting(&self) -> bool
     {
-        self.required_components.binary_search(&uid).is_ok()
+        self.excluded_components.iter().any(|excluded_comp_id| {
+            self.required_components
+                .binary_search(excluded_comp_id)
+                .is_ok()
+        })
+    }
+
+    fn archetype_contains_all_required(&self, archetype: &Archetype) -> bool
+    {
+        self.required_components
+            .iter()
+            .all(|comp_id| archetype.contains_matching_component(*comp_id))
     }
 }
 
@@ -61,11 +90,7 @@ impl Storage
     {
         let archetype_id = ArchetypeId::new(&search_terms.required_components);
 
-        if search_terms
-            .excluded_components
-            .iter()
-            .any(|excluded_comp_id| search_terms.required_contains(*excluded_comp_id))
-        {
+        if search_terms.contains_conflicting() {
             return ArchetypeRefIter {
                 storage: self,
                 pre_iter: Either::B(Vec::new().into_iter()),
@@ -81,7 +106,15 @@ impl Storage
                 .borrow_mut()
                 .push(ImaginaryArchetype {
                     id: archetype_id,
-                    component_ids: search_terms.required_components.to_vec(),
+                    component_ids: search_terms
+                        .required_components
+                        .iter()
+                        .copied()
+                        .filter(|required_comp_id| {
+                            required_comp_id.kind() != UidKind::Pair
+                                || required_comp_id.target_component() != Uid::wildcard()
+                        })
+                        .collect(),
                 });
 
             let found_archetypes = self.find_all_archetype_with_comps(&search_terms);
@@ -175,7 +208,7 @@ impl Storage
 
         if archetype_node
             .archetype()
-            .has_component_with_id(component_id)
+            .contains_component_with_exact_id(component_id)
         {
             return Err(Error::ComponentAlreadyInEntity {
                 entity: entity_uid,
@@ -266,7 +299,7 @@ impl Storage
 
         if !archetype_node
             .archetype()
-            .has_component_with_id(component_id)
+            .contains_component_with_exact_id(component_id)
         {
             return Err(Error::ComponentNotFoundInEntity {
                 entity: entity_uid,
@@ -367,11 +400,7 @@ impl Storage
                 continue;
             }
 
-            if !search_terms
-                .required_components
-                .iter()
-                .all(|comp_id| node.archetype().has_component_with_id(*comp_id))
-            {
+            if !search_terms.archetype_contains_all_required(node.archetype()) {
                 continue;
             }
 
diff --git a/ecs/src/component/storage/archetype.rs b/ecs/src/component/storage/archetype.rs
index b8ac826..10a665e 100644
--- a/ecs/src/component/storage/archetype.rs
+++ b/ecs/src/component/storage/archetype.rs
@@ -1,5 +1,8 @@
 use std::any::Any;
+use std::array::IntoIter as ArrayIntoIter;
 use std::hash::{DefaultHasher, Hash, Hasher};
+use std::iter::{Enumerate, Filter, Map, RepeatN, Zip};
+use std::option::IntoIter as OptionIntoIter;
 use std::slice::Iter as SliceIter;
 
 use hashbrown::HashMap;
@@ -7,7 +10,7 @@ use hashbrown::HashMap;
 use crate::component::Metadata as ComponentMetadata;
 use crate::lock::Lock;
 use crate::uid::{Kind as UidKind, Uid};
-use crate::util::HashMapExt;
+use crate::util::{Either, HashMapExt};
 
 #[derive(Debug)]
 pub struct Archetype
@@ -117,8 +120,54 @@ impl Archetype
         self.component_index_lookup.len()
     }
 
+    pub fn get_matching_component_indices(
+        &self,
+        component_id: Uid,
+    ) -> MatchingComponentIter
+    {
+        assert!(
+            component_id.kind() == UidKind::Component
+                || component_id.kind() == UidKind::Pair
+        );
+
+        if component_id.kind() == UidKind::Pair
+            && component_id.target_component() == Uid::wildcard()
+        {
+            return MatchingComponentIter {
+                inner: Either::A(
+                    self.component_ids
+                        .iter()
+                        .enumerate()
+                        .zip(std::iter::repeat_n(component_id, self.component_ids.len()))
+                        .filter(
+                            (|((_, other_comp_id), component_id)| {
+                                other_comp_id.kind() == UidKind::Pair
+                                    && other_comp_id.has_same_relation_as(*component_id)
+                            })
+                                as MatchingComponentIterFilterFn,
+                        )
+                        .map(|((index, other_comp_id), _)| (*other_comp_id, index)),
+                ),
+            };
+        }
+
+        MatchingComponentIter {
+            inner: Either::B(
+                [component_id]
+                    .into_iter()
+                    .zip(self.get_index_for_component(component_id).into_iter()),
+            ),
+        }
+    }
+
     pub fn get_index_for_component(&self, component_id: Uid) -> Option<usize>
     {
+        assert!(
+            component_id.kind() == UidKind::Component
+                || (component_id.kind() == UidKind::Pair
+                    && component_id.target_component() != Uid::wildcard())
+        );
+
         self.component_index_lookup.get(&component_id).copied()
     }
 
@@ -132,14 +181,69 @@ impl Archetype
         self.component_ids.iter().copied()
     }
 
-    pub fn has_component_with_id(&self, component_id: Uid) -> bool
+    pub fn contains_matching_component(&self, component_id: Uid) -> bool
+    {
+        let component_id_kind = component_id.kind();
+
+        debug_assert!(
+            component_id_kind == UidKind::Component || component_id_kind == UidKind::Pair
+        );
+
+        if component_id.kind() == UidKind::Pair
+            && component_id.target_component() == Uid::wildcard()
+        {
+            return self.component_ids.iter().any(|other_comp_id| {
+                other_comp_id.kind() == UidKind::Pair
+                    && other_comp_id.has_same_relation_as(component_id)
+            });
+        }
+
+        self.contains_component_with_exact_id(component_id)
+    }
+
+    pub fn contains_component_with_exact_id(&self, component_id: Uid) -> bool
     {
-        debug_assert_eq!(component_id.kind(), UidKind::Component);
+        let component_id_kind = component_id.kind();
+
+        debug_assert!(
+            component_id_kind == UidKind::Component
+                || (component_id_kind == UidKind::Pair
+                    && component_id.target_component() != Uid::wildcard())
+        );
 
         self.component_index_lookup.contains_key(&component_id)
     }
 }
 
+type MatchingComponentIterFilterFn = fn(&((usize, &Uid), Uid)) -> bool;
+
+type MatchingComponentIterMapFn = fn(((usize, &Uid), Uid)) -> (Uid, usize);
+
+#[derive(Debug)]
+pub struct MatchingComponentIter<'archetype>
+{
+    inner: Either<
+        Map<
+            Filter<
+                Zip<Enumerate<SliceIter<'archetype, Uid>>, RepeatN<Uid>>,
+                MatchingComponentIterFilterFn,
+            >,
+            MatchingComponentIterMapFn,
+        >,
+        Zip<ArrayIntoIter<Uid, 1>, OptionIntoIter<usize>>,
+    >,
+}
+
+impl Iterator for MatchingComponentIter<'_>
+{
+    type Item = (Uid, usize);
+
+    fn next(&mut self) -> Option<Self::Item>
+    {
+        self.inner.next()
+    }
+}
+
 #[derive(Debug)]
 pub struct EntityIter<'archetype>
 {
diff --git a/ecs/src/component/storage/graph.rs b/ecs/src/component/storage/graph.rs
index d38223a..29fa937 100644
--- a/ecs/src/component/storage/graph.rs
+++ b/ecs/src/component/storage/graph.rs
@@ -147,7 +147,11 @@ impl Graph
         let uniq_comp_id = target_node
             .archetype()
             .component_ids_sorted()
-            .find(|id| !subset_node.archetype().has_component_with_id(*id))
+            .find(|id| {
+                !subset_node
+                    .archetype()
+                    .contains_component_with_exact_id(*id)
+            })
             .unwrap();
 
         subset_node
@@ -177,7 +181,7 @@ impl Graph
                 .find(|other_archetype_comp_id| {
                     !target_node
                         .archetype()
-                        .has_component_with_id(*other_archetype_comp_id)
+                        .contains_component_with_exact_id(*other_archetype_comp_id)
                 })
                 .or_else(|| {
                     if target_node.archetype().component_cnt() != 0 {
@@ -204,7 +208,11 @@ impl Graph
         let extra_comp_id = superset_node
             .archetype()
             .component_ids_unsorted()
-            .find(|comp_id| !target_node.archetype().has_component_with_id(*comp_id))
+            .find(|comp_id| {
+                !target_node
+                    .archetype()
+                    .contains_component_with_exact_id(*comp_id)
+            })
             .expect("Archetype should contain one extra component ID");
 
         superset_node
@@ -242,7 +250,10 @@ impl ArchetypeNode
         insert_fn: impl FnOnce() -> ArchetypeEdges,
     ) -> &mut ArchetypeEdges
     {
-        debug_assert_eq!(component_id.kind(), UidKind::Component);
+        debug_assert!(matches!(
+            component_id.kind(),
+            UidKind::Component | UidKind::Pair
+        ));
 
         self.edges.entry(component_id).or_insert_with(insert_fn)
     }