rustbasic_core_macro/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{parse_macro_input, Item, TraitItem, ImplItem, Signature, ReturnType, FnArg, DeriveInput, Meta, Expr, Lit};
4
5#[proc_macro_attribute]
6pub fn async_trait(_attr: TokenStream, item: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(item as Item);
8    match input {
9        Item::Trait(mut trait_item) => {
10            for item in &mut trait_item.items {
11                if let TraitItem::Fn(method) = item {
12                    if method.sig.asyncness.is_some() {
13                        let original_body = method.default.clone();
14                        transform_signature(&mut method.sig);
15                        if let Some(body) = original_body {
16                            method.default = Some(syn::parse2(quote! {
17                                {
18                                    ::std::boxed::Box::pin(async move {
19                                        #body
20                                    })
21                                }
22                            }).unwrap());
23                        }
24                    }
25                }
26            }
27            TokenStream::from(quote!(#trait_item))
28        }
29        Item::Impl(mut impl_item) => {
30            for item in &mut impl_item.items {
31                if let ImplItem::Fn(method) = item {
32                    if method.sig.asyncness.is_some() {
33                        let original_body = method.block.clone();
34                        transform_signature(&mut method.sig);
35                        method.block = syn::parse2(quote! {
36                            {
37                                ::std::boxed::Box::pin(async move {
38                                    #original_body
39                                })
40                            }
41                        }).unwrap();
42                    }
43                }
44            }
45            TokenStream::from(quote!(#impl_item))
46        }
47        _ => TokenStream::from(quote!(#input)),
48    }
49}
50
51fn transform_signature(sig: &mut Signature) {
52    sig.asyncness = None;
53    let ret_type = match &sig.output {
54        ReturnType::Default => quote!(()),
55        ReturnType::Type(_, ty) => quote!(#ty),
56    };
57
58    // Find any lifetime in the generics. If none, default to '_
59    let mut lifetime_str = quote!('_);
60    for param in &sig.generics.params {
61        if let syn::GenericParam::Lifetime(lt) = param {
62            let lt_ident = &lt.lifetime;
63            lifetime_str = quote!(#lt_ident);
64            break;
65        }
66    }
67
68    // If there is self, and we have a specific lifetime like 'a, bind self to 'a
69    if lifetime_str.to_string() != "'_" {
70        for arg in &mut sig.inputs {
71            if let FnArg::Receiver(receiver) = arg {
72                if let Some((_, ref mut opt_lifetime)) = receiver.reference {
73                    if opt_lifetime.is_none() {
74                        let lt: syn::Lifetime = syn::parse2(quote!(#lifetime_str)).unwrap();
75                        *opt_lifetime = Some(lt);
76                    }
77                }
78            }
79        }
80    }
81
82    sig.output = syn::parse2(quote! {
83        -> ::std::pin::Pin<::std::boxed::Box<dyn ::std::future::Future<Output = #ret_type> + ::std::marker::Send + #lifetime_str>>
84    }).unwrap();
85}
86
87#[proc_macro_derive(RustEmbed, attributes(folder))]
88pub fn derive_rust_embed(input: TokenStream) -> TokenStream {
89    let ast = parse_macro_input!(input as DeriveInput);
90    let name = &ast.ident;
91
92    let mut folder_path = None;
93    for attr in &ast.attrs {
94        if attr.path().is_ident("folder") {
95            if let Meta::NameValue(meta) = &attr.meta {
96                if let Expr::Lit(expr_lit) = &meta.value {
97                    if let Lit::Str(lit_str) = &expr_lit.lit {
98                        folder_path = Some(lit_str.value());
99                    }
100                }
101            }
102        }
103    }
104
105    let folder = folder_path.expect("attribute #[folder = \"...\"] is required for RustEmbed");
106    
107    let manifest_dir = std::env::var("CARGO_MANIFEST_DIR").unwrap();
108    let abs_path = std::path::Path::new(&manifest_dir).join(&folder);
109
110    // Recursively read all files in directory
111    let mut files = Vec::new();
112    if abs_path.exists() {
113        files = get_files_recursively(&abs_path, &abs_path);
114    }
115
116    let file_paths: Vec<String> = files.iter().map(|(p, _)| p.clone()).collect();
117    let file_bytes: Vec<Vec<u8>> = files.iter().map(|(_, b)| b.clone()).collect();
118
119    let file_bytes_tokens: Vec<_> = file_bytes.iter().map(|bytes| {
120        quote! { &[#(#bytes),*] }
121    }).collect();
122
123    let gen_code = quote! {
124        impl #name {
125            pub fn get(file_path: &str) -> Option<rustbasic_core::rust_embed::EmbeddedFile> {
126                let normalized_path = file_path.replace("\\", "/");
127                match normalized_path.as_str() {
128                    #(
129                        #file_paths => Some(rustbasic_core::rust_embed::EmbeddedFile {
130                            data: std::borrow::Cow::Borrowed(#file_bytes_tokens),
131                            metadata: rustbasic_core::rust_embed::Metadata {
132                                last_modified: None,
133                                created: None,
134                                sha256_hash: [0; 32],
135                            }
136                        }),
137                    )*
138                    _ => None,
139                }
140            }
141
142            pub fn iter() -> impl Iterator<Item = std::borrow::Cow<'static, str>> {
143                let items: Vec<std::borrow::Cow<'static, str>> = vec![
144                    #( std::borrow::Cow::Borrowed(#file_paths) ),*
145                ];
146                items.into_iter()
147            }
148        }
149    };
150
151    gen_code.into()
152}
153
154fn get_files_recursively(dir: &std::path::Path, base_dir: &std::path::Path) -> Vec<(String, Vec<u8>)> {
155    let mut files = Vec::new();
156    if let Ok(entries) = std::fs::read_dir(dir) {
157        for entry in entries.flatten() {
158            let path = entry.path();
159            if path.is_dir() {
160                files.extend(get_files_recursively(&path, base_dir));
161            } else if path.is_file() {
162                if let Ok(content) = std::fs::read(&path) {
163                    let rel_path = path.strip_prefix(base_dir).unwrap().to_string_lossy().to_string();
164                    let rel_path = rel_path.replace("\\", "/");
165                    files.push((rel_path, content));
166                }
167            }
168        }
169    }
170    files
171}
rustbasic_core_macro/lib.rs

rustbasic_core_macro/
lib.rs
