rust_try_catch_macros/
lib.rs

1use proc_macro::TokenStream as TokenStream1;
2use proc_macro2::TokenStream;
3use quote::{quote_spanned, ToTokens};
4use syn::{parse_macro_input, parse_quote, ExprClosure, ItemFn};
5use syn::spanned::Spanned;
6use crate::parse::FnOrClosure;
7
8mod parse;
9
10
11#[proc_macro_attribute]
12pub fn throw_guard(attr: TokenStream1, item: TokenStream1) -> TokenStream1 {
13    if !attr.is_empty() {
14        let message = format!("Unexpected attributes \"{attr}\"");
15        return syn::Error::new(TokenStream::from(attr).span(), message)
16            .into_compile_error()
17            .into()
18    }
19    
20    let tokens = match parse_macro_input!(item as FnOrClosure) {
21        FnOrClosure::Function(func) => wrap_func(func),
22        FnOrClosure::Closure(closure) => wrap_closure(closure),
23    };
24    
25    tokens.into()
26}
27
28
29#[proc_macro]
30pub fn closure_throw_guard(item: TokenStream1) -> TokenStream1 {
31    wrap_closure(parse_macro_input!(item as ExprClosure)).into()
32}
33
34fn throw_guard_body(async_ness: Option<&syn::Token![async]>, fn_body: TokenStream) -> TokenStream {
35    let span = fn_body.span();
36    match &async_ness {
37        Some(_) => quote_spanned! { span=>
38            {
39                let mut fut = ::core::pin::pin!(async { #fn_body });
40                ::core::future::poll_fn(move |cx| {
41                    ::rust_try_catch::__throw_driver(|| {
42                        ::core::future::Future::poll(::core::pin::Pin::as_mut(&mut fut), cx)
43                    })
44                }).await
45            }
46        },
47        None => quote_spanned! { span=> ::rust_try_catch::__throw_driver(|| #fn_body) },
48    }
49}
50
51fn wrap_closure(closure: ExprClosure) -> TokenStream {
52    if let Some(token) = closure.constness {
53        return syn::Error::new(token.span, "can't drive try catch logic in const")
54            .into_compile_error()
55    }
56    let body_tokens = throw_guard_body(closure.asyncness.as_ref(), closure.body.into_token_stream());
57    
58    let closure = ExprClosure {
59        body: parse_quote! { #body_tokens },
60        ..closure
61    };
62
63    closure.into_token_stream()
64}
65
66fn wrap_func(input: ItemFn) -> TokenStream {
67    if let Some(token) = input.sig.constness {
68        return syn::Error::new(token.span, "can't drive try catch logic in const")
69            .into_compile_error()
70    }
71    
72    if let Some(variadic) = input.sig.variadic { 
73        return syn::Error::new(variadic.span(), "using variadic arguments would cause UB!!!")
74            .into_compile_error()
75    }
76    
77    let outer_fn_body = throw_guard_body(input.sig.asyncness.as_ref(), input.block.into_token_stream());
78    
79    let new_fn = ItemFn {
80        attrs: input.attrs,
81        vis: input.vis,
82        sig: input.sig,
83        block: parse_quote!({ #outer_fn_body }),
84    };
85    
86    new_fn.into_token_stream()
87}