1 files changed, 162 insertions, 0 deletions
diff --git a/proc-macro-error/proc-macro-error-attr/src/lib.rs b/proc-macro-error/proc-macro-error-attr/src/lib.rs
new file mode 100644
index 0000000..3c1013b
--- /dev/null
+++ b/proc-macro-error/proc-macro-error-attr/src/lib.rs
@@ -0,0 +1,162 @@
+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, Token,
+};
+use syn_mid::{Block, ItemFn};
+
+use self::Setting::*;
+
+#[proc_macro_attribute]
+pub fn proc_macro_error(attr: TokenStream, input: TokenStream) -> TokenStream {
+    let input = parse_macro_input!(input as ItemFn);
+    let mut settings = match syn::parse::<Settings>(attr) {
+        Ok(settings) => settings,
+        Err(err) => {
+            let err = err.to_compile_error();
+            return quote!(#input #err).into();
+        }
+    };
+
+    let is_proc_macro = is_proc_macro(&input.attrs);
+    if is_proc_macro {
+        settings.set(AssertUnwindSafe);
+    }
+
+    if detect_proc_macro_hack(&input.attrs) {
+        settings.set(ProcMacroHack);
+    }
+
+    if !(settings.is_set(AllowNotMacro) || is_proc_macro) {
+        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,
+        constness,
+        asyncness,
+        unsafety,
+        abi,
+        fn_token,
+        ident,
+        generics,
+        inputs,
+        output,
+        block,
+        ..
+    } = input;
+
+    let body = gen_body(block, settings);
+
+    quote!(
+        #(#attrs)*
+        #vis
+        #constness
+        #asyncness
+        #unsafety
+        #abi
+        #fn_token
+        #ident
+        #generics
+        (#inputs)
+        #output
+
+        { #body }
+    )
+    .into()
+}
+
+#[derive(PartialEq)]
+enum Setting {
+    AssertUnwindSafe,
+    AllowNotMacro,
+    ProcMacroHack,
+}
+
+impl Parse for Setting {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        let ident: Ident = input.parse()?;
+        match &*ident.to_string() {
+            "assert_unwind_safe" => Ok(AssertUnwindSafe),
+            "allow_not_macro" => Ok(AllowNotMacro),
+            "proc_macro_hack" => Ok(ProcMacroHack),
+            _ => Err(syn::Error::new(
+                ident.span(),
+                format!(
+                    "unknown setting `{}`, expected one of \
+                     `assert_unwind_safe`, `allow_not_macro`, `proc_macro_hack`",
+                    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)))
+    }
+}
+
+impl Settings {
+    fn is_set(&self, setting: Setting) -> bool {
+        self.0.iter().any(|s| *s == setting)
+    }
+
+    fn set(&mut self, setting: Setting) {
+        self.0.push(setting)
+    }
+}
+
+#[rustversion::since(1.37)]
+fn gen_body(block: Block, settings: Settings) -> proc_macro2::TokenStream {
+    let is_proc_macro_hack = settings.is_set(ProcMacroHack);
+    let closure = if settings.is_set(AssertUnwindSafe) {
+        quote!(::std::panic::AssertUnwindSafe(|| #block ))
+    } else {
+        quote!(|| #block)
+    };
+
+    quote!( ::proc_macro_error::entry_point(#closure, #is_proc_macro_hack) )
+}
+
+// FIXME:
+// proc_macro::TokenStream does not implement UnwindSafe until 1.37.0.
+// Considering this is the closure's return type the unwind safety check would fail
+// for virtually every closure possible, the check is meaningless.
+#[rustversion::before(1.37)]
+fn gen_body(block: Block, settings: Settings) -> proc_macro2::TokenStream {
+    let is_proc_macro_hack = settings.is_set(ProcMacroHack);
+    let closure = quote!(::std::panic::AssertUnwindSafe(|| #block ));
+    quote!( ::proc_macro_error::entry_point(#closure, #is_proc_macro_hack) )
+}
+
+fn detect_proc_macro_hack(attrs: &[Attribute]) -> bool {
+    attrs
+        .iter()
+        .any(|attr| attr.path.is_ident("proc_macro_hack"))
+}
+
+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")
+    })
+}