1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151

//!
//! orion-async-macros提供orion-async的过程宏实现部分的功能
//!

#[allow(unused_extern_crates)]
extern crate proc_macro;

use proc_macro::{ TokenStream };
use quote::{ quote, ToTokens };
use proc_macro2::{ Ident, Span };
use syn::{
    parse_macro_input,
    ItemFn, AttributeArgs,
    NestedMeta, Meta,
    Lit,
    spanned::{ Spanned },
};

struct Config {
    send: bool,
}

impl Config {
    const fn new() -> Self {
        Config {
            send: false,
        }
    }
    fn build(&mut self, arg: &NestedMeta) -> Result<(), syn::Error> {
        match arg {
            NestedMeta::Meta(Meta::NameValue(namevalue)) => {
                let ident = namevalue.path.get_ident();
                if ident.is_none() {
                    return Err(syn::Error::new_spanned(&namevalue, "should have specified ident"));
                }
                let ident = ident.unwrap().to_string().to_lowercase();
                match ident.as_str() {
                    "body_send" => {
                        self.send = parse_bool(namevalue.lit.clone(), Spanned::span(&namevalue.lit))?;
                    },
                    _ => {
                        let msg = "unknown attribute, expected is `body_send`";
                        return Err(syn::Error::new_spanned(namevalue, msg));
                    }
                }
            },
            _ => {
                let msg = "unknown attribute";
                return Err(syn::Error::new_spanned(arg, msg));
            }
        }
        Ok(())
    }
}

fn parse_bool(val: Lit, span: Span) -> Result<bool, syn::Error> {
    match val {
        Lit::Bool(b) => Ok(b.value),
        _ => Err(syn::Error::new(span, "value should be true or false")),
    }
}

///
/// 如下的实现的async函数内部使用了Rc，因为不支持Send，无法利用tokio::spawn调度
/// 
/// ```rust
/// async fn foo() -> i32 {
///     let id = Rc::new(0);
///     bar(*id).await;
/// }
/// ```
/// 如下定义即可保证Rc可在Future内部正常使用，同时保证安全和性能
/// ```rust
/// #[orion_async::future(body_send = true)]
/// async fn foo() -> i32 {
///     let id = Rc::new(0);
///     bar(*id).await;
/// }
/// ```
///

#[proc_macro_attribute]
pub fn future(args: TokenStream, item: TokenStream) -> TokenStream {
    let item_copy = item.clone();
    let args = parse_macro_input!(args as AttributeArgs);
    let mut input = parse_macro_input!(item as ItemFn);

    if input.sig.asyncness.is_none() {
        let msg = "should add `async` keyword";
        return syn::Error::new_spanned(input.sig.fn_token, msg).into_compile_error().into();
    }

    let mut conf = Config::new();

    for arg in args {
        if let Err(error) = conf.build(&arg) {
            return error.into_compile_error().into(); 
        }
    }

    if conf.send {
        send_future(&mut input)
    } else {
        item_copy
    }
}

fn send_future(input: &mut ItemFn) -> TokenStream {
    let mut inner = input.clone();

    let name = format!("_1234_{}", inner.sig.ident.to_string());
    let fn_ident = Ident::new(&name, inner.sig.ident.span());

    inner.sig.ident = fn_ident.clone();
    inner.attrs.clear(); 
    inner.vis = syn::Visibility::Inherited;
    
    input.block.stmts.clear();
    input.block.stmts.push(syn::Stmt::Item(syn::Item::Fn(inner)));

    let args: Vec<proc_macro2::TokenStream> = input.sig.inputs.iter().map(|arg| {
        match arg {
            syn::FnArg::Receiver(_) => {
                let ident = Ident::new("self", input.sig.ident.span());
                quote!(#ident)
            },
            syn::FnArg::Typed(ref arg) => {
                match &*arg.pat {
                    syn::Pat::Ident(ident) => {
                        let ident = &ident.ident;
                        quote!(#ident)
                    }
                    _ => {
                        let msg = "failed to parse argument";
                        return syn::Error::new_spanned(arg, msg).into_compile_error().into();
                    }
                }
            },
        }
    }).collect();

    let tokens: TokenStream = quote!({
        let future = unsafe { orion_async::SendFuture::new(#fn_ident(#(#args),*)) };
        future.await
    }).into();

    let wrapper = parse_macro_input!(tokens as syn::Stmt);
    input.block.stmts.push(wrapper);

    input.to_token_stream().into()
}