- Adding an ecs_iter_remove! macro that will remove the currently ite…

…ration's entity if the closure returns true

macros/src/generate/query.rs

@@ -5,7 +5,14 @@ use proc_macro2::{Ident, Span, TokenStream};
 use quote::{format_ident, quote, quote_spanned};
 
 use crate::data::{DataArchetype, DataWorld};
-use crate::parse::{ParseQueryFind, ParseQueryIter, ParseQueryParam, ParseQueryParamType};
+
+use crate::parse::{
+    ParseQueryFind, //.
+    ParseQueryIter,
+    ParseQueryIterRemove,
+    ParseQueryParam,
+    ParseQueryParamType,
+};
 
 // NOTE: We should avoid using panics to express errors in queries when generating.
 // Doing so will attribute the error to the ecs_world! declaration (due to the redirect
@@ -18,7 +25,10 @@ pub enum FetchMode {
 }
 
 #[allow(non_snake_case)]
-pub fn generate_query_find(mode: FetchMode, query: ParseQueryFind) -> syn::Result<TokenStream> {
+pub fn generate_query_find(
+    mode: FetchMode, //.
+    query: ParseQueryFind,
+) -> syn::Result<TokenStream> {
     let world_data = DataWorld::from_base64(&query.world_data);
     let bound_params = bind_query_params(&world_data, &query.params)?;
 
@@ -203,7 +213,10 @@ fn find_bind_borrow(param: &ParseQueryParam) -> TokenStream {
 }
 
 #[allow(non_snake_case)]
-pub fn generate_query_iter(mode: FetchMode, query: ParseQueryIter) -> syn::Result<TokenStream> {
+pub fn generate_query_iter(
+    mode: FetchMode, //.
+    query: ParseQueryIter,
+) -> syn::Result<TokenStream> {
     let world_data = DataWorld::from_base64(&query.world_data);
     let bound_params = bind_query_params(&world_data, &query.params)?;
 
@@ -283,6 +296,96 @@ pub fn generate_query_iter(mode: FetchMode, query: ParseQueryIter) -> syn::Resul
     }
 }
 
+#[allow(non_snake_case)]
+pub fn generate_query_iter_remove(
+    mode: FetchMode,
+    query: ParseQueryIterRemove,
+) -> syn::Result<TokenStream> {
+    let world_data = DataWorld::from_base64(&query.world_data);
+    let bound_params = bind_query_params(&world_data, &query.params)?;
+
+    // NOTE: Beyond this point, query.params is only safe to use for information that
+    // does not change depending on the type of the parameter (e.g. mutability). Anything
+    // that might change after OneOf binding etc. must use the bound query params in
+    // bound_params for the given archetype. Note that it's faster to use query.params
+    // where available, since it avoids redundant computation for each archetype.
+
+    // TODO PERF: We could avoid binding entirely if we know that the params have no OneOf.
+
+    // Variables and fields
+    let world = &query.world;
+    let body = &query.body;
+    let arg = query.params.iter().map(to_name).collect::<Vec<_>>();
+
+    // Special cases
+    let maybe_mut = query.params.iter().map(to_maybe_mut).collect::<Vec<_>>();
+
+    let mut queries = Vec::<TokenStream>::new();
+    for archetype in world_data.archetypes {
+        debug_assert!(archetype.build_data.is_none());
+
+        if let Some(bound_params) = bound_params.get(&archetype.name) {
+            // Types and traits
+            let Archetype = format_ident!("{}", archetype.name);
+            let Type = bound_params
+                .iter()
+                .map(|p| to_type(p, &archetype))
+                .collect::<Vec<_>>(); // Bind-dependent!
+
+            #[rustfmt::skip]
+            let get_archetype = match mode {
+                FetchMode::Borrow => panic!("borrow unsupported for iter_remove"),
+                FetchMode::Mut => quote!(#world.archetype_mut::<#Archetype>()),
+            };
+
+            #[rustfmt::skip]
+            let get_slices = match mode {
+                FetchMode::Borrow => panic!("borrow unsupported for iter_remove"),
+                FetchMode::Mut => quote!(archetype.get_all_slices_mut()),
+            };
+
+            #[rustfmt::skip]
+            let bind = match mode {
+                FetchMode::Borrow => panic!("borrow unsupported for iter_remove"),
+                FetchMode::Mut => bound_params.iter().map(iter_bind_mut).collect::<Vec<_>>(),
+            };
+
+            queries.push(quote!(
+                {
+                    // Alias the current archetype for use in the closure
+                    type MatchedArchetype = #Archetype;
+                    // The closure needs to be made per-archetype because of OneOf types
+                    let mut closure = //FnMut(#(&#maybe_mut #Type),*) -> bool
+                        |#(#arg: &#maybe_mut #Type),*| #body;
+
+                    let archetype = #get_archetype;
+                    let version = archetype.version();
+                    let len = archetype.len();
+
+                    // Iterate in reverse order to still visit each entity once.
+                    // Note: This assumes that we remove entities by swapping.
+                    for idx in (0..len).rev() {
+                        let slices = #get_slices;
+                        if closure(#(#bind),*) {
+                            let entity = slices.entity[idx];
+                            archetype.destroy(entity);
+                        }
+                    }
+                }
+            ));
+        }
+    }
+
+    if queries.is_empty() {
+        Err(syn::Error::new_spanned(
+            world,
+            "query matched no archetypes in world",
+        ))
+    } else {
+        Ok(quote!(#(#queries)*))
+    }
+}
+
 #[rustfmt::skip]
 fn iter_bind_mut(param: &ParseQueryParam) -> TokenStream {
     match &param.param_type {

macros/src/generate/world.rs

@@ -68,6 +68,7 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
     let __ecs_find_borrow_unique = format_ident!("__ecs_find_borrow_{}", unique_hash);
     let __ecs_iter_unique = format_ident!("__ecs_iter_{}", unique_hash);
     let __ecs_iter_borrow_unique = format_ident!("__ecs_iter_borrow_{}", unique_hash);
+    let __ecs_iter_remove_unique = format_ident!("__ecs_iter_remove_{}", unique_hash);
 
     quote!(
         #( pub use #ecs_world_sealed::#Archetype; )*
@@ -434,6 +435,14 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
             }
         }
 
+        #[macro_export]
+        #[doc(hidden)]
+        macro_rules! #__ecs_iter_remove_unique {
+            ($($args:tt)*) => {
+                ::gecs::__internal::__ecs_iter_remove!(#WORLD_DATA, $($args)*);
+            }
+        }
+
         #[doc(inline)]
         pub use #__ecs_find_unique as ecs_find;
         #[doc(inline)]
@@ -442,6 +451,8 @@ pub fn generate_world(world_data: &DataWorld, raw_input: &str) -> TokenStream {
         pub use #__ecs_iter_unique as ecs_iter;
         #[doc(inline)]
         pub use #__ecs_iter_borrow_unique as ecs_iter_borrow;
+        #[doc(inline)]
+        pub use #__ecs_iter_remove_unique as ecs_iter_remove;
     )
 }
 

macros/src/lib.rs

@@ -79,3 +79,14 @@ pub fn __ecs_iter_borrow(args: TokenStream) -> TokenStream {
         Err(err) => err.into_compile_error().into(),
     }
 }
+
+#[proc_macro]
+#[doc(hidden)]
+pub fn __ecs_iter_remove(args: TokenStream) -> TokenStream {
+    let query_parse = parse_macro_input!(args as ParseQueryIterRemove);
+
+    match generate::generate_query_iter_remove(FetchMode::Mut, query_parse) {
+        Ok(tokens) => tokens.into(),
+        Err(err) => err.into_compile_error().into(),
+    }
+}

macros/src/parse/query.rs

@@ -31,6 +31,14 @@ pub struct ParseQueryIter {
     pub body: Expr,
 }
 
+#[derive(Debug)]
+pub struct ParseQueryIterRemove {
+    pub world_data: String,
+    pub world: Expr,
+    pub params: Vec<ParseQueryParam>,
+    pub body: Expr,
+}
+
 #[derive(Clone, Debug)]
 pub struct ParseQueryParam {
     pub name: Ident,
@@ -114,6 +122,33 @@ impl Parse for ParseQueryIter {
     }
 }
 
+impl Parse for ParseQueryIterRemove {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        // Parse out the hidden serialized world data
+        let world_data = input.parse::<LitStr>()?;
+        input.parse::<Comma>()?;
+
+        // Parse out the meta-arguments for the query
+        let world = input.parse()?;
+        input.parse::<Comma>()?;
+
+        // Parse out the closure arguments
+        input.parse::<Token![|]>()?;
+        let params = parse_params(&input)?;
+        input.parse::<Token![|]>()?;
+
+        // Parse the rest of the body, including the braces (if any)
+        let body = input.parse::<Expr>()?;
+
+        Ok(Self {
+            world_data: world_data.value(),
+            world,
+            params,
+            body,
+        })
+    }
+}
+
 impl Parse for ParseQueryParam {
     fn parse(input: ParseStream) -> syn::Result<Self> {
         // Parse the name and following : token

src/lib.rs

@@ -603,7 +603,7 @@ mod macros {
 
     /// Variant of `ecs_iter!` that runtime-borrows data, for use with a non-mut world reference.
     ///
-    /// See [`ecs_iter`] for more information on find queries.
+    /// See [`ecs_iter`] for more information on iter queries.
     ///
     /// This version borrows each archetype's data on a component-by-component basis at runtime
     /// rather than at compile-time, allowing for situations where compile-time borrow checking
@@ -619,6 +619,66 @@ mod macros {
     macro_rules! ecs_iter_borrow {
         (...) => {};
     }
+
+    /// Variant of `ecs_iter!` that will remove the current entity if the closure returns `true`.
+    ///
+    /// See [`ecs_iter`] for more information on iter queries.
+    ///
+    /// This version works similarly to [`ecs_iter`], but if the inner closure returns `true`,
+    /// the iteration will immediately remove that entity after that iteration step. The entity
+    /// and its handle are not preserved after this process. Note that this iterates the world
+    /// in a different order from the normal `ecs_iter!` (which should not be relied upon for
+    /// deterministic iteration in any case. This is also slightly slower than `ecs_iter!`.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use gecs::prelude::*;
+    ///
+    /// pub struct CompA(pub u32);
+    /// pub struct CompB(pub u32);
+    /// pub struct CompC(pub u32);
+    ///
+    /// ecs_world! {
+    ///     ecs_archetype!(ArchFoo, 100, CompA, CompB);
+    ///     ecs_archetype!(ArchBar, 100, CompA, CompC);
+    /// }
+    ///
+    /// fn main() {
+    ///     let mut world = EcsWorld::default();
+    ///
+    ///     world.archetype_mut::<ArchFoo>().create((CompA(1), CompB(10)));
+    ///     world.archetype_mut::<ArchFoo>().create((CompA(2), CompB(20)));
+    ///     world.archetype_mut::<ArchFoo>().create((CompA(3), CompB(30)));
+    ///
+    ///     world.archetype_mut::<ArchBar>().create((CompA(4), CompC(10)));
+    ///     world.archetype_mut::<ArchBar>().create((CompA(5), CompC(10)));
+    ///     world.archetype_mut::<ArchBar>().create((CompA(6), CompC(10)));
+    ///
+    ///     let mut vec_a = Vec::<u32>::new();
+    ///     let mut vec_b = Vec::<u32>::new();
+    ///
+    ///     ecs_iter_remove!(world, |comp_a: &CompA| {
+    ///         if comp_a.0 & 1 == 0 {
+    ///             vec_a.push(comp_a.0);
+    ///             true // True to remove
+    ///         } else {
+    ///             false
+    ///         }
+    ///     });
+    ///
+    ///     ecs_iter!(world, |comp_a: &CompA| {
+    ///         vec_b.push(comp_a.0);
+    ///     });
+    ///
+    ///     assert_eq!(vec_a.iter().copied().sum::<u32>(), 2 + 4 + 6);
+    ///     assert_eq!(vec_b.iter().copied().sum::<u32>(), 1 + 3 + 5);
+    /// }
+    /// ```
+    #[macro_export]
+    macro_rules! ecs_iter_remove {
+        (...) => {};
+    }
 }
 
 /// A special parameter type for ECS query closures to match one of multiple components.
@@ -711,7 +771,7 @@ pub mod __internal {
 
     pub use gecs_macros::__ecs_finalize;
     pub use gecs_macros::{__ecs_find, __ecs_find_borrow};
-    pub use gecs_macros::{__ecs_iter, __ecs_iter_borrow};
+    pub use gecs_macros::{__ecs_iter, __ecs_iter_borrow, __ecs_iter_remove};
 
     pub use error::EcsError;
 

tests/test_iter_remove.rs

@@ -0,0 +1,58 @@
+use gecs::prelude::*;
+
+#[derive(Debug, PartialEq)]
+pub struct CompA(pub u32);
+#[derive(Debug, PartialEq)]
+pub struct CompB(pub u32);
+#[derive(Debug, PartialEq)]
+pub struct CompC(pub u32);
+
+ecs_world! {
+    #[archetype_id(3)]
+    ecs_archetype!(
+        ArchFoo,
+        5,
+        CompA,
+        CompB,
+    );
+
+    ecs_archetype!(
+        ArchBar,
+        5,
+        CompA,
+        CompC,
+    );
+}
+
+#[test]
+#[rustfmt::skip]
+fn test_one_of_basic() {
+    let mut world = EcsWorld::default();
+
+    world.archetype_mut::<ArchFoo>().create((CompA(1), CompB(10)));
+    world.archetype_mut::<ArchFoo>().create((CompA(2), CompB(20)));
+    world.archetype_mut::<ArchFoo>().create((CompA(3), CompB(30)));
+
+    world.archetype_mut::<ArchBar>().create((CompA(4), CompC(10)));
+    world.archetype_mut::<ArchBar>().create((CompA(5), CompC(10)));
+    world.archetype_mut::<ArchBar>().create((CompA(6), CompC(10)));
+
+    let mut vec_a = Vec::<u32>::new();
+    let mut vec_b = Vec::<u32>::new();
+
+    ecs_iter_remove!(world, |comp_a: &CompA| {
+        if comp_a.0 & 1 == 0 {
+            vec_a.push(comp_a.0);
+            true
+        } else {
+            false
+        }
+    });
+
+    ecs_iter!(world, |comp_a: &CompA| {
+        vec_b.push(comp_a.0);
+    });
+
+    assert_eq!(vec_a.iter().copied().sum::<u32>(), 2 + 4 + 6);
+    assert_eq!(vec_b.iter().copied().sum::<u32>(), 1 + 3 + 5);
+}