aionbot_macros/
lib.rs

1use std::hash::{DefaultHasher, Hash, Hasher};
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::{meta::ParseNestedMeta, parse_macro_input, Result};
6
7struct HandlerArgs {
8    priority: syn::LitInt,
9    router: Option<syn::Expr>,
10}
11
12impl Default for HandlerArgs {
13    fn default() -> Self {
14        Self {
15            priority: syn::LitInt::new("0", proc_macro2::Span::call_site()),
16            router: None,
17        }
18    }
19}
20
21impl HandlerArgs {
22    fn parse(&mut self, meta: ParseNestedMeta) -> Result<()> {
23        if let Some(ident) = meta.path.get_ident() {
24            match ident.to_string().as_str() {
25                "router" => {
26                    self.router = Some(meta.value()?.parse()?);
27                    Ok(())
28                }
29                "priority" => {
30                    self.priority = meta.value()?.parse()?;
31                    Ok(())
32                }
33                _ => Err(meta.error("msg")),
34            }
35        } else {
36            Err(meta.error("msg"))
37        }
38    }
39
40    fn is_empty(&self) -> bool {
41        self.router.is_none()
42    }
43}
44
45fn get_router(handler_args: &HandlerArgs) -> syn::Expr {
46    if let Some(router) = &handler_args.router {
47        router.clone()
48    } else {
49        syn::parse(quote! { "AnyRouter::default()" }.into()).unwrap()
50    }
51}
52
53fn get_hash_id(ident: &syn::Ident) -> String {
54    let mut hasher = DefaultHasher::new();
55    ident.hash(&mut hasher);
56    hasher.finish().to_string()
57}
58
59fn extract_fn_name_ident(item: &syn::Ident, hash_id: &str) -> syn::Ident {
60    let mut fn_name = String::from("__");
61    fn_name.push_str(&item.to_string());
62    fn_name.extend("_".chars().chain(hash_id.chars()));
63    syn::Ident::new(&fn_name, item.span())
64}
65
66#[proc_macro_attribute]
67pub fn register(attr: TokenStream, item: TokenStream) -> TokenStream {
68    let input = parse_macro_input!(item as syn::ItemFn);
69
70    let mut attrs = HandlerArgs::default();
71    let parser = syn::meta::parser(|meta| attrs.parse(meta));
72    parse_macro_input!(attr with parser);
73
74    let vis = &input.vis;
75    match vis {
76        syn::Visibility::Public(_) => {}
77        _ => {
78            return TokenStream::from(
79                quote! { compile_error!("Only public functions can be registered"); },
80            )
81        }
82    }
83
84    let origin_ident = &input.sig.ident;
85    let hash_id = get_hash_id(origin_ident);
86    let fn_name_ident = extract_fn_name_ident(origin_ident, &hash_id);
87
88    let fn_args = &input.sig.inputs;
89    let fn_body = &input.block;
90
91    let router = get_router(&attrs);
92    let priority = &attrs.priority;
93
94    if attrs.is_empty() {
95        return TokenStream::from(
96            quote! { compile_error!("Missing `#[register(router = \"...\")]` attribute"); },
97        );
98    };
99
100    let expanded = quote! {
101        use std::sync::*;
102        use std::cell::*;
103        use aionbot::prelude::*;
104
105        pub fn #fn_name_ident(#fn_args) -> HandlerCallback {
106            Box::pin(async move { #fn_body })
107        }
108
109        pub fn #origin_ident() -> Entry {
110            Entry {
111                id: #hash_id,
112                priority: #priority,
113                router: Arc::new(Box::new(#router)),
114                callback: Arc::new(#fn_name_ident),
115            }
116        }
117    };
118
119    TokenStream::from(expanded)
120}
aionbot_macros/lib.rs

aionbot_macros/
lib.rs
