feat(engine-ecs): add query traversal kind SelfOrUp

1 files changed, 57 insertions, 13 deletions

diff --git a/engine-ecs/src/component/storage.rs b/engine-ecs/src/component/storage.rs
index f6ab309..5746bea 100644
--- a/engine-ecs/src/component/storage.rs
+++ b/engine-ecs/src/component/storage.rs
@@ -80,6 +80,7 @@ impl<const TERM_CAP: usize> Traversal<TERM_CAP>
 pub enum TraversalKind
 {
     Up,
+    SelfOrUp,
 }
 
 #[derive(Debug)]
@@ -87,7 +88,15 @@ pub struct TraversalResult
 {
     pub relation: Uid,
     pub term_metadata: TermMetadata,
-    pub found_ent: Uid,
+    pub found: FoundByTraversal,
+}
+
+#[derive(Debug)]
+#[non_exhaustive]
+pub enum FoundByTraversal
+{
+    SelfEntity,
+    OtherEntity(Uid),
 }
 
 impl<const TERM_CAP: usize> SearchTerms<TERM_CAP>
@@ -869,20 +878,55 @@ fn do_traversals<const TERM_CAP: usize>(
     // TODO: Optimize traversals. Caching is needed especially
 
     for traversal in traversals {
-        if traversal.kind == TraversalKind::Up {
-            archetype
-                .get_matching_component_ids(traversal.relationship())
-                .find_map(|relationship| {
-                    traverse_dfs(
+        let traversal_terms = (traversal.terms)(traversal);
+
+        match traversal.kind {
+            TraversalKind::Up => {
+                archetype
+                    .get_matching_component_ids(traversal.relationship())
+                    .find_map(|relationship| {
+                        traverse_dfs(
+                            storage,
+                            relationship.target(),
+                            traversal,
+                            &traversal_terms,
+                            results,
+                        )
+                    })?;
+            }
+            TraversalKind::SelfOrUp => {
+                if traversal_terms.absent.iter().all(|absent_comp_id| {
+                    !archetype.contains_matching_component(*absent_comp_id)
+                }) && traversal_terms.archetype_contains_all_required(archetype)
+                    && do_traversals(
                         storage,
-                        relationship.target(),
-                        traversal,
-                        &(traversal.terms)(traversal),
+                        archetype,
+                        &traversal_terms.traverse,
                         results,
                     )
-                })?;
-        } else {
-            unimplemented!();
+                    .is_some()
+                {
+                    results.push(TraversalResult {
+                        relation: traversal.relation,
+                        term_metadata: traversal.term_metadata,
+                        found: FoundByTraversal::SelfEntity,
+                    });
+
+                    return Some(());
+                }
+
+                archetype
+                    .get_matching_component_ids(traversal.relationship())
+                    .find_map(|relationship| {
+                        traverse_dfs(
+                            storage,
+                            relationship.target(),
+                            traversal,
+                            &traversal_terms,
+                            results,
+                        )
+                    })?;
+            }
         }
     }
 
@@ -909,7 +953,7 @@ fn traverse_dfs<const TERM_CAP: usize>(
         results.push(TraversalResult {
             relation: traversal.relation,
             term_metadata: traversal.term_metadata,
-            found_ent: ent_id,
+            found: FoundByTraversal::OtherEntity(ent_id),
         });
 
         return Some(());