stuck_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::{Ident, Span};
5use quote::quote;
6
7#[derive(Default)]
8struct Configuration {
9    crate_name: Option<Ident>,
10    parallelism: Option<usize>,
11}
12
13impl Configuration {
14    fn set_crate_name(&mut self, lit: syn::Lit) -> Result<(), syn::Error> {
15        let span = lit.span();
16        if self.crate_name.is_some() {
17            return Err(syn::Error::new(span, "crate name already set"));
18        }
19        if let syn::Lit::Str(s) = lit {
20            if let Ok(path) = s.parse::<syn::Path>() {
21                if let Some(ident) = path.get_ident() {
22                    self.crate_name = Some(ident.clone());
23                    return Ok(());
24                }
25            }
26            return Err(syn::Error::new(span, format!("invalid crate name: {}", s.value())));
27        }
28        Err(syn::Error::new(span, "invalid crate name"))
29    }
30
31    fn set_parallelism(&mut self, lit: syn::Lit) -> Result<(), syn::Error> {
32        let span = lit.span();
33        if self.parallelism.is_some() {
34            return Err(syn::Error::new(span, "parallelism already set"));
35        }
36        if let syn::Lit::Int(lit) = lit {
37            let parallelism = lit.base10_parse::<usize>()?;
38            if parallelism > 0 {
39                self.parallelism = Some(parallelism);
40                return Ok(());
41            }
42        }
43        Err(syn::Error::new(span, "parallelism should be positive integer"))
44    }
45}
46
47fn parse_config(args: syn::AttributeArgs) -> Result<Configuration, syn::Error> {
48    let mut config = Configuration::default();
49    for arg in args.into_iter() {
50        match arg {
51            syn::NestedMeta::Meta(syn::Meta::NameValue(name_value)) => {
52                let name = name_value
53                    .path
54                    .get_ident()
55                    .ok_or_else(|| syn::Error::new_spanned(&name_value, "invalid attribute name"))?
56                    .to_string();
57                match name.as_str() {
58                    "parallelism" => config.set_parallelism(name_value.lit)?,
59                    "crate" => config.set_crate_name(name_value.lit)?,
60                    _ => return Err(syn::Error::new_spanned(&name_value, "unknown attribute name")),
61                }
62            },
63            _ => return Err(syn::Error::new_spanned(arg, "unknown attribute")),
64        }
65    }
66    Ok(config)
67}
68
69fn generate(is_test: bool, attr: TokenStream, item: TokenStream) -> TokenStream {
70    let args = syn::parse_macro_input!(attr as syn::AttributeArgs);
71    let config = parse_config(args).unwrap();
72    let input = syn::parse_macro_input!(item as syn::ItemFn);
73
74    let ret = &input.sig.output;
75    let inputs = &input.sig.inputs;
76    let name = &input.sig.ident;
77    let body = &input.block;
78    let attrs = &input.attrs;
79    let vis = &input.vis;
80
81    let macro_name = if is_test { "#[stuck::test]" } else { "#[stuck::main]" };
82
83    if input.sig.asyncness.is_some() {
84        let err =
85            syn::Error::new_spanned(input, format!("only synchronous function can be tagged with {}", macro_name));
86        return TokenStream::from(err.into_compile_error());
87    }
88
89    if !is_test && name != "main" {
90        let err = syn::Error::new_spanned(name, "only the main function can be tagged with #[stuck::main]");
91        return TokenStream::from(err.into_compile_error());
92    }
93
94    let header = if is_test {
95        quote! {
96            #[::core::prelude::v1::test]
97        }
98    } else {
99        quote! {}
100    };
101
102    let crate_name = config.crate_name.unwrap_or_else(|| Ident::new("stuck", Span::call_site()));
103    let parallelism = config.parallelism.unwrap_or(0);
104    let result = quote! {
105        #header
106        #(#attrs)*
107        #vis fn #name() #ret {
108            fn entry(#inputs) #ret {
109                #body
110            }
111
112            let mut builder = #crate_name::runtime::Builder::default();
113            if #parallelism != 0 {
114                builder.parallelism(#parallelism);
115            }
116            let mut runtime = builder.build();
117            let task = runtime.spawn(entry);
118            task.join().unwrap()
119        }
120    };
121
122    result.into()
123}
124
125/// Executes marked main function in configured runtime.
126///
127/// ## Options
128/// * `parallelism`: positive integer to specify parallelism for runtime scheduler
129///
130/// ## Examples
131/// ```rust
132/// #[stuck::main]
133/// fn main() {
134///     stuck::task::yield_now();
135/// }
136/// ```
137///
138/// ```rust
139/// #[stuck::main(parallelism = 1)]
140/// fn main() {
141///     stuck::task::yield_now();
142/// }
143/// ```
144#[cfg(not(test))]
145#[proc_macro_attribute]
146pub fn main(attr: TokenStream, item: TokenStream) -> TokenStream {
147    generate(false, attr, item)
148}
149
150/// Executes marked test function in configured runtime.
151///
152/// See [macro@main] for configurable options.
153#[proc_macro_attribute]
154pub fn test(attr: TokenStream, item: TokenStream) -> TokenStream {
155    generate(true, attr, item)
156}
stuck_macros/lib.rs

stuck_macros/
lib.rs
