diff --git a/codegen/src/lib.rs b/codegen/src/lib.rs index dab6d116..73a4dd0b 100644 --- a/codegen/src/lib.rs +++ b/codegen/src/lib.rs @@ -3,8 +3,8 @@ use proc_macro::TokenStream; use proc_macro2::Span; use quote::{format_ident, quote, quote_spanned}; use syn::{ - parse_macro_input, parse_quote, Attribute, Expr, GenericArgument, Generics, Ident, Index, - ItemFn, Lit, Meta, PathArguments, Signature, Type, Visibility, + parse_macro_input, parse_quote, spanned::Spanned, Attribute, Expr, GenericArgument, Generics, + Ident, Index, ItemFn, Lit, Meta, PathArguments, Signature, Type, TypePath, Visibility, }; /// Wraps a function in a system, and generates a new function which constructs that system. @@ -73,6 +73,20 @@ use syn::{ /// } /// ``` /// +/// Systems can declare queries. The above can also be written as: +/// +/// ```ignore +/// # use legion_codegen::system; +/// # use legion::{Schedule, world::SubWorld, Read, Write, IntoQuery, Query}; +/// # struct Time; +/// #[system] +/// fn run_query(world: &mut SubWorld, query: &mut Query<(&usize, &mut bool)>) { +/// for (a, b) in query.iter_mut(world) { +/// println!("{} {}", a, b); +/// } +/// } +/// ``` +/// /// `for_each` and `par_for_each` system types can be used to implement the query for you. /// References will be interpreted as `Read` and `Write`, while options of references /// (e.g. `Option<&Position>`) will be interpreted as `TryRead` and `TryWrite`. You can @@ -119,7 +133,7 @@ use syn::{ /// /// Schedule::builder() /// // initialize state when you construct the system -/// .add_system(stateful_system(5usize)) +/// .add_system(stateful_system(5_usize)) /// .build(); /// ``` /// @@ -194,11 +208,13 @@ enum Error { ExpectedFilterExpression(Span), #[error( "system does not request any component access (sub-world will have no permissions), \ - consider using #[read_compnent(T)] or #[write_component(T)]" + consider using #[read_component(T)] or #[write_component(T)], or add a Query to the system" )] SubworldWithoutPermissions, #[error("{0}")] Message(String), + #[error("Queries should be passed to system functions via mutable references")] + QueryShouldBeMutableReference(Span), } impl Error { @@ -210,6 +226,7 @@ impl Error { Error::InvalidArgument(span) => *span, Error::ExpectedComponentType(span) => *span, Error::ExpectedFilterExpression(span) => *span, + Error::QueryShouldBeMutableReference(span) => *span, _ => Span::call_site(), } } @@ -307,48 +324,55 @@ impl Sig { syn::FnArg::Receiver(_) => return Err(Error::SelfNotAllowed), syn::FnArg::Typed(arg) => { match arg.ty.as_ref() { - Type::Path(ty_path) => { - let ident = &ty_path.path.segments[0].ident; - if ident == "Option" { - match &ty_path.path.segments[0].arguments { - PathArguments::AngleBracketed(bracketed) => { - let arg = bracketed.args.iter().next().unwrap(); - match arg { - GenericArgument::Type(ty) => { - match ty { - Type::Reference(ty) => { - let mutable = ty.mutability.is_some(); - parameters.push(Parameter::Component( - query.len(), - )); - let elem = &ty.elem; - if mutable { - query.push( - parse_quote!(::legion::TryWrite<#elem>), - ); - } else { - query.push( - parse_quote!(::legion::TryRead<#elem>), - ); - } - } - _ => { - return Err(Error::InvalidOptionArgument( - ident.span(), - quote!(#ty).to_string(), - )) + Type::Path(ty_path) if ty_path.path.segments[0].ident == "Option" => { + let segment = &ty_path.path.segments[0]; + match &segment.arguments { + PathArguments::AngleBracketed(bracketed) => { + let arg = bracketed.args.iter().next().unwrap(); + match arg { + GenericArgument::Type(ty) => { + match ty { + Type::Reference(ty) => { + let mutable = ty.mutability.is_some(); + parameters + .push(Parameter::Component(query.len())); + let elem = &ty.elem; + if mutable { + query.push( + parse_quote!(::legion::TryWrite<#elem>), + ); + } else { + query.push( + parse_quote!(::legion::TryRead<#elem>), + ); } } + _ => { + return Err(Error::InvalidOptionArgument( + segment.ident.span(), + quote!(#ty).to_string(), + )) + } } - _ => panic!(), } + _ => panic!(), } - _ => panic!(), } - } else { - return Err(Error::InvalidArgument(ident.span())); + _ => panic!(), } } + Type::Path(ty_path) + if path_match(ty_path, &["Query"]) + || path_match(ty_path, &["legion", "Query"]) + || path_match(ty_path, &["legion", "query", "Query"]) => + { + return Err(Error::QueryShouldBeMutableReference(ty_path.span())); + } + Type::Path(ty_path) => { + return Err(Error::InvalidArgument( + ty_path.path.segments[0].ident.span(), + )); + } Type::Reference(ty) if is_type(&ty.elem, &["CommandBuffer"]) || is_type(&ty.elem, &["legion", "CommandBuffer"]) @@ -379,6 +403,17 @@ impl Sig { parameters.push(Parameter::Component(query.len())); query.push(parse_quote!(::legion::Entity)); } + Type::Reference(ty) + if is_type(&ty.elem, &["Query"]) + || is_type(&ty.elem, &["legion", "Query"]) + || is_type(&ty.elem, &["legion", "query", "Query"]) => + { + if ty.mutability.is_none() { + return Err(Error::QueryShouldBeMutableReference(ty.span())); + } + + parameters.push(Parameter::Query(ty.elem.clone())); + } Type::Reference(ty) => { let mutable = ty.mutability.is_some(); let attribute = Self::find_remove_arg_attr(&mut arg.attrs); @@ -454,15 +489,19 @@ enum ArgAttr { fn is_type(ty: &Type, segments: &[&str]) -> bool { if let Type::Path(path) = ty { - segments - .iter() - .zip(path.path.segments.iter()) - .all(|(a, b)| b.ident == *a) + path_match(path, segments) } else { false } } +fn path_match(path: &TypePath, segments: &[&str]) -> bool { + segments + .iter() + .zip(path.path.segments.iter()) + .all(|(a, b)| b.ident == *a) +} + #[derive(Copy, Clone, PartialEq)] enum SystemType { Simple, @@ -480,6 +519,12 @@ impl SystemType { } } +impl Default for SystemType { + fn default() -> Self { + SystemType::Simple + } +} + enum Parameter { CommandBuffer, CommandBufferMut, @@ -490,6 +535,7 @@ enum Parameter { ResourceMut(usize), State(usize), StateMut(usize), + Query(Box), } struct Config { @@ -586,7 +632,12 @@ impl Config { .any(|p| matches!(p, Parameter::SubWorldMut)); let has_components = !self.read_components.is_empty() || !self.write_components.is_empty(); - if (has_subworld || has_subworld_mut) && !has_components { + let has_queries = self + .signature + .parameters + .iter() + .any(|param| matches!(param, Parameter::Query(_))); + if (has_subworld || has_subworld_mut) && (!has_components && !has_queries) { return Err(Error::SubworldWithoutPermissions); } } @@ -640,7 +691,7 @@ impl Config { signature, } = self; - let system_type = attr.system_type.unwrap_or(SystemType::Simple); + let system_type = attr.system_type.unwrap_or_default(); // declare query let query = if system_type.requires_query() { @@ -659,9 +710,16 @@ impl Config { let has_query = !signature.query.is_empty(); let single_resource = (signature.read_resources.len() + signature.write_resources.len()) == 1; + let single_query = signature + .parameters + .iter() + .filter(|param| matches!(param, Parameter::Query(_))) + .count() + == 1; let mut call_params = Vec::new(); let mut fn_params = Vec::new(); let mut world = None; + let mut queries = Vec::new(); for param in &signature.parameters { match param { Parameter::CommandBuffer => call_params.push(quote!(cmd)), @@ -719,6 +777,15 @@ impl Config { call_params.push(quote!(&mut #arg_name)); fn_params.push(quote!(mut #arg_name: #arg_type)); } + Parameter::Query(ty) if single_query => { + call_params.push(quote!(query)); + queries.push(quote!(#ty)); + } + Parameter::Query(ty) => { + let idx = Index::from(queries.len()); + call_params.push(quote!(&mut query.#idx)); + queries.push(quote!(#ty)); + } } } @@ -777,6 +844,7 @@ impl Config { #(.read_resource::<#read_resources>())* #(.write_resource::<#write_resources>())* #query + #(.with_query(<#queries>::new()))* .build(move |cmd, world, resources, query| { #body }) diff --git a/src/lib.rs b/src/lib.rs index 48ac7126..b30cd9f7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -205,8 +205,8 @@ pub use legion_codegen::system; pub use crate::serialize::Registry; pub use crate::{ query::{ - any, component, maybe_changed, passthrough, Fetch, IntoQuery, Read, TryRead, TryWrite, - Write, + any, component, maybe_changed, passthrough, Fetch, IntoQuery, Query, Read, TryRead, + TryWrite, Write, }, storage::{GroupSource, IntoSoa}, systems::{Resources, Schedule, SystemBuilder}, diff --git a/tests/systems_basic.rs b/tests/systems_basic.rs index fd19755b..51e76c5c 100644 --- a/tests/systems_basic.rs +++ b/tests/systems_basic.rs @@ -1,10 +1,16 @@ #[cfg(feature = "codegen")] mod tests { - use std::fmt::Debug; + use std::{ + fmt::Debug, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, + }; use legion::{ - storage::Component, system, systems::CommandBuffer, world::SubWorld, IntoQuery, Read, - Schedule, Write, + storage::Component, system, systems::CommandBuffer, world::SubWorld, IntoQuery, Query, + Read, Resources, Schedule, World, Write, }; #[test] @@ -158,4 +164,43 @@ mod tests { Schedule::builder().add_system(basic_system(false)).build(); } + + #[test] + fn with_query() { + #[system] + fn basic(_: &mut SubWorld, _: &mut Query<(&u8, &mut i8)>) {} + + Schedule::builder().add_system(basic_system()).build(); + } + + #[test] + fn with_two_queries() { + #[system] + fn basic( + #[state] a_count: &Arc, + #[state] b_count: &Arc, + world: &mut SubWorld, + a: &mut Query<(&u8, &mut usize)>, + b: &mut Query<(&usize, &mut bool)>, + ) { + a_count.store(a.iter_mut(world).count(), Ordering::SeqCst); + b_count.store(b.iter_mut(world).count(), Ordering::SeqCst); + } + + let a_count = Arc::new(AtomicUsize::new(0)); + let b_count = Arc::new(AtomicUsize::new(0)); + let mut schedule = Schedule::builder() + .add_system(basic_system(a_count.clone(), b_count.clone())) + .build(); + + let mut world = World::default(); + let mut resources = Resources::default(); + + world.extend(vec![(1_usize, false), (2_usize, false)]); + + schedule.execute(&mut world, &mut resources); + + assert_eq!(0, a_count.load(Ordering::SeqCst)); + assert_eq!(2, b_count.load(Ordering::SeqCst)); + } }