conerror_macro/
lib.rs

1use proc_macro::TokenStream;
2
3use quote::quote;
4use syn::parse::discouraged::Speculative;
5use syn::parse::{Parse, ParseStream};
6use syn::spanned::Spanned;
7use syn::visit_mut::{visit_expr_try_mut, visit_impl_item_fn_mut, VisitMut};
8use syn::{parse_macro_input, parse_quote_spanned, ExprTry, ImplItemFn, ItemFn, ItemImpl, Type};
9
10#[proc_macro_attribute]
11pub fn conerror(_: TokenStream, input: TokenStream) -> TokenStream {
12    match parse_macro_input!(input as Item) {
13        Item::Fn(mut f) => {
14            MapErr::new(None, Some(f.sig.ident.to_string())).visit_item_fn_mut(&mut f);
15            quote!(#f).into()
16        }
17        Item::Impl(mut i) => {
18            MapErr::new(Some(i.self_ty.clone()), None).visit_item_impl_mut(&mut i);
19            quote!(#i).into()
20        }
21    }
22}
23
24enum Item {
25    Fn(ItemFn),
26    Impl(ItemImpl),
27}
28
29impl Parse for Item {
30    fn parse(input: ParseStream) -> syn::Result<Self> {
31        let ahead = input.fork();
32        match ahead.parse::<ItemFn>() {
33            Ok(v) => {
34                input.advance_to(&ahead);
35                Ok(Item::Fn(v))
36            }
37            Err(e) => match input.parse::<ItemImpl>() {
38                Ok(v) => Ok(Item::Impl(v)),
39                Err(mut e1) => {
40                    e1.combine(e);
41                    Err(e1)
42                }
43            },
44        }
45    }
46}
47
48struct MapErr {
49    self_ty: Option<Box<Type>>,
50    ident: Option<String>,
51}
52
53impl MapErr {
54    fn new(self_ty: Option<Box<Type>>, ident: Option<String>) -> Self {
55        Self { self_ty, ident }
56    }
57}
58
59impl VisitMut for MapErr {
60    fn visit_expr_try_mut(&mut self, i: &mut ExprTry) {
61        let ident = self.ident.as_ref().unwrap();
62        let module = match self.self_ty {
63            Some(ref v) => quote!(std::any::type_name::<#v>()),
64            None => quote!(module_path!()),
65        };
66        let expr = &i.expr;
67        *i.expr = parse_quote_spanned! {expr.span() =>
68            #expr.map_err(|err| conerror::Error::chain(err, file!(), line!(), #ident, #module))
69        };
70        visit_expr_try_mut(self, i);
71    }
72
73    fn visit_impl_item_fn_mut(&mut self, i: &mut ImplItemFn) {
74        let mut indices = vec![];
75        for (i, attr) in i.attrs.iter().enumerate() {
76            if attr.path().is_ident("conerror") {
77                indices.push(i);
78            }
79        }
80        if indices.is_empty() {
81            return;
82        }
83
84        for idx in indices {
85            i.attrs.remove(idx);
86        }
87        self.ident = Some(i.sig.ident.to_string());
88        visit_impl_item_fn_mut(self, i);
89    }
90}