smol_attributes/
lib.rs

1#![forbid(unsafe_code, future_incompatible, rust_2018_idioms)]
2#![deny(missing_debug_implementations, nonstandard_style)]
3#![recursion_limit = "512"]
4
5use proc_macro::TokenStream;
6use quote::{quote, quote_spanned};
7use syn::spanned::Spanned;
8
9/// Enables an async main function.
10///
11/// # Examples
12///
13/// ```ignore
14/// #[smol::main]
15/// async fn main() -> std::io::Result<()> {
16///     Ok(())
17/// }
18/// ```
19#[cfg(not(test))] // NOTE: exporting main breaks tests, we should file an issue.
20#[proc_macro_attribute]
21pub fn main(_attr: TokenStream, item: TokenStream) -> TokenStream {
22    let input = syn::parse_macro_input!(item as syn::ItemFn);
23
24    let ret = &input.sig.output;
25    let inputs = &input.sig.inputs;
26    let name = &input.sig.ident;
27    let body = &input.block;
28    let attrs = &input.attrs;
29
30    if name != "main" {
31        return TokenStream::from(quote_spanned! { name.span() =>
32            compile_error!("only the main function can be tagged with #[smol::main]"),
33        });
34    }
35
36    if input.sig.asyncness.is_none() {
37        return TokenStream::from(quote_spanned! { input.span() =>
38            compile_error!("the async keyword is missing from the function declaration"),
39        });
40    }
41
42    let result = quote! {
43        fn main() #ret {
44            #(#attrs)*
45            async fn main(#inputs) #ret {
46                #body
47            }
48
49            smol::run(async {
50                main().await
51            })
52        }
53
54    };
55
56    result.into()
57}
58
59/// Enables an async test function.
60///
61/// # Examples
62///
63/// ```ignore
64/// #[smol::test]
65/// async fn my_test() -> std::io::Result<()> {
66///     assert_eq!(2 * 2, 4);
67///     Ok(())
68/// }
69/// ```
70#[proc_macro_attribute]
71pub fn test(_attr: TokenStream, item: TokenStream) -> TokenStream {
72    let input = syn::parse_macro_input!(item as syn::ItemFn);
73
74    let ret = &input.sig.output;
75    let name = &input.sig.ident;
76    let body = &input.block;
77    let attrs = &input.attrs;
78
79    if input.sig.asyncness.is_none() {
80        return TokenStream::from(quote_spanned! { input.span() =>
81            compile_error!("the async keyword is missing from the function declaration"),
82        });
83    }
84
85    let result = quote! {
86        #[test]
87        #(#attrs)*
88        fn #name() #ret {
89            smol::run(async { #body })
90        }
91    };
92
93    result.into()
94}
95
96/// Enables an async benchmark function.
97///
98/// # Examples
99///
100/// ```ignore
101/// #![feature(test)]
102/// extern crate test;
103///
104/// #[smol::bench]
105/// async fn bench() {
106///     println!("hello world");
107/// }
108/// ```
109#[proc_macro_attribute]
110pub fn bench(_attr: TokenStream, item: TokenStream) -> TokenStream {
111    let input = syn::parse_macro_input!(item as syn::ItemFn);
112
113    let ret = &input.sig.output;
114    let args = &input.sig.inputs;
115    let name = &input.sig.ident;
116    let body = &input.block;
117    let attrs = &input.attrs;
118
119    if input.sig.asyncness.is_none() {
120        return TokenStream::from(quote_spanned! { input.span() =>
121            compile_error!("the async keyword is missing from the function declaration"),
122        });
123    }
124
125    if !args.is_empty() {
126        return TokenStream::from(quote_spanned! { args.span() =>
127            compile_error!("async benchmarks don't take any arguments"),
128        });
129    }
130
131    let result = quote! {
132        #[bench]
133        #(#attrs)*
134        fn #name(b: &mut test::Bencher) #ret {
135            let _ = b.iter(|| {
136                smol::block_on(async {
137                    #body
138                })
139            });
140        }
141    };
142
143    result.into()
144}
smol_attributes/lib.rs

smol_attributes/
lib.rs
