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 let aeth_entry_crate = aeth_entry_crate()?;
50
51 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 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#[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}