macros 1.1: Allow proc_macro functions to declare attributes to be mark as used

src/libproc_macro/lib.rs

@@ -95,7 +95,8 @@ pub mod __internal {
     pub trait Registry {
         fn register_custom_derive(&mut self,
                                   trait_name: &str,
-                                  expand: fn(TokenStream) -> TokenStream);
+                                  expand: fn(TokenStream) -> TokenStream,
+                                  attributes: &[&'static str]);
     }
 
     // Emulate scoped_thread_local!() here essentially

src/librustc_lint/unused.rs

@@ -19,7 +19,7 @@ use std::collections::hash_map::Entry::{Occupied, Vacant};
 
 use syntax::ast;
 use syntax::attr;
-use syntax::feature_gate::{KNOWN_ATTRIBUTES, AttributeType};
+use syntax::feature_gate::{BUILTIN_ATTRIBUTES, AttributeType};
 use syntax::parse::token::keywords;
 use syntax::ptr::P;
 use syntax_pos::Span;
@@ -245,7 +245,7 @@ impl LateLintPass for UnusedAttributes {
         debug!("checking attribute: {:?}", attr);
 
         // Note that check_name() marks the attribute as used if it matches.
-        for &(ref name, ty, _) in KNOWN_ATTRIBUTES {
+        for &(ref name, ty, _) in BUILTIN_ATTRIBUTES {
             match ty {
                 AttributeType::Whitelisted if attr.check_name(name) => {
                     debug!("{:?} is Whitelisted", name);
@@ -267,7 +267,7 @@ impl LateLintPass for UnusedAttributes {
             debug!("Emitting warning for: {:?}", attr);
             cx.span_lint(UNUSED_ATTRIBUTES, attr.span, "unused attribute");
             // Is it a builtin attribute that must be used at the crate level?
-            let known_crate = KNOWN_ATTRIBUTES.iter()
+            let known_crate = BUILTIN_ATTRIBUTES.iter()
                 .find(|&&(name, ty, _)| attr.name() == name && ty == AttributeType::CrateLevel)
                 .is_some();
 

src/librustc_metadata/creader.rs

@@ -624,8 +624,12 @@ impl<'a> CrateLoader<'a> {
         impl Registry for MyRegistrar {
             fn register_custom_derive(&mut self,
                                       trait_name: &str,
-                                      expand: fn(TokenStream) -> TokenStream) {
-                let derive = SyntaxExtension::CustomDerive(Box::new(CustomDerive::new(expand)));
+                                      expand: fn(TokenStream) -> TokenStream,
+                                      attributes: &[&'static str]) {
+                let attrs = attributes.iter().map(|s| InternedString::new(s)).collect();
+                let derive = SyntaxExtension::CustomDerive(
+                    Box::new(CustomDerive::new(expand, attrs))
+                );
                 self.0.push((intern(trait_name), derive));
             }
         }

src/libsyntax/attr.rs

@@ -32,7 +32,8 @@ use std::cell::{RefCell, Cell};
 use std::collections::HashSet;
 
 thread_local! {
-    static USED_ATTRS: RefCell<Vec<u64>> = RefCell::new(Vec::new())
+    static USED_ATTRS: RefCell<Vec<u64>> = RefCell::new(Vec::new());
+    static KNOWN_ATTRS: RefCell<Vec<u64>> = RefCell::new(Vec::new());
 }
 
 enum AttrError {
@@ -81,6 +82,29 @@ pub fn is_used(attr: &Attribute) -> bool {
     })
 }
 
+pub fn mark_known(attr: &Attribute) {
+    debug!("Marking {:?} as known.", attr);
+    let AttrId(id) = attr.node.id;
+    KNOWN_ATTRS.with(|slot| {
+        let idx = (id / 64) as usize;
+        let shift = id % 64;
+        if slot.borrow().len() <= idx {
+            slot.borrow_mut().resize(idx + 1, 0);
+        }
+        slot.borrow_mut()[idx] |= 1 << shift;
+    });
+}
+
+pub fn is_known(attr: &Attribute) -> bool {
+    let AttrId(id) = attr.node.id;
+    KNOWN_ATTRS.with(|slot| {
+        let idx = (id / 64) as usize;
+        let shift = id % 64;
+        slot.borrow().get(idx).map(|bits| bits & (1 << shift) != 0)
+            .unwrap_or(false)
+    })
+}
+
 impl NestedMetaItem {
     /// Returns the MetaItem if self is a NestedMetaItemKind::MetaItem.
     pub fn meta_item(&self) -> Option<&P<MetaItem>> {

src/libsyntax/feature_gate.rs

@@ -417,11 +417,11 @@ macro_rules! cfg_fn {
 }
 
 pub fn deprecated_attributes() -> Vec<&'static (&'static str, AttributeType, AttributeGate)> {
-    KNOWN_ATTRIBUTES.iter().filter(|a| a.2.is_deprecated()).collect()
+    BUILTIN_ATTRIBUTES.iter().filter(|a| a.2.is_deprecated()).collect()
 }
 
 // Attributes that have a special meaning to rustc or rustdoc
-pub const KNOWN_ATTRIBUTES: &'static [(&'static str, AttributeType, AttributeGate)] = &[
+pub const BUILTIN_ATTRIBUTES: &'static [(&'static str, AttributeType, AttributeGate)] = &[
     // Normal attributes
 
     ("warn", Normal, Ungated),
@@ -790,12 +790,12 @@ impl<'a> Context<'a> {
     fn check_attribute(&self, attr: &ast::Attribute, is_macro: bool) {
         debug!("check_attribute(attr = {:?})", attr);
         let name = &*attr.name();
-        for &(n, ty, ref gateage) in KNOWN_ATTRIBUTES {
+        for &(n, ty, ref gateage) in BUILTIN_ATTRIBUTES {
             if n == name {
                 if let &Gated(_, ref name, ref desc, ref has_feature) = gateage {
                     gate_feature_fn!(self, has_feature, attr.span, name, desc);
                 }
-                debug!("check_attribute: {:?} is known, {:?}, {:?}", name, ty, gateage);
+                debug!("check_attribute: {:?} is builtin, {:?}, {:?}", name, ty, gateage);
                 return;
             }
         }
@@ -815,6 +815,8 @@ impl<'a> Context<'a> {
                            are reserved for internal compiler diagnostics");
         } else if name.starts_with("derive_") {
             gate_feature!(self, custom_derive, attr.span, EXPLAIN_DERIVE_UNDERSCORE);
+        } else if attr::is_known(attr) {
+            debug!("check_attribute: {:?} is known", name);
         } else {
             // Only run the custom attribute lint during regular
             // feature gate checking. Macro gating runs

src/libsyntax_ext/deriving/custom.rs

@@ -12,28 +12,43 @@ use std::panic;
 
 use errors::FatalError;
 use proc_macro::{TokenStream, __internal};
-use syntax::ast::{self, ItemKind};
-use syntax::codemap::{ExpnInfo, MacroAttribute, NameAndSpan, Span};
+use syntax::ast::{self, ItemKind, Attribute};
+use syntax::attr::{mark_used, mark_known};
+use syntax::codemap::Span;
 use syntax::ext::base::*;
 use syntax::fold::Folder;
-use syntax::parse::token::intern;
-use syntax::print::pprust;
+use syntax::parse::token::InternedString;
+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 }
     }
 }
 
 impl MultiItemModifier for CustomDerive {
     fn expand(&self,
               ecx: &mut ExtCtxt,
               span: Span,
-              meta_item: &ast::MetaItem,
+              _meta_item: &ast::MetaItem,
               item: Annotatable)
               -> Vec<Annotatable> {
         let item = match item {
@@ -47,31 +62,23 @@ 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");
                 return Vec::new()
             }
         }
 
-        let input_span = Span {
-            expn_id: ecx.codemap().record_expansion(ExpnInfo {
-                call_site: span,
-                callee: NameAndSpan {
-                    format: MacroAttribute(intern(&pprust::meta_item_to_string(meta_item))),
-                    span: Some(span),
-                    allow_internal_unstable: true,
-                },
-            }),
-            ..item.span
-        };
-        let input = __internal::new_token_stream(item);
+        // Mark attributes as known, and used.
+        MarkAttrs(&self.attrs).visit_item(&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 +95,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| {
-            ChangeSpan { span: input_span }.fold_item(item)
-        }).map(Annotatable::Item).collect()
+        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: span }.fold_item(item)
+        }).map(Annotatable::Item));
+        res
     }
 }
 

src/libsyntax_ext/proc_macro_registrar.rs

@@ -30,6 +30,7 @@ struct CustomDerive {
     trait_name: InternedString,
     function_name: Ident,
     span: Span,
+    attrs: Vec<InternedString>,
 }
 
 struct CollectCustomDerives<'a> {
@@ -133,7 +134,8 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
         }
 
         // Once we've located the `#[proc_macro_derive]` attribute, verify
-        // that it's of the form `#[proc_macro_derive(Foo)]`
+        // that it's of the form `#[proc_macro_derive(Foo)]` or
+        // `#[proc_macro_derive(Foo, attributes(A, ..))]`
         let list = match attr.meta_item_list() {
             Some(list) => list,
             None => {
@@ -143,38 +145,69 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
                 return
             }
         };
-        if list.len() != 1 {
+        if list.len() != 1 && list.len() != 2 {
             self.handler.span_err(attr.span(),
-                                  "attribute must only have one argument");
+                                  "attribute must have either one or two arguments");
             return
         }
-        let attr = &list[0];
-        let trait_name = match attr.name() {
+        let trait_attr = &list[0];
+        let attributes_attr = list.get(1);
+        let trait_name = match trait_attr.name() {
             Some(name) => name,
             _ => {
-                self.handler.span_err(attr.span(), "not a meta item");
+                self.handler.span_err(trait_attr.span(), "not a meta item");
                 return
             }
         };
-        if !attr.is_word() {
-            self.handler.span_err(attr.span(), "must only be one word");
+        if !trait_attr.is_word() {
+            self.handler.span_err(trait_attr.span(), "must only be one word");
         }
 
         if deriving::is_builtin_trait(&trait_name) {
-            self.handler.span_err(attr.span(),
+            self.handler.span_err(trait_attr.span(),
                                   "cannot override a built-in #[derive] mode");
         }
 
         if self.derives.iter().any(|d| d.trait_name == trait_name) {
-            self.handler.span_err(attr.span(),
+            self.handler.span_err(trait_attr.span(),
                                   "derive mode defined twice in this crate");
         }
 
+        let proc_attrs: Vec<_> = if let Some(attr) = attributes_attr {
+            if !attr.check_name("attributes") {
+                self.handler.span_err(attr.span(), "second argument must be `attributes`")
+            }
+            attr.meta_item_list().unwrap_or_else(|| {
+                self.handler.span_err(attr.span(),
+                                      "attribute must be of form: \
+                                       `attributes(foo, bar)`");
+                &[]
+            }).into_iter().filter_map(|attr| {
+                let name = match attr.name() {
+                    Some(name) => name,
+                    _ => {
+                        self.handler.span_err(attr.span(), "not a meta item");
+                        return None;
+                    },
+                };
+
+                if !attr.is_word() {
+                    self.handler.span_err(attr.span(), "must only be one word");
+                    return None;
+                }
+
+                Some(name)
+            }).collect()
+        } else {
+            Vec::new()
+        };
+
         if self.in_root {
             self.derives.push(CustomDerive {
                 span: item.span,
                 trait_name: trait_name,
                 function_name: item.ident,
+                attrs: proc_attrs,
             });
         } else {
             let msg = "functions tagged with `#[proc_macro_derive]` must \
@@ -208,8 +241,8 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
 //
 //          #[plugin_registrar]
 //          fn registrar(registrar: &mut Registry) {
-//              registrar.register_custom_derive($name_trait1, ::$name1);
-//              registrar.register_custom_derive($name_trait2, ::$name2);
+//              registrar.register_custom_derive($name_trait1, ::$name1, &[]);
+//              registrar.register_custom_derive($name_trait2, ::$name2, &["attribute_name"]);
 //              // ...
 //          }
 //      }
@@ -238,14 +271,18 @@ fn mk_registrar(cx: &mut ExtCtxt,
     let stmts = custom_derives.iter().map(|cd| {
         let path = cx.path_global(cd.span, vec![cd.function_name]);
         let trait_name = cx.expr_str(cd.span, cd.trait_name.clone());
-        (path, trait_name)
-    }).map(|(path, trait_name)| {
+        let attrs = cx.expr_vec_slice(
+            span,
+            cd.attrs.iter().map(|s| cx.expr_str(cd.span, s.clone())).collect::<Vec<_>>()
+        );
+        (path, trait_name, attrs)
+    }).map(|(path, trait_name, attrs)| {
         let registrar = cx.expr_ident(span, registrar);
         let ufcs_path = cx.path(span, vec![proc_macro, __internal, registry,
                                            register_custom_derive]);
         cx.expr_call(span,
                      cx.expr_path(ufcs_path),
-                     vec![registrar, trait_name, cx.expr_path(path)])
+                     vec![registrar, trait_name, cx.expr_path(path), attrs])
     }).map(|expr| {
         cx.stmt_expr(expr)
     }).collect::<Vec<_>>();