implicit_await_macro/
lib.rs

1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4use proc_macro2::{TokenStream as TokenStream2, TokenTree, Span};
5use syn::{parse_macro_input, parse_quote, ItemFn, Type, AngleBracketedGenericArguments, PathArguments, Ident};
6use syn::fold::Fold;
7
8#[proc_macro_attribute]
9pub fn implicit_await(_attr: TokenStream, input: TokenStream) -> TokenStream {
10    let item_fn = parse_macro_input!(input as ItemFn);
11    let mut folder = ImplicitAwaitFold { defer: 0 };
12    let transformed_fn = folder.fold_item_fn(item_fn);
13    quote::quote!(#transformed_fn).into()
14}
15
16struct ImplicitAwaitFold {
17    defer: usize
18}
19
20impl Fold for ImplicitAwaitFold {
21    fn fold_expr(&mut self, expr: syn::Expr) -> syn::Expr {
22        let subfolded = syn::fold::fold_expr(self, expr);
23        match (self.defer, &subfolded) {
24            (0, syn::Expr::Call(call)) => {
25                parse_quote!({{ #call }.as_future().await})
26            },
27            (0, syn::Expr::MethodCall(method_call)) => {
28                parse_quote!({{ #method_call }.as_future().await})
29            },
30            _ => subfolded
31        }
32    }
33
34    fn fold_impl_item_method(&mut self, method: syn::ImplItemMethod) -> syn::ImplItemMethod {
35        let mut subfolded = syn::fold::fold_impl_item_method(self, method);
36        subfolded.block.stmts.insert(0, parse_quote!(use implicit_await::as_future::FutureAsFuture;));
37        subfolded.block.stmts.insert(0, parse_quote!(use implicit_await::as_future::NonFutureAsFuture;));
38        subfolded
39    }
40
41    fn fold_item_fn(&mut self, func: syn::ItemFn) -> syn::ItemFn {
42        let mut subfolded = syn::fold::fold_item_fn(self, func);
43        subfolded.block.stmts.insert(0, parse_quote!(use implicit_await::as_future::FutureAsFuture;));
44        subfolded.block.stmts.insert(0, parse_quote!(use implicit_await::as_future::NonFutureAsFuture;));
45        subfolded
46    }
47
48    fn fold_expr_macro(&mut self, mac: syn::ExprMacro) -> syn::ExprMacro {
49        let is_defer = mac.mac.path.segments.iter()
50            .last()
51            .map_or(false, |segment| segment.ident.to_string() == "defer");
52        if is_defer {
53            self.defer += 1;
54        }
55        let subfolded = syn::fold::fold_expr_macro(self, mac);
56        if is_defer {
57            self.defer -= 1;
58        }
59        subfolded
60    }
61}
62
63// Meant to be used by third-party crates outside of implicit-await.
64#[proc_macro]
65pub fn as_future(input: TokenStream) -> TokenStream {
66    as_future_impl(TokenStream2::from(input), "implicit_await")
67}
68
69// Meant to be used by the implicit-await crate. Different from as_future due to lack of $crate support.
70#[proc_macro]
71pub fn as_future_internal(input: TokenStream) -> TokenStream {
72    as_future_impl(TokenStream2::from(input), "crate")
73}
74
75// Note: want this to return a TokenStream2, but the parse_macro_input! macro seems to only work in
76// functions which return a TokenStream, so I'll need to replace usage of that macro first.
77fn as_future_impl(input: TokenStream2, crate_prefix: &'static str) -> TokenStream {
78    let crate_prefix = Ident::new(crate_prefix, Span::call_site());
79    let mut types: Vec<Type> = Vec::new();
80    let mut current_type: TokenStream2 = TokenStream2::new();
81    for token in input {
82        if !is_delimiter(&token) {
83            current_type.extend(vec![token.clone()].drain(..));
84            continue;
85        }
86        let current_type_as_proc_macro_stream = TokenStream::from(current_type);
87        let parsed_type: Type = parse_macro_input!(current_type_as_proc_macro_stream as Type);
88        types.push(parsed_type);
89        current_type = TokenStream2::new();
90    }
91    if !current_type.is_empty() {
92        let current_type_as_proc_macro_stream = TokenStream::from(current_type);
93        let parsed_type: Type = parse_macro_input!(current_type_as_proc_macro_stream as Type);
94        types.push(parsed_type);
95    }
96    let mut output = TokenStream2::new();
97    for impl_type in &types {
98        let generic_args = get_generic_args(&impl_type);
99        output.extend(quote::quote!(
100            impl #generic_args #crate_prefix::as_future::NonFutureAsFuture for #impl_type {
101                fn as_future(self) -> #crate_prefix::as_future::Ready<Self> {
102                    #crate_prefix::as_future::ready(self)
103                }
104            }
105        ));
106    }
107    output.into()
108}
109
110fn is_delimiter(token: &TokenTree) -> bool {
111    match token {
112        TokenTree::Punct(punct) => punct.as_char() == ',',
113        _ => false
114    }
115}
116
117fn get_generic_args(impl_type: &Type) -> Option<AngleBracketedGenericArguments> {
118    if let Type::Path(type_path) = impl_type {
119        type_path.path.segments.iter()
120            .filter_map(|path_segment| match &path_segment.arguments {
121                PathArguments::AngleBracketed(args) => Some(args),
122                _ => None
123            })
124            .last()
125            .cloned()
126    } else {
127        None
128    }
129}
implicit_await_macro/lib.rs

implicit_await_macro/
lib.rs
