bevy_trait_resource_macro/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::TokenStream as TokenStream2;
3use quote::quote;
4use syn::{parse_quote, ItemTrait, Result};
5
6/// When added to a trait declaration, generates the impls required to find a resource that implements a specific trait.
7#[proc_macro_attribute]
8pub fn trait_resource(attr: TokenStream, item: TokenStream) -> TokenStream {
9    impl_trait_resource(attr, item)
10        .unwrap_or_else(syn::Error::into_compile_error)
11        .into()
12}
13
14fn impl_trait_resource(arg: TokenStream, item: TokenStream) -> Result<TokenStream2> {
15    let _ = arg;
16    let trait_definition = syn::parse::<ItemTrait>(item)?;
17    let trait_name = trait_definition.ident.clone();
18
19    let mut impl_generics_list = vec![];
20    let mut trait_generics_list = vec![];
21    let where_clause = trait_definition.generics.where_clause.clone();
22
23    for param in &trait_definition.generics.params {
24        impl_generics_list.push(param.clone());
25        match param {
26            syn::GenericParam::Type(param) => {
27                let ident = &param.ident;
28                trait_generics_list.push(quote! { #ident });
29            }
30            syn::GenericParam::Lifetime(param) => {
31                let ident = &param.lifetime;
32                trait_generics_list.push(quote! { #ident });
33            }
34            syn::GenericParam::Const(param) => {
35                let ident = &param.ident;
36                trait_generics_list.push(quote! { #ident });
37            }
38        }
39    }
40
41    let impl_generics = quote! { <#( #impl_generics_list ,)*> };
42    let trait_generics = quote! { <#( #trait_generics_list ,)*> };
43
44    let trait_object = quote! { dyn #trait_name #trait_generics };
45
46    let my_crate = proc_macro_crate::crate_name("bevy-trait-resource").unwrap();
47    let my_crate = match my_crate {
48        proc_macro_crate::FoundCrate::Itself => quote! { crate },
49        proc_macro_crate::FoundCrate::Name(x) => {
50            let ident = quote::format_ident!("{x}");
51            quote! { #ident }
52        }
53    };
54
55    let imports = quote! { #my_crate::imports };
56
57    let trait_resource = quote! { #my_crate::TraitResource };
58
59    let mut marker_impl_generics_list = impl_generics_list.clone();
60    marker_impl_generics_list.push(parse_quote!(__Resource: #trait_name #trait_generics + #imports::Resource));
61
62    let marker_impl_generics = quote! { <#( #marker_impl_generics_list ,)*> };
63
64    let marker_impl_code = quote! {
65        impl #impl_generics #trait_resource for #trait_object #where_clause {}
66
67        impl #marker_impl_generics #my_crate::TraitResourceMarker::<#trait_object> for (__Resource,)
68        #where_clause
69        {
70            type Covered = __Resource;
71            fn cast(ptr: *mut u8) -> *mut #trait_object {
72                ptr as *mut __Resource as *mut _
73            }
74        }
75    };
76
77    Ok(quote! {
78        #trait_definition
79        #marker_impl_code
80    })
81}