use crate::utils::{Cat3, CommaPunctuated, IdentAssignment};
use proc_macro::TokenStream;
use proc_macro2::TokenStream as TokenStream2;
use syn::{
Error as SynError, Expr, ItemFn,
parse::{Parse, ParseStream},
};
pub struct MethodApplications {
method_applications: Vec<TokenStream2>,
}
impl MethodApplications {
fn method_application(Cat3(ident, _comma, expr): IdentAssignment<Expr>) -> TokenStream2 {
quote::quote! { .#ident(#expr) }
}
}
impl Parse for MethodApplications {
fn parse(parse_stream: ParseStream) -> Result<Self, SynError> {
let assignments = CommaPunctuated::<IdentAssignment<Expr>>::parse_terminated(parse_stream)?;
let method_applications = assignments.into_iter().map(Self::method_application).collect();
let method_applications = Self { method_applications };
Ok(method_applications)
}
}
pub fn tokio_main(attr_args_token_stream: TokenStream, input_token_stream: TokenStream) -> TokenStream {
let MethodApplications { method_applications } = syn::parse_macro_input!(attr_args_token_stream);
let ItemFn {
attrs,
vis,
mut sig,
block,
} = syn::parse_macro_input!(input_token_stream);
sig.asyncness = None;
let tokio_main_fn_token_stream = quote::quote! {
#(#attrs)*
#vis #sig {
::tokio::runtime::Builder::new_multi_thread()
.enable_all()
#(#method_applications)*
.build()
.unwrap()
.block_on(async #block)
}
};
tokio_main_fn_token_stream.into()
}