init_hook_macros/
lib.rs

1use darling::{ast::NestedMeta, error::Accumulator};
2use quote::quote;
3use syn::{Ident, ItemFn, parse_macro_input, spanned::Spanned};
4
5/// Register function to be called during initialization
6#[proc_macro_attribute]
7pub fn call_on_init(
8    args: proc_macro::TokenStream,
9    input: proc_macro::TokenStream,
10) -> proc_macro::TokenStream {
11    let mut errors = Accumulator::default();
12
13    let attr_args = match NestedMeta::parse_meta_list(args.into()) {
14        Ok(v) => v,
15        Err(err) => {
16            return darling::Error::from(err).write_errors().into();
17        }
18    };
19
20    for nested_meta in attr_args {
21        errors.push(darling::Error::custom("args not supported").with_span(&nested_meta.span()));
22    }
23
24    let input = parse_macro_input!(input as ItemFn);
25
26    if input.sig.asyncness.is_some() {
27        errors.push(
28            darling::Error::custom(
29                "the `async` keywork cannot be used on the function declaration",
30            )
31            .with_span(&input.sig.fn_token.span()),
32        );
33    }
34
35    if let Err(err) = errors.finish() {
36        return err.write_errors().into();
37    }
38
39    let ident = input.sig.ident.clone();
40    let static_ident = Ident::new(&ident.to_string().to_uppercase(), ident.span());
41
42    let expanded = if input.sig.unsafety.is_some() {
43        quote! {
44            #[init_hook::linkme::distributed_slice(init_hook::__private::UNSAFE_INIT_FNS)]
45            #[linkme(crate = init_hook::linkme)]
46            static #static_ident: unsafe fn() = #ident;
47
48            #input
49        }
50    } else {
51        quote! {
52            #[init_hook::linkme::distributed_slice(init_hook::__private::INIT_FNS)]
53            #[linkme(crate = init_hook::linkme)]
54            static #static_ident: fn() = #ident;
55
56            #input
57        }
58    };
59
60    expanded.into()
61}