Skip to main content

asynchelp_macros/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::ext::IdentExt;
4use syn::{parse_quote, Ident, ImplItemMethod, ReturnType, AttributeArgs};
5use darling::FromMeta;
6
7#[derive(Debug, FromMeta)]
8struct StackfutureArgs {
9    #[darling(default)]
10    size: usize,
11}
12
13#[allow(unused)]
14fn snake_to_camel(ident_str: &str) -> String {
15    let mut camel_ty = String::with_capacity(ident_str.len());
16
17    let mut last_char_was_underscore = true;
18    for c in ident_str.chars() {
19        match c {
20            '_' => last_char_was_underscore = true,
21            c if last_char_was_underscore => {
22                camel_ty.extend(c.to_uppercase());
23                last_char_was_underscore = false;
24            }
25            c => camel_ty.extend(c.to_lowercase()),
26        }
27    }
28
29    camel_ty.shrink_to_fit();
30    camel_ty
31}
32
33#[allow(unused)]
34fn associated_type_for_rpc(method: &ImplItemMethod) -> String {
35    snake_to_camel(&method.sig.ident.unraw().to_string()) + "Fut"
36}
37
38#[cfg(feature = "stackfuture")]
39#[proc_macro_attribute]
40pub fn tarpc_stackfuture(attr: TokenStream, input: TokenStream) -> TokenStream {
41    let attr_args = syn::parse_macro_input!(attr as AttributeArgs);
42    let mut method = syn::parse_macro_input!(input as ImplItemMethod);
43
44    let StackfutureArgs { size } = match StackfutureArgs::from_list(&attr_args) {
45        Ok(v) => v,
46        Err(e) => { return TokenStream::from(e.write_errors()); }
47    };
48
49
50    method.sig.asyncness = None;
51
52    // get either the return type or ().
53    let ret = match &method.sig.output {
54        ReturnType::Default => quote!(()),
55        ReturnType::Type(_, ret) => quote!(#ret),
56    };
57
58    let fut_name = associated_type_for_rpc(&method);
59    let fut_name_ident = Ident::new(&fut_name, method.sig.ident.span());
60
61    // generate the updated return signature.
62    method.sig.output = parse_quote! {
63        -> ::stackfuture::StackFuture<'static, #ret, #size>
64    };
65
66    // transform the body of the method into Box::pin(async move { body }).
67    let block = method.block.clone();
68    method.block = parse_quote! [{
69        ::stackfuture::StackFuture::from(async move
70            #block
71        )
72    }];
73
74    // generate and return type declaration for return type.
75    TokenStream::from(quote! {
76        type #fut_name_ident = ::stackfuture::StackFuture<'static, #ret, #size>;
77        #method
78    })
79}