azalea_ecs_macros/
lib.rs

1//! A fork of bevy_ecs_macros that uses azalea_ecs instead of bevy_ecs.
2
3extern crate proc_macro;
4
5mod component;
6mod fetch;
7pub(crate) mod utils;
8
9use crate::fetch::derive_world_query_impl;
10use proc_macro::TokenStream;
11use proc_macro2::Span;
12use quote::{format_ident, quote};
13use syn::{
14    parse::{Parse, ParseStream},
15    parse_macro_input,
16    punctuated::Punctuated,
17    spanned::Spanned,
18    token::Comma,
19    DeriveInput, Field, GenericParam, Ident, Index, LitInt, Meta, MetaList, NestedMeta, Result,
20    Token, TypeParam,
21};
22use utils::{derive_label, get_named_struct_fields, BevyManifest};
23
24struct AllTuples {
25    macro_ident: Ident,
26    start: usize,
27    end: usize,
28    idents: Vec<Ident>,
29}
30
31impl Parse for AllTuples {
32    fn parse(input: ParseStream) -> Result<Self> {
33        let macro_ident = input.parse::<Ident>()?;
34        input.parse::<Comma>()?;
35        let start = input.parse::<LitInt>()?.base10_parse()?;
36        input.parse::<Comma>()?;
37        let end = input.parse::<LitInt>()?.base10_parse()?;
38        input.parse::<Comma>()?;
39        let mut idents = vec![input.parse::<Ident>()?];
40        while input.parse::<Comma>().is_ok() {
41            idents.push(input.parse::<Ident>()?);
42        }
43
44        Ok(AllTuples {
45            macro_ident,
46            start,
47            end,
48            idents,
49        })
50    }
51}
52
53#[proc_macro]
54pub fn all_tuples(input: TokenStream) -> TokenStream {
55    let input = parse_macro_input!(input as AllTuples);
56    let len = input.end - input.start;
57    let mut ident_tuples = Vec::with_capacity(len);
58    for i in input.start..=input.end {
59        let idents = input
60            .idents
61            .iter()
62            .map(|ident| format_ident!("{}{}", ident, i));
63        if input.idents.len() < 2 {
64            ident_tuples.push(quote! {
65                #(#idents)*
66            });
67        } else {
68            ident_tuples.push(quote! {
69                (#(#idents),*)
70            });
71        }
72    }
73
74    let macro_ident = &input.macro_ident;
75    let invocations = (input.start..=input.end).map(|i| {
76        let ident_tuples = &ident_tuples[..i];
77        quote! {
78            #macro_ident!(#(#ident_tuples),*);
79        }
80    });
81    TokenStream::from(quote! {
82        #(
83            #invocations
84        )*
85    })
86}
87
88enum BundleFieldKind {
89    Component,
90    Ignore,
91}
92
93const BUNDLE_ATTRIBUTE_NAME: &str = "bundle";
94const BUNDLE_ATTRIBUTE_IGNORE_NAME: &str = "ignore";
95
96#[proc_macro_derive(Bundle, attributes(bundle))]
97pub fn derive_bundle(input: TokenStream) -> TokenStream {
98    let ast = parse_macro_input!(input as DeriveInput);
99    let ecs_path = azalea_ecs_path();
100
101    let named_fields = match get_named_struct_fields(&ast.data) {
102        Ok(fields) => &fields.named,
103        Err(e) => return e.into_compile_error().into(),
104    };
105
106    let mut field_kind = Vec::with_capacity(named_fields.len());
107
108    'field_loop: for field in named_fields.iter() {
109        for attr in &field.attrs {
110            if attr.path.is_ident(BUNDLE_ATTRIBUTE_NAME) {
111                if let Ok(Meta::List(MetaList { nested, .. })) = attr.parse_meta() {
112                    if let Some(&NestedMeta::Meta(Meta::Path(ref path))) = nested.first() {
113                        if path.is_ident(BUNDLE_ATTRIBUTE_IGNORE_NAME) {
114                            field_kind.push(BundleFieldKind::Ignore);
115                            continue 'field_loop;
116                        }
117
118                        return syn::Error::new(
119                            path.span(),
120                            format!(
121                                "Invalid bundle attribute. Use `{BUNDLE_ATTRIBUTE_IGNORE_NAME}`"
122                            ),
123                        )
124                        .into_compile_error()
125                        .into();
126                    }
127
128                    return syn::Error::new(attr.span(), format!("Invalid bundle attribute. Use `#[{BUNDLE_ATTRIBUTE_NAME}({BUNDLE_ATTRIBUTE_IGNORE_NAME})]`")).into_compile_error().into();
129                }
130            }
131        }
132
133        field_kind.push(BundleFieldKind::Component);
134    }
135
136    let field = named_fields
137        .iter()
138        .map(|field| field.ident.as_ref().unwrap())
139        .collect::<Vec<_>>();
140    let field_type = named_fields
141        .iter()
142        .map(|field| &field.ty)
143        .collect::<Vec<_>>();
144
145    let mut field_component_ids = Vec::new();
146    let mut field_get_components = Vec::new();
147    let mut field_from_components = Vec::new();
148    for ((field_type, field_kind), field) in
149        field_type.iter().zip(field_kind.iter()).zip(field.iter())
150    {
151        match field_kind {
152            BundleFieldKind::Component => {
153                field_component_ids.push(quote! {
154                <#field_type as #ecs_path::bundle::BevyBundle>::component_ids(components, storages, &mut *ids);
155                });
156                field_get_components.push(quote! {
157                    self.#field.get_components(&mut *func);
158                });
159                field_from_components.push(quote! {
160                    #field: <#field_type as #ecs_path::bundle::BevyBundle>::from_components(ctx, &mut *func),
161                });
162            }
163
164            BundleFieldKind::Ignore => {
165                field_from_components.push(quote! {
166                    #field: ::std::default::Default::default(),
167                });
168            }
169        }
170    }
171    let generics = ast.generics;
172    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
173    let struct_name = &ast.ident;
174
175    TokenStream::from(quote! {
176        /// SAFETY: ComponentId is returned in field-definition-order. [from_components] and [get_components] use field-definition-order
177        unsafe impl #impl_generics #ecs_path::bundle::BevyBundle for #struct_name #ty_generics #where_clause {
178            fn component_ids(
179                components: &mut #ecs_path::component::Components,
180                storages: &mut #ecs_path::storage::Storages,
181                ids: &mut impl FnMut(#ecs_path::component::ComponentId)
182            ){
183                #(#field_component_ids)*
184            }
185
186            #[allow(unused_variables, non_snake_case)]
187            unsafe fn from_components<__T, __F>(ctx: &mut __T, func: &mut __F) -> Self
188            where
189                __F: FnMut(&mut __T) -> #ecs_path::ptr::OwningPtr<'_>
190            {
191                Self {
192                    #(#field_from_components)*
193                }
194            }
195
196            #[allow(unused_variables)]
197            fn get_components(self, func: &mut impl FnMut(#ecs_path::ptr::OwningPtr<'_>)) {
198                #(#field_get_components)*
199            }
200        }
201    })
202}
203
204fn get_idents(fmt_string: fn(usize) -> String, count: usize) -> Vec<Ident> {
205    (0..count)
206        .map(|i| Ident::new(&fmt_string(i), Span::call_site()))
207        .collect::<Vec<Ident>>()
208}
209
210#[proc_macro]
211pub fn impl_param_set(_input: TokenStream) -> TokenStream {
212    let mut tokens = TokenStream::new();
213    let max_params = 8;
214    let params = get_idents(|i| format!("P{i}"), max_params);
215    let params_fetch = get_idents(|i| format!("PF{i}"), max_params);
216    let metas = get_idents(|i| format!("m{i}"), max_params);
217    let mut param_fn_muts = Vec::new();
218    for (i, param) in params.iter().enumerate() {
219        let fn_name = Ident::new(&format!("p{i}"), Span::call_site());
220        let index = Index::from(i);
221        param_fn_muts.push(quote! {
222            pub fn #fn_name<'a>(&'a mut self) -> <#param::Fetch as SystemParamFetch<'a, 'a>>::Item {
223                // SAFETY: systems run without conflicts with other systems.
224                // Conflicting params in ParamSet are not accessible at the same time
225                // ParamSets are guaranteed to not conflict with other SystemParams
226                unsafe {
227                    <#param::Fetch as SystemParamFetch<'a, 'a>>::get_param(&mut self.param_states.#index, &self.system_meta, self.world, self.change_tick)
228                }
229            }
230        });
231    }
232
233    for param_count in 1..=max_params {
234        let param = &params[0..param_count];
235        let param_fetch = &params_fetch[0..param_count];
236        let meta = &metas[0..param_count];
237        let param_fn_mut = &param_fn_muts[0..param_count];
238        tokens.extend(TokenStream::from(quote! {
239            impl<'w, 's, #(#param: SystemParam,)*> SystemParam for ParamSet<'w, 's, (#(#param,)*)>
240            {
241                type Fetch = ParamSetState<(#(#param::Fetch,)*)>;
242            }
243
244            // SAFETY: All parameters are constrained to ReadOnlyFetch, so World is only read
245
246            unsafe impl<#(#param_fetch: for<'w1, 's1> SystemParamFetch<'w1, 's1>,)*> ReadOnlySystemParamFetch for ParamSetState<(#(#param_fetch,)*)>
247            where #(#param_fetch: ReadOnlySystemParamFetch,)*
248            { }
249
250            // SAFETY: Relevant parameter ComponentId and ArchetypeComponentId access is applied to SystemMeta. If any ParamState conflicts
251            // with any prior access, a panic will occur.
252
253            unsafe impl<#(#param_fetch: for<'w1, 's1> SystemParamFetch<'w1, 's1>,)*> SystemParamState for ParamSetState<(#(#param_fetch,)*)>
254            {
255                fn init(world: &mut World, system_meta: &mut SystemMeta) -> Self {
256                    #(
257                        // Pretend to add each param to the system alone, see if it conflicts
258                        let mut #meta = system_meta.clone();
259                        #meta.component_access_set.clear();
260                        #meta.archetype_component_access.clear();
261                        #param_fetch::init(world, &mut #meta);
262                        let #param = #param_fetch::init(world, &mut system_meta.clone());
263                    )*
264                    #(
265                        system_meta
266                            .component_access_set
267                            .extend(#meta.component_access_set);
268                        system_meta
269                            .archetype_component_access
270                            .extend(&#meta.archetype_component_access);
271                    )*
272                    ParamSetState((#(#param,)*))
273                }
274
275                fn new_archetype(&mut self, archetype: &Archetype, system_meta: &mut SystemMeta) {
276                    let (#(#param,)*) = &mut self.0;
277                    #(
278                        #param.new_archetype(archetype, system_meta);
279                    )*
280                }
281
282                fn apply(&mut self, world: &mut World) {
283                    self.0.apply(world)
284                }
285            }
286
287
288
289            impl<'w, 's, #(#param_fetch: for<'w1, 's1> SystemParamFetch<'w1, 's1>,)*> SystemParamFetch<'w, 's> for ParamSetState<(#(#param_fetch,)*)>
290            {
291                type Item = ParamSet<'w, 's, (#(<#param_fetch as SystemParamFetch<'w, 's>>::Item,)*)>;
292
293                #[inline]
294                unsafe fn get_param(
295                    state: &'s mut Self,
296                    system_meta: &SystemMeta,
297                    world: &'w World,
298                    change_tick: u32,
299                ) -> Self::Item {
300                    ParamSet {
301                        param_states: &mut state.0,
302                        system_meta: system_meta.clone(),
303                        world,
304                        change_tick,
305                    }
306                }
307            }
308
309            impl<'w, 's, #(#param: SystemParam,)*> ParamSet<'w, 's, (#(#param,)*)>
310            {
311
312                #(#param_fn_mut)*
313            }
314        }));
315    }
316
317    tokens
318}
319
320#[derive(Default)]
321struct SystemParamFieldAttributes {
322    pub ignore: bool,
323}
324
325static SYSTEM_PARAM_ATTRIBUTE_NAME: &str = "system_param";
326
327/// Implement `SystemParam` to use a struct as a parameter in a system
328#[proc_macro_derive(SystemParam, attributes(system_param))]
329pub fn derive_system_param(input: TokenStream) -> TokenStream {
330    let ast = parse_macro_input!(input as DeriveInput);
331    let fields = match get_named_struct_fields(&ast.data) {
332        Ok(fields) => &fields.named,
333        Err(e) => return e.into_compile_error().into(),
334    };
335    let path = azalea_ecs_path();
336
337    let field_attributes = fields
338        .iter()
339        .map(|field| {
340            (
341                field,
342                field
343                    .attrs
344                    .iter()
345                    .find(|a| *a.path.get_ident().as_ref().unwrap() == SYSTEM_PARAM_ATTRIBUTE_NAME)
346                    .map_or_else(SystemParamFieldAttributes::default, |a| {
347                        syn::custom_keyword!(ignore);
348                        let mut attributes = SystemParamFieldAttributes::default();
349                        a.parse_args_with(|input: ParseStream| {
350                            if input.parse::<Option<ignore>>()?.is_some() {
351                                attributes.ignore = true;
352                            }
353                            Ok(())
354                        })
355                        .expect("Invalid 'system_param' attribute format.");
356
357                        attributes
358                    }),
359            )
360        })
361        .collect::<Vec<(&Field, SystemParamFieldAttributes)>>();
362    let mut fields = Vec::new();
363    let mut field_indices = Vec::new();
364    let mut field_types = Vec::new();
365    let mut ignored_fields = Vec::new();
366    let mut ignored_field_types = Vec::new();
367    for (i, (field, attrs)) in field_attributes.iter().enumerate() {
368        if attrs.ignore {
369            ignored_fields.push(field.ident.as_ref().unwrap());
370            ignored_field_types.push(&field.ty);
371        } else {
372            fields.push(field.ident.as_ref().unwrap());
373            field_types.push(&field.ty);
374            field_indices.push(Index::from(i));
375        }
376    }
377
378    let generics = ast.generics;
379    let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
380
381    let lifetimeless_generics: Vec<_> = generics
382        .params
383        .iter()
384        .filter(|g| matches!(g, GenericParam::Type(_)))
385        .collect();
386
387    let mut punctuated_generics = Punctuated::<_, Token![,]>::new();
388    punctuated_generics.extend(lifetimeless_generics.iter().map(|g| match g {
389        GenericParam::Type(g) => GenericParam::Type(TypeParam {
390            default: None,
391            ..g.clone()
392        }),
393        _ => unreachable!(),
394    }));
395
396    let mut punctuated_generic_idents = Punctuated::<_, Token![,]>::new();
397    punctuated_generic_idents.extend(lifetimeless_generics.iter().map(|g| match g {
398        GenericParam::Type(g) => &g.ident,
399        _ => unreachable!(),
400    }));
401
402    let struct_name = &ast.ident;
403    let fetch_struct_visibility = &ast.vis;
404
405    TokenStream::from(quote! {
406        // We define the FetchState struct in an anonymous scope to avoid polluting the user namespace.
407        // The struct can still be accessed via SystemParam::Fetch, e.g. EventReaderState can be accessed via
408        // <EventReader<'static, 'static, T> as SystemParam>::Fetch
409        const _: () = {
410            impl #impl_generics #path::system::SystemParam for #struct_name #ty_generics #where_clause {
411                type Fetch = FetchState <(#(<#field_types as #path::system::SystemParam>::Fetch,)*), #punctuated_generic_idents>;
412            }
413
414            #[doc(hidden)]
415            #fetch_struct_visibility struct FetchState <TSystemParamState, #punctuated_generic_idents> {
416                state: TSystemParamState,
417                marker: std::marker::PhantomData<fn()->(#punctuated_generic_idents)>
418            }
419
420            unsafe impl<TSystemParamState: #path::system::SystemParamState, #punctuated_generics> #path::system::SystemParamState for FetchState <TSystemParamState, #punctuated_generic_idents> #where_clause {
421                fn init(world: &mut #path::world::World, system_meta: &mut #path::system::SystemMeta) -> Self {
422                    Self {
423                        state: TSystemParamState::init(world, system_meta),
424                        marker: std::marker::PhantomData,
425                    }
426                }
427
428                fn new_archetype(&mut self, archetype: &#path::archetype::Archetype, system_meta: &mut #path::system::SystemMeta) {
429                    self.state.new_archetype(archetype, system_meta)
430                }
431
432                fn apply(&mut self, world: &mut #path::world::World) {
433                    self.state.apply(world)
434                }
435            }
436
437            impl #impl_generics #path::system::SystemParamFetch<'w, 's> for FetchState <(#(<#field_types as #path::system::SystemParam>::Fetch,)*), #punctuated_generic_idents> #where_clause {
438                type Item = #struct_name #ty_generics;
439                unsafe fn get_param(
440                    state: &'s mut Self,
441                    system_meta: &#path::system::SystemMeta,
442                    world: &'w #path::world::World,
443                    change_tick: u32,
444                ) -> Self::Item {
445                    #struct_name {
446                        #(#fields: <<#field_types as #path::system::SystemParam>::Fetch as #path::system::SystemParamFetch>::get_param(&mut state.state.#field_indices, system_meta, world, change_tick),)*
447                        #(#ignored_fields: <#ignored_field_types>::default(),)*
448                    }
449                }
450            }
451
452            // Safety: The `ParamState` is `ReadOnlySystemParamFetch`, so this can only read from the `World`
453            unsafe impl<TSystemParamState: #path::system::SystemParamState + #path::system::ReadOnlySystemParamFetch, #punctuated_generics> #path::system::ReadOnlySystemParamFetch for FetchState <TSystemParamState, #punctuated_generic_idents> #where_clause {}
454        };
455    })
456}
457
458/// Implement `WorldQuery` to use a struct as a parameter in a query
459#[proc_macro_derive(WorldQuery, attributes(world_query))]
460pub fn derive_world_query(input: TokenStream) -> TokenStream {
461    let ast = parse_macro_input!(input as DeriveInput);
462    derive_world_query_impl(ast)
463}
464
465/// Generates an impl of the `SystemLabel` trait.
466///
467/// This works only for unit structs, or enums with only unit variants.
468/// You may force a struct or variant to behave as if it were fieldless with
469/// `#[system_label(ignore_fields)]`.
470#[proc_macro_derive(SystemLabel, attributes(system_label))]
471pub fn derive_system_label(input: TokenStream) -> TokenStream {
472    let input = parse_macro_input!(input as DeriveInput);
473    let mut trait_path = azalea_ecs_path();
474    trait_path.segments.push(format_ident!("schedule").into());
475    trait_path
476        .segments
477        .push(format_ident!("SystemLabel").into());
478    derive_label(input, &trait_path, "system_label")
479}
480
481/// Generates an impl of the `StageLabel` trait.
482///
483/// This works only for unit structs, or enums with only unit variants.
484/// You may force a struct or variant to behave as if it were fieldless with
485/// `#[stage_label(ignore_fields)]`.
486#[proc_macro_derive(StageLabel, attributes(stage_label))]
487pub fn derive_stage_label(input: TokenStream) -> TokenStream {
488    let input = parse_macro_input!(input as DeriveInput);
489    let mut trait_path = azalea_ecs_path();
490    trait_path.segments.push(format_ident!("schedule").into());
491    trait_path.segments.push(format_ident!("StageLabel").into());
492    derive_label(input, &trait_path, "stage_label")
493}
494
495/// Generates an impl of the `RunCriteriaLabel` trait.
496///
497/// This works only for unit structs, or enums with only unit variants.
498/// You may force a struct or variant to behave as if it were fieldless with
499/// `#[run_criteria_label(ignore_fields)]`.
500#[proc_macro_derive(RunCriteriaLabel, attributes(run_criteria_label))]
501pub fn derive_run_criteria_label(input: TokenStream) -> TokenStream {
502    let input = parse_macro_input!(input as DeriveInput);
503    let mut trait_path = azalea_ecs_path();
504    trait_path.segments.push(format_ident!("schedule").into());
505    trait_path
506        .segments
507        .push(format_ident!("RunCriteriaLabel").into());
508    derive_label(input, &trait_path, "run_criteria_label")
509}
510
511pub(crate) fn azalea_ecs_path() -> syn::Path {
512    BevyManifest::default().get_path("azalea_ecs")
513}
514
515#[proc_macro_derive(Resource)]
516pub fn derive_resource(input: TokenStream) -> TokenStream {
517    component::derive_resource(input)
518}
519
520#[proc_macro_derive(Component, attributes(component))]
521pub fn derive_component(input: TokenStream) -> TokenStream {
522    component::derive_component(input)
523}
azalea_ecs_macros/lib.rs

azalea_ecs_macros/
lib.rs
