try_macro_proc_macro/
lib.rs

1use std::mem::replace;
2
3use proc_macro::{Delimiter::Brace, Group, Span, TokenStream, TokenTree};
4use quote::{quote_spanned, ToTokens};
5use syn::{
6    spanned::Spanned,
7    visit_mut::{visit_expr_mut, VisitMut},
8    Block, Expr, ExprTry, Item,
9};
10
11fn default_expr_try() -> ExprTry {
12    ExprTry {
13        attrs: vec![],
14        expr: Box::new(Expr::PLACEHOLDER),
15        question_token: Default::default(),
16    }
17}
18
19struct Visitor;
20
21impl VisitMut for Visitor {
22    fn visit_expr_mut(&mut self, node: &mut Expr) {
23        let Expr::Try(expr_try) = node else {
24            visit_expr_mut(self, node);
25            return;
26        };
27        let ExprTry {
28            attrs,
29            mut expr,
30            question_token,
31        } = replace(expr_try, default_expr_try());
32
33        self.visit_expr_mut(&mut expr);
34
35        *node = syn::parse(
36            quote_spanned! { question_token.span() =>
37                #(#attrs)*
38                match ::try_macro::Try::branch(#expr) {
39                    ::core::ops::ControlFlow::Continue(value) => value,
40                    ::core::ops::ControlFlow::Break(err) => {
41                        return ::try_macro::FromResidual::from_residual(err);
42                    }
43                }
44            }.into(),
45        ).unwrap();
46    }
47}
48
49/// Replace `?` for user `Try::branch`
50///
51/// Apply on Item, e.g `#[try_macro] fn foo() {}`
52#[proc_macro_attribute]
53pub fn try_macro(attr: TokenStream, item: TokenStream) -> TokenStream {
54    if let Some(attr) = attr.into_iter().next() {
55        return syn::Error::new(
56            attr.span().into(),
57            "invalid attribute input",
58        ).into_compile_error().into();
59    }
60    syn::parse::<Item>(item)
61        .map(|mut item| {
62            Visitor.visit_item_mut(&mut item);
63            item.into_token_stream()
64        })
65        .unwrap_or_else(|e| e.into_compile_error())
66        .into()
67}
68
69/// Replace `?` for user `Try::branch` in block
70///
71/// try_macro_block! {}
72#[proc_macro]
73pub fn try_macro_block(stream: TokenStream) -> TokenStream {
74    let braced = TokenStream::from_iter([
75        TokenTree::from(Group::new(Brace, stream))
76    ]);
77    syn::parse::<Block>(braced)
78        .map(|mut block| {
79            Visitor.visit_block_mut(&mut block);
80            block.into_token_stream()
81        })
82        .unwrap_or_else(|e| e.into_compile_error())
83        .into()
84}
85
86
87/// ```
88/// try_macro_proc_macro::__test!();
89/// ```
90#[doc(hidden)]
91#[proc_macro]
92pub fn __test(_: TokenStream) -> TokenStream {
93    use quote::quote;
94
95    let mut fails = vec![];
96    macro_rules! eq {
97        ($name:ident : {$($a:tt)*} => {$($b:tt)*}) => { #[allow(unused)] fn $name() {} {
98            let name = stringify!($name);
99
100            let (a, b) = (quote! { $($a)* }, quote! { $($b)* });
101            let a = TokenStream::from(try_macro(TokenStream::new(), a.into()));
102            let b = TokenStream::from(b);
103
104            let a = a.to_string();
105            let b = b.to_string();
106
107            if a != b {
108                fails.push(format!("`{name}` failed\n left: {a}\nright: {b}"))
109            }
110        }};
111    }
112
113    eq!(it_works: {
114        fn foo() {
115            x?
116        }
117    } => {
118        fn foo() {
119            match ::try_macro::Try::branch(x) {
120                ::core::ops::ControlFlow::Continue(value) => value,
121                ::core::ops::ControlFlow::Break(err) => {
122                    return ::try_macro::FromResidual::from_residual(err);
123                }
124            }
125        }
126    });
127
128    eq!(nesting: {
129        fn foo() {
130            x??
131        }
132    } => {
133        fn foo() {
134            match ::try_macro::Try::branch(match ::try_macro::Try::branch(x) {
135                ::core::ops::ControlFlow::Continue(value) => value,
136                ::core::ops::ControlFlow::Break(err) => {
137                    return ::try_macro::FromResidual::from_residual(err);
138                }
139            }) {
140                ::core::ops::ControlFlow::Continue(value) => value,
141                ::core::ops::ControlFlow::Break(err) => {
142                    return ::try_macro::FromResidual::from_residual(err);
143                }
144            }
145        }
146    });
147
148    eq!(inner: {
149        fn foo() {
150            x?+2
151        }
152    } => {
153        fn foo() {
154            (match ::try_macro::Try::branch(x) {
155                ::core::ops::ControlFlow::Continue(value) => value,
156                ::core::ops::ControlFlow::Break(err) => {
157                    return ::try_macro::FromResidual::from_residual(err);
158                }
159            })+2
160        }
161    });
162
163    if fails.is_empty() {
164        return TokenStream::new();
165    }
166    syn::Error::new(Span::call_site().into(), fails.join("\n\n"))
167        .into_compile_error()
168        .into()
169}
try_macro_proc_macro/lib.rs

try_macro_proc_macro/
lib.rs
