kioto_uring_executor_macros/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use quote::quote;
5use syn::ItemFn;
6
7fn token_stream_with_error(mut tokens: TokenStream, error: syn::Error) -> TokenStream {
8    tokens.extend(TokenStream::from(error.into_compile_error()));
9    tokens
10}
11
12#[proc_macro_attribute]
13pub fn test(_args: TokenStream, item: TokenStream) -> TokenStream {
14    let mut input: ItemFn = match syn::parse2(item.clone().into()) {
15        Ok(it) => it,
16        Err(e) => return token_stream_with_error(item, e),
17    };
18
19    if input.sig.asyncness.is_none() {
20        panic!("the `async` keyword is missing from the test declaration");
21    }
22
23    input.sig.asyncness = None;
24
25    let body = &input.block;
26    let brace_token = input.block.brace_token;
27
28    let header = quote! {
29        #[::core::prelude::v1::test]
30    };
31
32    let tokio_expr = quote! {
33        let runtime = kioto_uring_executor::Runtime::new();
34        unsafe {
35                runtime.unsafe_block_on(async {
36                    #body
37                })
38        }
39    };
40
41    input.block = syn::parse2(quote! {
42        {
43            #tokio_expr
44        }
45    })
46    .expect("Parsing failure");
47    input.block.brace_token = brace_token;
48
49    let result = quote! {
50        #header
51        #input
52    };
53
54    result.into()
55}
56
57#[proc_macro_attribute]
58pub fn main(_args: TokenStream, item: TokenStream) -> TokenStream {
59    let mut input: ItemFn = match syn::parse2(item.clone().into()) {
60        Ok(it) => it,
61        Err(e) => return token_stream_with_error(item, e),
62    };
63
64    if input.sig.asyncness.is_none() {
65        panic!("the `async` keyword is missing from the main function declaration");
66    }
67
68    if input.sig.ident != "main" {
69        panic!("Main function must be called `main`");
70    }
71
72    input.sig.asyncness = None;
73
74    let body = &input.block;
75    let brace_token = input.block.brace_token;
76
77    let tokio_expr = quote! {
78        unsafe {
79            let runtime = kioto_uring_executor::Runtime::new();
80            runtime.unsafe_block_on(async {
81                #body
82            })
83        }
84    };
85
86    input.block = syn::parse2(quote! {
87        {
88            #tokio_expr
89        }
90    })
91    .expect("Parsing failure");
92    input.block.brace_token = brace_token;
93
94    quote! {
95        #input
96    }
97    .into()
98}
kioto_uring_executor_macros/lib.rs

kioto_uring_executor_macros/
lib.rs
