orion_async_macros/
lib.rs

1//!
2//! orion-async-macros提供orion-async的过程宏实现部分的功能
3//!
4
5#[allow(unused_extern_crates)]
6extern crate proc_macro;
7
8use proc_macro::{ TokenStream };
9use quote::{ quote, ToTokens };
10use proc_macro2::{ Span };
11use syn::{
12    parse_macro_input,
13    ItemFn, AttributeArgs,
14    NestedMeta, Meta,
15    Lit,
16    spanned::{ Spanned },
17};
18
19struct Config {
20    send: bool,
21}
22
23impl Config {
24    const fn new() -> Self {
25        Config {
26            send: false,
27        }
28    }
29    fn build(&mut self, arg: &NestedMeta) -> Result<(), syn::Error> {
30        match arg {
31            NestedMeta::Meta(Meta::NameValue(namevalue)) => {
32                let ident = namevalue.path.get_ident();
33                if ident.is_none() {
34                    return Err(syn::Error::new_spanned(&namevalue, "should have specified ident"));
35                }
36                let ident = ident.unwrap().to_string().to_lowercase();
37                match ident.as_str() {
38                    "body_send" => {
39                        self.send = parse_bool(namevalue.lit.clone(), Spanned::span(&namevalue.lit))?;
40                    },
41                    _ => {
42                        let msg = "unknown attribute, expected is `body_send`";
43                        return Err(syn::Error::new_spanned(namevalue, msg));
44                    }
45                }
46            },
47            _ => {
48                let msg = "unknown attribute";
49                return Err(syn::Error::new_spanned(arg, msg));
50            }
51        }
52        Ok(())
53    }
54}
55
56fn parse_bool(val: Lit, span: Span) -> Result<bool, syn::Error> {
57    match val {
58        Lit::Bool(b) => Ok(b.value),
59        _ => Err(syn::Error::new(span, "value should be true or false")),
60    }
61}
62
63///
64/// 如下的实现的async函数内部使用了Rc,因为不支持Send,无法利用tokio::spawn调度
65/// 
66/// ```rust
67/// async fn foo() -> i32 {
68///     let id = Rc::new(0);
69///     bar(*id).await;
70/// }
71/// ```
72/// 如下定义即可保证Rc可在Future内部正常使用,同时保证安全和性能
73/// ```rust
74/// #[orion_async::future(body_send = true)]
75/// async fn foo() -> i32 {
76///     let id = Rc::new(0);
77///     bar(*id).await;
78/// }
79/// ```
80///
81
82#[proc_macro_attribute]
83pub fn future(args: TokenStream, item: TokenStream) -> TokenStream {
84    let item_copy = item.clone();
85    let args = parse_macro_input!(args as AttributeArgs);
86    let mut input = parse_macro_input!(item as ItemFn);
87
88    if input.sig.asyncness.is_none() {
89        let msg = "should add `async` keyword";
90        return syn::Error::new_spanned(input.sig.fn_token, msg).into_compile_error().into();
91    }
92
93    let mut conf = Config::new();
94
95    for arg in args {
96        if let Err(error) = conf.build(&arg) {
97            return error.into_compile_error().into(); 
98        }
99    }
100
101    if conf.send {
102        send_future(&mut input)
103    } else {
104        item_copy
105    }
106}
107
108fn send_future(input: &mut ItemFn) -> TokenStream {
109    let body = input.block.to_token_stream();
110    let tokens: TokenStream = quote!({
111        let future = unsafe {
112            orion_async::SendFuture::new(async { #body })
113        };
114        future.await
115    }).into();
116
117    let wrapper = parse_macro_input!(tokens as syn::Stmt);
118    input.block.stmts.clear();
119    input.block.stmts.push(wrapper);
120
121    input.to_token_stream().into()
122}