susync_macros/
lib.rs

1//! Support crate that contains the function-like procedural macros for [susync].
2//!
3//! All documentation lives in that crate.
4//!
5//! [susync]: https://docs.rs/susync
6
7extern crate proc_macro;
8
9use proc_macro2::{Span, TokenStream};
10use quote::{quote, quote_spanned};
11use syn::{
12    parse_macro_input, punctuated::Punctuated, Expr, ExprCall, ExprMethodCall, Pat, PatIdent,
13    PatType,
14};
15
16/// Generate the boilerplate for the use case where the future output is equal, or similar, to the callback arguments.
17///
18/// See full [documentation] for more details.
19///
20/// [documentation]: https://docs.rs/susync
21#[proc_macro]
22pub fn sus(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
23    let input = parse_macro_input!(input as Expr);
24    // Extract method/function call arguments
25    let args = match &input {
26        Expr::Call(expr_call) => &expr_call.args,
27        Expr::MethodCall(method_call) => &method_call.args,
28        _ => panic!("suspend macro can only be applied to function/method calls"),
29    };
30    // Extract the last closure argument from the function call
31    let _ = args
32        .iter()
33        .rev()
34        .find(|arg| matches!(arg, syn::Expr::Closure(_)))
35        .expect("expected closure as an argument");
36
37    let handle_ident = quote_spanned!(Span::mixed_site()=> handle);
38    // Leave normal args the same and modify closure argument
39    let mut found_first_closure = false;
40    let gen_args = args
41        .into_iter()
42        .rev()
43        .map(|arg| match arg {
44            Expr::Closure(expr_closure) if !found_first_closure => {
45                found_first_closure = true;
46                Expr::Closure(generate_closure(&handle_ident, expr_closure.clone()))
47            }
48            _ => arg.clone(),
49        })
50        .collect::<Vec<_>>();
51
52    let call_fn = match input {
53        Expr::Call(expr_call) => {
54            let call = ExprCall {
55                args: Punctuated::from_iter(gen_args.into_iter().rev()),
56                ..expr_call
57            };
58            quote! {
59                #call;
60            }
61        }
62        Expr::MethodCall(method_call) => {
63            let call = ExprMethodCall {
64                args: Punctuated::from_iter(gen_args.into_iter().rev()),
65                ..method_call
66            };
67            quote! {
68                #call;
69            }
70        }
71        _ => panic!("suspend macro can only be applied to function/method calls"),
72    };
73
74    quote! {
75        ::susync::suspend(|#handle_ident| {
76            let _ = #call_fn;
77        })
78    }
79    .into()
80}
81
82// Helper function to generate the closure
83fn generate_closure(captured_handle: &TokenStream, closure: syn::ExprClosure) -> syn::ExprClosure {
84    let args = closure
85        .inputs
86        .iter()
87        .flat_map(|arg_pat| match arg_pat {
88            Pat::Ident(PatIdent { ident, .. }) => Some(ident),
89            Pat::Type(PatType { pat, .. }) => match pat.as_ref() {
90                Pat::Ident(PatIdent { ident, .. }) => Some(ident),
91                _ => None,
92            },
93            Pat::Wild(_) => None,
94            _ => panic!("invalid closure arguments"),
95        })
96        .collect::<Vec<_>>();
97
98    let handle_stmt = if args.len() == 1 {
99        quote! {
100            #captured_handle.complete(#(#args.to_owned())*)
101        }
102    } else {
103        quote! {
104            #captured_handle.complete(( #(#args.to_owned(),)* ))
105        }
106    };
107
108    let body = &closure.body;
109    let body = quote! {
110        {
111            let expr_result = #body;
112            #handle_stmt;
113            expr_result
114        }
115    };
116
117    syn::ExprClosure {
118        body: Box::new(syn::Expr::Verbatim(body)),
119        ..closure
120    }
121}