1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

extern crate proc_macro;

use proc_macro::TokenStream;
use proc_macro2::Ident;
use quote::quote;
use std::iter::FromIterator;
use syn::{
    parse::{Parse, ParseStream},
    parse_macro_input,
    punctuated::Punctuated,
    Attribute, Block, ItemFn, Token,
};

use self::Setting::*;

/// **Either this attribute or [`proc_macro_error::entry_point`][ep] MUST be present
/// on the top level of your macro.**
///
/// This attribute helps you build the right [`entry_point`][ep] invocation while
/// keeping the indentation level.
///
/// # Syntax
///
/// `#[proc_macro_error]` or `#[proc_macro_error(settings...)]`, where `settings...`
/// is a comma-separated list of:
///
/// - `assert_unwind_safe`:
///
///     Tells `proc_macro_error` that the code passed to [`proc_macro_error::entry_point`][ep]
///     should we wrapped with [`AssertUnwindSafe`].
///
/// - `allow_not_macro`:
///
///     By default, the attribute checks that it's applied to a proc-macro.
///     If none of `#[proc_macro]`, `#[proc_macro_derive]` nor `#[proc_macro_attribute]` are
///     present it will panic. It's the intention - this crate is supposed to be used only with
///     proc-macros. By applying this setting you're bypassing this check, useful in certain
///     circumstances.
///
/// [ep]: https://docs.rs/proc-macro-error/0.3/proc_macro_error/fn.entry_point.html
/// [`AssertUnwindSafe`]: https://doc.rust-lang.org/std/panic/struct.AssertUnwindSafe.html
#[proc_macro_attribute]
pub fn proc_macro_error(attr: TokenStream, input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as ItemFn);
    let settings = match syn::parse::<Settings>(attr) {
        Ok(settings) => settings.0,
        Err(err) => {
            let err = err.to_compile_error();
            return quote!(#input #err).into();
        }
    };

    let allow_not_macro = settings.iter().find(|s| **s == AllowNotMacro).is_some();
    let assert_unwind_safe = settings.iter().find(|s| **s == AssertUnwindSafe).is_some();

    if !(allow_not_macro || is_proc_macro(&input.attrs)) {
        return quote!(
            #input
            compile_error!("#[proc_macro_error] attribute can be used only with a proc-macro\n\n  hint: if you are really sure that #[proc_macro_error] should be applied to this exact function use #[proc_macro_error(allow_not_macro)]\n");
        )
        .into();
    }

    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = input;

    let body = gen_body(block, assert_unwind_safe);

    quote!( #(#attrs)* #vis #sig { #body } ).into()
}

#[derive(PartialEq)]
enum Setting {
    AssertUnwindSafe,
    AllowNotMacro,
}

impl Parse for Setting {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let ident: Ident = input.parse()?;
        match &*ident.to_string() {
            "assert_unwind_safe" => Ok(Setting::AssertUnwindSafe),
            "allow_not_macro" => Ok(Setting::AllowNotMacro),
            _ => Err(syn::Error::new(
                ident.span(),
                format!(
                    "unknown setting `{}`, expected one of `assert_unwind_safe`, `allow_not_macro`",
                    ident
                ),
            )),
        }
    }
}

struct Settings(Vec<Setting>);
impl Parse for Settings {
    fn parse(input: ParseStream) -> syn::Result<Self> {
        let punct = Punctuated::<Setting, Token![,]>::parse_terminated(input)?;
        Ok(Settings(Vec::from_iter(punct)))
    }
}

#[rustversion::since(1.37)]
fn gen_body(block: Box<Block>, assert_unwind_safe: bool) -> proc_macro2::TokenStream {
    if assert_unwind_safe {
        quote!( ::proc_macro_error::entry_point(::std::panic::AssertUnwindSafe(|| #block )) )
    } else {
        quote!( ::proc_macro_error::entry_point(|| #block ) )
    }
}

#[rustversion::before(1.37)]
fn gen_body(block: Box<Block>, _assert_unwind_safe: bool) -> proc_macro2::TokenStream {
    quote! {
        // FIXME:
        // proc_macro::TokenStream does not implement UnwindSafe until 1.37.0.
        // Considering this is the closure's return type the safety check would fail
        // for virtually every closure possible, the check is meaningless.
        ::proc_macro_error::entry_point(::std::panic::AssertUnwindSafe(|| #block ))
    }
}

fn is_proc_macro(attrs: &[Attribute]) -> bool {
    attrs.iter().any(|attr| {
        attr.path.is_ident("proc_macro")
            || attr.path.is_ident("proc_macro_derive")
            || attr.path.is_ident("proc_macro_attribute")
    })
}