diff --git a/macros/src/generate/query.rs b/macros/src/generate/query.rs index e637c1b..e4d97e7 100644 --- a/macros/src/generate/query.rs +++ b/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 { +pub fn generate_query_find( + mode: FetchMode, //. + query: ParseQueryFind, +) -> syn::Result { 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 { +pub fn generate_query_iter( + mode: FetchMode, //. + query: ParseQueryIter, +) -> syn::Result { 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 { + 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::>(); + + // Special cases + let maybe_mut = query.params.iter().map(to_maybe_mut).collect::>(); + + let mut queries = Vec::::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::>(); // 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::>(), + }; + + 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 ¶m.param_type { diff --git a/macros/src/generate/world.rs b/macros/src/generate/world.rs index fad5201..73cfa78 100644 --- a/macros/src/generate/world.rs +++ b/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; ) } diff --git a/macros/src/lib.rs b/macros/src/lib.rs index 790eed8..a723c00 100644 --- a/macros/src/lib.rs +++ b/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(), + } +} diff --git a/macros/src/parse/query.rs b/macros/src/parse/query.rs index 9aaec00..62850a0 100644 --- a/macros/src/parse/query.rs +++ b/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, + 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 { + // Parse out the hidden serialized world data + let world_data = input.parse::()?; + input.parse::()?; + + // Parse out the meta-arguments for the query + let world = input.parse()?; + input.parse::()?; + + // Parse out the closure arguments + input.parse::()?; + let params = parse_params(&input)?; + input.parse::()?; + + // Parse the rest of the body, including the braces (if any) + let body = input.parse::()?; + + Ok(Self { + world_data: world_data.value(), + world, + params, + body, + }) + } +} + impl Parse for ParseQueryParam { fn parse(input: ParseStream) -> syn::Result { // Parse the name and following : token diff --git a/src/lib.rs b/src/lib.rs index e2be4fa..dc971c0 100644 --- a/src/lib.rs +++ b/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::().create((CompA(1), CompB(10))); + /// world.archetype_mut::().create((CompA(2), CompB(20))); + /// world.archetype_mut::().create((CompA(3), CompB(30))); + /// + /// world.archetype_mut::().create((CompA(4), CompC(10))); + /// world.archetype_mut::().create((CompA(5), CompC(10))); + /// world.archetype_mut::().create((CompA(6), CompC(10))); + /// + /// let mut vec_a = Vec::::new(); + /// let mut vec_b = Vec::::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::(), 2 + 4 + 6); + /// assert_eq!(vec_b.iter().copied().sum::(), 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; diff --git a/tests/test_iter_remove.rs b/tests/test_iter_remove.rs new file mode 100644 index 0000000..6ca69ee --- /dev/null +++ b/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::().create((CompA(1), CompB(10))); + world.archetype_mut::().create((CompA(2), CompB(20))); + world.archetype_mut::().create((CompA(3), CompB(30))); + + world.archetype_mut::().create((CompA(4), CompC(10))); + world.archetype_mut::().create((CompA(5), CompC(10))); + world.archetype_mut::().create((CompA(6), CompC(10))); + + let mut vec_a = Vec::::new(); + let mut vec_b = Vec::::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::(), 2 + 4 + 6); + assert_eq!(vec_b.iter().copied().sum::(), 1 + 3 + 5); +}