Adding test cases for ArchetypeSelectEntity and ArchetypeSelectEntity…

…Direct. Adding a contains function to World to make that test easier as well.
recatek · Jan 15, 2025 · b0fb2cb · b0fb2cb
1 parent 41a38a5
commit b0fb2cb
Show file tree

Hide file tree

Showing 3 changed files with 145 additions and 0 deletions.
diff --git a/macros/src/generate/world.rs b/macros/src/generate/world.rs
@@ -277,6 +277,14 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
                 }
 
                 impl WorldCanResolve<Entity<#Archetype>> for #World {
+                    #[inline(always)]
+                    fn resolve_contains(
+                        &self,
+                        entity: Entity<#Archetype>,
+                    ) -> bool {
+                        self.archetype::<#Archetype>().contains(entity)
+                    }
+
                     #[inline(always)]
                     fn resolve_direct(
                         &self,
@@ -295,6 +303,14 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
                 }
 
                 impl WorldCanResolve<EntityDirect<#Archetype>> for #World {
+                    #[inline(always)]
+                    fn resolve_contains(
+                        &self,
+                        entity: EntityDirect<#Archetype>,
+                    ) -> bool {
+                        self.archetype::<#Archetype>().contains(entity)
+                    }
+
                     #[inline(always)]
                     fn resolve_direct(
                         &self,
@@ -325,6 +341,20 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
             }
 
             impl WorldCanResolve<EntityAny> for #World {
+                #[inline(always)]
+                fn resolve_contains(
+                    &self,
+                    entity: EntityAny,
+                ) -> bool {
+                    match entity.try_into() {
+                        #(
+                            Ok(#ArchetypeSelectEntity::#Archetype(entity)) =>
+                                self.#archetype.contains(entity),
+                        )*
+                        Err(_) => panic!("invalid entity type"),
+                    }
+                }
+
                 #[inline(always)]
                 fn resolve_direct(
                     &self,
@@ -355,6 +385,20 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
             }
 
             impl WorldCanResolve<EntityDirectAny> for #World {
+                #[inline(always)]
+                fn resolve_contains(
+                    &self,
+                    entity: EntityDirectAny,
+                ) -> bool {
+                    match entity.try_into() {
+                        #(
+                            Ok(#ArchetypeSelectEntityDirect::#Archetype(entity)) =>
+                                self.#archetype.contains(entity),
+                        )*
+                        Err(_) => panic!("invalid entity type"),
+                    }
+                }
+
                 #[inline(always)]
                 fn resolve_direct(
                     &self,

diff --git a/src/traits.rs b/src/traits.rs
@@ -94,6 +94,21 @@ pub trait World: Sized {
         <Self as WorldHas<A>>::resolve_create_within_capacity(self, components)
     }
 
+    /// Returns true if this world contains the given entity key.
+    ///
+    ///
+    /// # Panics
+    ///
+    /// This will panic if given a dynamic key ([`EntityAny`] or [`EntityDirectAny`]) with an
+    /// archetype ID that is not known to this ECS world.
+    #[inline(always)]
+    fn contains<K: EntityKey>(&self, entity: K) -> bool
+    where
+        Self: WorldCanResolve<K>,
+    {
+        <Self as WorldCanResolve<K>>::resolve_contains(self, entity)
+    }
+
     /// If the entity exists in the world, this returns a direct entity pointing to its data.
     /// See [`EntityDirect`] and [`EntityDirectAny`] for more information.
     ///
@@ -637,6 +652,9 @@ pub trait BorrowHas<C>: Borrow {
 ///
 /// This is implemented for [`Entity`], [`EntityDirect`]. [`EntityAny`], and [`EntityDirectAny`].
 pub trait WorldCanResolve<K: EntityKey> {
+    #[doc(hidden)]
+    fn resolve_contains(&self, entity: K) -> bool;
+
     #[doc(hidden)]
     fn resolve_direct(&self, entity: K) -> Option<K::DirectOutput>;
 

diff --git a/tests/test_select.rs b/tests/test_select.rs
@@ -0,0 +1,83 @@
+use gecs::prelude::*;
+
+pub struct CompA(pub u32);
+pub struct CompB(pub u32);
+pub struct CompZ; // ZST
+
+ecs_world! {
+    ecs_archetype!(
+        ArchFoo,
+        CompA,
+        CompZ,
+    );
+
+    ecs_archetype!(
+        ArchBar,
+        CompB,
+        CompZ,
+    );
+}
+
+#[test]
+fn test_select() {
+    let mut world = EcsWorld::default();
+
+    let entity: Entity<ArchFoo> = world
+        .archetype_mut::<ArchFoo>()
+        .create((CompA(1), CompZ));
+
+    match entity.try_into() {
+        Ok(ArchetypeSelectEntity::ArchFoo(entity)) => assert!(world.contains(entity)),
+        Ok(ArchetypeSelectEntity::ArchBar(_)) => panic!(),
+        Err(_) => panic!(),
+    };
+}
+
+#[test]
+fn test_select_direct() {
+    let mut world = EcsWorld::default();
+
+    let entity = world
+        .archetype_mut::<ArchFoo>()
+        .create((CompA(1), CompZ));
+    let entity: EntityDirect<ArchFoo> = world.resolve_direct(entity).unwrap();
+
+    match entity.try_into() {
+        Ok(ArchetypeSelectEntityDirect::ArchFoo(entity)) => assert!(world.contains(entity)),
+        Ok(ArchetypeSelectEntityDirect::ArchBar(_)) => panic!(),
+        Err(_) => panic!(),
+    };
+}
+
+
+#[test]
+fn test_select_any() {
+    let mut world = EcsWorld::default();
+
+    let entity: EntityAny = world
+        .archetype_mut::<ArchFoo>()
+        .create((CompA(1), CompZ))
+        .into_any();
+
+    match entity.try_into() {
+        Ok(ArchetypeSelectEntity::ArchFoo(entity)) => assert!(world.contains(entity)),
+        Ok(ArchetypeSelectEntity::ArchBar(_)) => panic!(),
+        Err(_) => panic!(),
+    };
+}
+
+#[test]
+fn test_select_any_direct() {
+    let mut world = EcsWorld::default();
+
+    let entity = world
+        .archetype_mut::<ArchFoo>()
+        .create((CompA(1), CompZ));
+    let entity: EntityDirectAny = world.resolve_direct(entity).unwrap().into_any();
+
+    match entity.try_into() {
+        Ok(ArchetypeSelectEntityDirect::ArchFoo(entity)) => assert!(world.contains(entity)),
+        Ok(ArchetypeSelectEntityDirect::ArchBar(_)) => panic!(),
+        Err(_) => panic!(),
+    };
+}