1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use quote::quote;
use syn::{parse_quote, ItemTrait, Result};

/// When added to a trait declaration, generates the impls required to find a resource that implements a specific trait.
#[proc_macro_attribute]
pub fn trait_resource(attr: TokenStream, item: TokenStream) -> TokenStream {
    impl_trait_resource(attr, item)
        .unwrap_or_else(syn::Error::into_compile_error)
        .into()
}

fn impl_trait_resource(arg: TokenStream, item: TokenStream) -> Result<TokenStream2> {
    let _ = arg;
    let trait_definition = syn::parse::<ItemTrait>(item)?;
    let trait_name = trait_definition.ident.clone();

    let mut impl_generics_list = vec![];
    let mut trait_generics_list = vec![];
    let where_clause = trait_definition.generics.where_clause.clone();

    for param in &trait_definition.generics.params {
        impl_generics_list.push(param.clone());
        match param {
            syn::GenericParam::Type(param) => {
                let ident = &param.ident;
                trait_generics_list.push(quote! { #ident });
            }
            syn::GenericParam::Lifetime(param) => {
                let ident = &param.lifetime;
                trait_generics_list.push(quote! { #ident });
            }
            syn::GenericParam::Const(param) => {
                let ident = &param.ident;
                trait_generics_list.push(quote! { #ident });
            }
        }
    }

    let impl_generics = quote! { <#( #impl_generics_list ,)*> };
    let trait_generics = quote! { <#( #trait_generics_list ,)*> };

    let trait_object = quote! { dyn #trait_name #trait_generics };

    let my_crate = proc_macro_crate::crate_name("bevy-trait-resource").unwrap();
    let my_crate = match my_crate {
        proc_macro_crate::FoundCrate::Itself => quote! { crate },
        proc_macro_crate::FoundCrate::Name(x) => {
            let ident = quote::format_ident!("{x}");
            quote! { #ident }
        }
    };

    let imports = quote! { #my_crate::imports };

    let trait_resource = quote! { #my_crate::TraitResource };

    let mut marker_impl_generics_list = impl_generics_list.clone();
    marker_impl_generics_list.push(parse_quote!(__Resource: #trait_name #trait_generics + #imports::Resource));

    let marker_impl_generics = quote! { <#( #marker_impl_generics_list ,)*> };

    let marker_impl_code = quote! {
        impl #impl_generics #trait_resource for #trait_object #where_clause {}

        impl #marker_impl_generics #my_crate::TraitResourceMarker::<#trait_object> for (__Resource,)
        #where_clause
        {
            type Covered = __Resource;
            fn cast(ptr: *mut u8) -> *mut #trait_object {
                ptr as *mut __Resource as *mut _
            }
        }
    };

    Ok(quote! {
        #trait_definition
        #marker_impl_code
    })
}