Skip to content

Commit

Permalink
Merge #237
Browse files Browse the repository at this point in the history 
237: Allow Queries to be requested as parameters in #[system] r=TomGillen a=TomGillen

This PR allows users to request queries as parameters in basic `#[system]` macro invocations. This allows a system to request multiple queries without the user having to construct them themselves. It also adds these queries to the system (via `SystemBuilder`), which means (in most cases) users won't need to use `#[read_component(T)]`/`#[write_component(T)]` attributes, and the queries persist between each run of the system, allowing for some level of caching between runs.

For example, prior to this PR:

```rust
#[system]
#[read_component(A, B)]
#[write_component(C)]
fn my_system(world: &mut SubWorld) {
    let query_a = <(&A, &mut C)>::query();
    for (a, c) in query_a.iter_mut(world) {
        ...
    }

    let query_b = <(&B, &mut C)>::query();
    for (b, c) in query_b.iter_mut(world) {
        ...
    }
}
```

Can now be written as:

```rust
#[system]
fn my_system(
    world: &mut SubWorld,
    query_a: &mut Query<(&A, &mut C)>,
    query_b: &mut Query<(&B, &mut C)>,
) {
    for (a, c) in query_a.iter_mut(world) {
        ...
    }

    for (b, c) in query_b.iter_mut(world) {
        ...
    }
}
```

Co-authored-by: Tom Gillen <[email protected]>
  • Loading branch information
@amethyst-bors @TomGillen
amethyst-bors and TomGillen committed Feb 25, 2021
2 parents 5712f4d + e7676ce commit 7f8a680
Show file tree
Hide file tree
Showing 3 changed files with 162 additions and 49 deletions.
156 changes: 112 additions & 44 deletions codegen/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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<T>` and `Write<T>`, while options of references
/// (e.g. `Option<&Position>`) will be interpreted as `TryRead<T>` and `TryWrite<T>`. You can
Expand Down Expand Up @@ -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();
/// ```
///
Expand Down Expand Up @@ -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 {
Expand All @@ -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(),
}
}
Expand Down Expand Up @@ -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"])
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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,
Expand All @@ -480,6 +519,12 @@ impl SystemType {
}
}

impl Default for SystemType {
fn default() -> Self {
SystemType::Simple
}
}

enum Parameter {
CommandBuffer,
CommandBufferMut,
Expand All @@ -490,6 +535,7 @@ enum Parameter {
ResourceMut(usize),
State(usize),
StateMut(usize),
Query(Box<Type>),
}

struct Config {
Expand Down Expand Up @@ -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);
}
}
Expand Down Expand Up @@ -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() {
Expand All @@ -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)),
Expand Down Expand Up @@ -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));
}
}
}

Expand Down Expand Up @@ -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
})
Expand Down
4 changes: 2 additions & 2 deletions src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down
51 changes: 48 additions & 3 deletions tests/systems_basic.rs
View file
Original file line number Diff line number Diff line change
@@ -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]
Expand Down Expand Up @@ -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<AtomicUsize>,
#[state] b_count: &Arc<AtomicUsize>,
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));
}
}

0 comments on commit 7f8a680

Please sign in to comment.