1extern 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#[proc_macro]
22pub fn sus(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
23 let input = parse_macro_input!(input as Expr);
24 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 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 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
82fn 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}