bones_ecs_macros_core/
lib.rs

1use proc_macro2::{Span, TokenStream};
2use quote::quote;
3use syn::{parse2, punctuated::Punctuated, Fields, GenericParam, Index, ItemStruct, Token};
4
5macro_rules! err {
6    ($target:expr, $message:expr) => {
7        return Err(::syn::Error::new(
8            ::syn::spanned::Spanned::span(&$target),
9            $message,
10        ))
11    };
12}
13
14pub fn generate_system_param_impl(input: TokenStream) -> TokenStream {
15    match _generate_system_param_impl(input) {
16        Ok(output) => output,
17        Err(err) => err.to_compile_error(),
18    }
19}
20
21fn _generate_system_param_impl(input: TokenStream) -> syn::Result<TokenStream> {
22    let item_struct: ItemStruct = parse2(input)?;
23
24    let Some(GenericParam::Lifetime(lifetime)) =
25        get_single_punctuated(&item_struct.generics.params)
26    else {
27        err!(
28            item_struct,
29            "struct must have a single generic lifetime parameter"
30        );
31    };
32
33    let ident = &item_struct.ident;
34
35    let fields = match &item_struct.fields {
36        Fields::Unit => err!(item_struct, "unit structs are not supported"),
37        Fields::Unnamed(_) => err!(item_struct, "structs with unnamed fields are not supported"),
38        Fields::Named(fields) => fields,
39    };
40
41    let state_types: Punctuated<TokenStream, Token![,]> =
42        Punctuated::from_iter(fields.named.iter().map(|field| {
43            let ty = &field.ty;
44            quote! { <#ty as ::bones_ecs::prelude::SystemParam>::State }
45        }));
46
47    let get_state_items: Punctuated<TokenStream, Token![,]> =
48        Punctuated::from_iter(fields.named.iter().map(|field| {
49            let ty = &field.ty;
50            quote! { <#ty as ::bones_ecs::prelude::SystemParam>::get_state(world) }
51        }));
52
53    let borrow_param_fields: Punctuated<TokenStream, Token![,]> = fields
54        .named
55        .iter()
56        .enumerate()
57        .map(|(index, field)| {
58            let ident = field.ident.as_ref().unwrap();
59            let ty = &field.ty;
60            let index = Index {
61                index: index as u32,
62                span: Span::call_site(),
63            };
64            quote! {
65                #ident: <#ty as ::bones_ecs::prelude::SystemParam>::borrow(world, &mut state.#index)
66            }
67        })
68        .collect();
69
70    Ok(quote! {
71        impl<#lifetime> ::bones_ecs::prelude::SystemParam for #ident<#lifetime> {
72            type State = ( #state_types );
73            type Param<'p> = #ident<'p>;
74            fn get_state(world: &::bones_ecs::prelude::World) -> Self::State {
75                ( #get_state_items )
76            }
77            fn borrow<'s>(
78                world: &'s ::bones_ecs::prelude::World,
79                state: &'s mut Self::State,
80            ) -> Self::Param<'s> {
81                Self::Param { #borrow_param_fields }
82            }
83        }
84    })
85}
86
87fn get_single_punctuated<T, P>(punctuated: &Punctuated<T, P>) -> Option<&T> {
88    match punctuated.first() {
89        single @ Some(_) if punctuated.len() == 1 => single,
90        _ => None,
91    }
92}
93
94#[cfg(test)]
95mod tests {
96    use pretty_assertions::assert_eq;
97    use quote::quote;
98
99    use super::*;
100
101    fn assert_tokens_eq(expected: TokenStream, actual: TokenStream) {
102        let expected = expected.to_string();
103        let actual = actual.to_string();
104        assert_eq!(expected, actual);
105    }
106
107    #[test]
108    fn correct_system_param_impl() {
109        let expected = quote! {
110            impl<'a> ::bones_ecs::prelude::SystemParam for MySystemParam<'a> {
111                type State = (
112                    <Commands<'a> as ::bones_ecs::prelude::SystemParam>::State,
113                    <ResMut<'a, Entities> as ::bones_ecs::prelude::SystemParam>::State
114                );
115                type Param<'p> = MySystemParam<'p>;
116                fn get_state(world: &::bones_ecs::prelude::World) -> Self::State {
117                    (
118                        <Commands<'a> as ::bones_ecs::prelude::SystemParam>::get_state(world),
119                        <ResMut<'a, Entities> as ::bones_ecs::prelude::SystemParam>::get_state(world)
120                    )
121                }
122                fn borrow<'s>(
123                    world: &'s ::bones_ecs::prelude::World,
124                    state: &'s mut Self::State,
125                ) -> Self::Param<'s> {
126                    Self::Param {
127                        commands: <Commands<'a> as ::bones_ecs::prelude::SystemParam>::borrow(world, &mut state.0),
128                        entities: <ResMut<'a, Entities> as ::bones_ecs::prelude::SystemParam>::borrow(world, &mut state.1)
129                    }
130                }
131            }
132        };
133        let input = quote! {
134            struct MySystemParam<'a> {
135                commands: Commands<'a>,
136                entities: ResMut<'a, Entities>,
137            }
138        };
139        let actual = generate_system_param_impl(input);
140        assert_tokens_eq(expected, actual);
141    }
142}