Allow proc_macro functions to whitelist specific attributes

By using a second attribute `attributes(Bar)` on
proc_macro_derive, whitelist any attributes with
the name `Bar` in the deriving item. This allows
a proc_macro function to use custom attribtues
without a custom attribute error or unused attribute
lint.

1 files changed, 33 insertions, 11 deletions

diff --git a/src/libsyntax_ext/deriving/custom.rs b/src/libsyntax_ext/deriving/custom.rs
index f8cb1294a66..2b80deded0a 100644
--- a/src/libsyntax_ext/deriving/custom.rs
+++ b/src/libsyntax_ext/deriving/custom.rs
@@ -12,20 +12,37 @@ use std::panic;
 
 use errors::FatalError;
 use proc_macro::{TokenStream, __internal};
-use syntax::ast::{self, ItemKind};
+use syntax::ast::{self, ItemKind, Attribute};
+use syntax::attr::{mark_used, mark_known};
 use syntax::codemap::{ExpnInfo, MacroAttribute, NameAndSpan, Span};
 use syntax::ext::base::*;
 use syntax::fold::Folder;
+use syntax::parse::token::InternedString;
 use syntax::parse::token::intern;
 use syntax::print::pprust;
+use syntax::visit::Visitor;
+
+struct MarkAttrs<'a>(&'a [InternedString]);
+
+impl<'a> Visitor for MarkAttrs<'a> {
+    fn visit_attribute(&mut self, attr: &Attribute) {
+        if self.0.contains(&attr.name()) {
+            mark_used(attr);
+            mark_known(attr);
+        }
+    }
+}
 
 pub struct CustomDerive {
     inner: fn(TokenStream) -> TokenStream,
+    attrs: Vec<InternedString>,
 }
 
 impl CustomDerive {
-    pub fn new(inner: fn(TokenStream) -> TokenStream) -> CustomDerive {
-        CustomDerive { inner: inner }
+    pub fn new(inner: fn(TokenStream) -> TokenStream,
+               attrs: Vec<InternedString>)
+               -> CustomDerive {
+        CustomDerive { inner: inner, attrs: attrs }
     }
 }
 
@@ -47,7 +64,7 @@ impl MultiItemModifier for CustomDerive {
         };
         match item.node {
             ItemKind::Struct(..) |
-            ItemKind::Enum(..) => {}
+            ItemKind::Enum(..) => {},
             _ => {
                 ecx.span_err(span, "custom derive attributes may only be \
                                     applied to struct/enum items");
@@ -55,6 +72,9 @@ impl MultiItemModifier for CustomDerive {
             }
         }
 
+        // Mark attributes as known, and used.
+        MarkAttrs(&self.attrs).visit_item(&item);
+
         let input_span = Span {
             expn_id: ecx.codemap().record_expansion(ExpnInfo {
                 call_site: span,
@@ -66,12 +86,13 @@ impl MultiItemModifier for CustomDerive {
             }),
             ..item.span
         };
-        let input = __internal::new_token_stream(item);
+
+        let input = __internal::new_token_stream(item.clone());
         let res = __internal::set_parse_sess(&ecx.parse_sess, || {
             let inner = self.inner;
             panic::catch_unwind(panic::AssertUnwindSafe(|| inner(input)))
         });
-        let item = match res {
+        let new_items = match res {
             Ok(stream) => __internal::token_stream_items(stream),
             Err(e) => {
                 let msg = "custom derive attribute panicked";
@@ -88,12 +109,13 @@ impl MultiItemModifier for CustomDerive {
             }
         };
 
-        // Right now we have no knowledge of spans at all in custom derive
-        // macros, everything is just parsed as a string. Reassign all spans to
-        // the input `item` for better errors here.
-        item.into_iter().flat_map(|item| {
+        let mut res = vec![Annotatable::Item(item)];
+        // Reassign spans of all expanded items to the input `item`
+        // for better errors here.
+        res.extend(new_items.into_iter().flat_map(|item| {
             ChangeSpan { span: input_span }.fold_item(item)
-        }).map(Annotatable::Item).collect()
+        }).map(Annotatable::Item));
+        res
     }
 }