Skip to main content

aeth_entry_macros/
lib.rs

1use anyhow::{Result, anyhow};
2use proc_macro::TokenStream;
3use proc_macro_crate::{Error, FoundCrate, crate_name};
4use proc_macro2::Span;
5use quote::quote;
6use syn::spanned::Spanned;
7use syn::{Ident, ItemFn, ReturnType, Type, parse_macro_input};
8
9fn find_optional_crate(name: &str) -> Result<Option<FoundCrate>> {
10    let result = crate_name(name);
11    match result {
12        Ok(result) => Ok(Some(result)),
13        Err(err) => match err {
14            Error::CrateNotFound {
15                crate_name: _,
16                path: _,
17            } => Ok(None),
18            _ => Err(err.into()),
19        },
20    }
21}
22
23fn aeth_entry_crate() -> Result<proc_macro2::TokenStream> {
24    if let Some(found) = find_optional_crate("aeth-entry")? {
25        let aeth_entry_crate = match found {
26            FoundCrate::Itself => quote!(crate),
27            FoundCrate::Name(name) => {
28                let ident = Ident::new(name.as_str(), Span::call_site());
29                quote!(::#ident)
30            }
31        };
32        return Ok(aeth_entry_crate);
33    }
34    if let Some(found) = find_optional_crate("aeth")? {
35        let aeth_entry_crate = match found {
36            FoundCrate::Itself => quote!(crate::entry),
37            FoundCrate::Name(name) => {
38                let ident = Ident::new(name.as_str(), Span::call_site());
39                quote!(::#ident::entry)
40            }
41        };
42        return Ok(aeth_entry_crate);
43    }
44    Err(anyhow!("Neither `aeth` nor `aeth-entry` crate is found."))
45}
46
47fn main_result(mut item_fn: ItemFn) -> Result<TokenStream> {
48    // Resolve the identifier of the aeth_entry crate.
49    let aeth_entry_crate = aeth_entry_crate()?;
50
51    // Make sure main function is well-formed.
52    let main_ident = item_fn.sig.ident;
53    if main_ident != "main" {
54        return Err(anyhow!("Must only decorate the main function."));
55    }
56    if item_fn.sig.asyncness.is_none() {
57        return Err(anyhow!("The main function must be async."));
58    }
59
60    // If the result is not present or `()` then we treat
61    // the main function as returning unit type. Otherwise
62    // we treat the main function as returning result.
63    let return_unit = match &item_fn.sig.output {
64        ReturnType::Default => true,
65        ReturnType::Type(_, ty) => match ty.as_ref() {
66            Type::Tuple(tpl) => tpl.elems.is_empty(),
67            Type::Never(_) => {
68                return Err(anyhow!(
69                    "We need the main function to return, in order to shutdown the event loop."
70                ));
71            }
72            _ => false,
73        },
74    };
75
76    let inner_ident = Ident::new("__inner", main_ident.span());
77    item_fn.sig.ident = inner_ident.clone();
78    let vis = item_fn.vis;
79    let sig = item_fn.sig;
80    let body = item_fn.block;
81    let future_block = if return_unit {
82        quote! {
83            async move {
84                Ok(#inner_ident().await)
85            }
86        }
87    } else {
88        quote! {
89            async move {
90                Ok(#inner_ident().await?)
91            }
92        }
93    };
94    Ok(TokenStream::from(quote! {
95        fn main() -> #aeth_entry_crate::Result<()> {
96            #vis #sig {
97                #body
98            }
99            #aeth_entry_crate::entrypoint(#future_block)
100        }
101    }))
102}
103
104/// Generate the main function for aeth-rs.
105///
106/// In order to use this macro, the user must
107/// depend on either `aeth-entry` or `aeth`.
108///
109/// This procedural macro should wrap an **async**
110/// main function. That is, an async function
111/// with identifier `main` and optional return
112/// type extending [`Result`].
113#[proc_macro_attribute]
114pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream {
115    let item_fn = parse_macro_input!(item as ItemFn);
116    let span = item_fn.span();
117    match main_result(item_fn) {
118        Ok(result) => result,
119        Err(err) => {
120            let err = syn::Error::new(span, err);
121            err.to_compile_error().into()
122        }
123    }
124}