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

3 files changed, 83 insertions, 21 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(());
diff --git a/engine-ecs/src/query.rs b/engine-ecs/src/query.rs
index 555380e..c323092 100644
--- a/engine-ecs/src/query.rs
+++ b/engine-ecs/src/query.rs
@@ -17,11 +17,12 @@ use crate::tuple::{Concat, Tuple};
 use crate::uid::Uid;
 use crate::util::array_vec::ArrayVec;
 use crate::util::Array;
-use crate::{pair, World};
+use crate::World;
 
 mod reexports
 {
     pub use crate::component::storage::{
+        FoundByTraversal,
         SearchTerms as Terms,
         TermMetadata,
         Traversal,
@@ -268,10 +269,7 @@ impl_terms_builder! {
     {
         self.traverse.extend(traversals.clone());
 
-        self.present(
-            traversals
-                .map(|traversal| pair!({ traversal.relation }, *).id())
-        )
+        self
     }
 }
 
diff --git a/engine-ecs/src/query/term.rs b/engine-ecs/src/query/term.rs
index 1bcb3dd..453a750 100644
--- a/engine-ecs/src/query/term.rs
+++ b/engine-ecs/src/query/term.rs
@@ -7,6 +7,7 @@ use crate::component::{
     HandleMut as ComponentHandleMut,
 };
 use crate::query::{
+    FoundByTraversal,
     SearchResult,
     Term,
     TermMetadata,
@@ -212,14 +213,21 @@ where
             unreachable!();
         };
 
-        let Some(found_ent) = world.get_entity(traversal_result.found_ent) else {
-            unreachable!();
+        let ent = match traversal_result.found {
+            FoundByTraversal::SelfEntity => search_result.entity_handle.clone(),
+            FoundByTraversal::OtherEntity(found_ent_id) => {
+                let Some(found_ent) = world.get_entity(found_ent_id) else {
+                    unreachable!();
+                };
+
+                found_ent
+            }
         };
 
         TermsT::get_fields(
             world,
             &SearchResult {
-                entity_handle: found_ent,
+                entity_handle: ent,
                 traversal_results: search_result.traversal_results,
             },
             term_metadata
@@ -247,6 +255,18 @@ impl TraversalKindType for TraverseUp
 
 impl sealed::Sealed for TraverseUp {}
 
+pub struct TraverseSelfOrUp;
+
+impl TraversalKindType for TraverseSelfOrUp
+{
+    fn traversal_kind() -> TraversalKind
+    {
+        TraversalKind::SelfOrUp
+    }
+}
+
+impl sealed::Sealed for TraverseSelfOrUp {}
+
 mod sealed
 {
     pub trait Sealed {}