Rollup merge of #37614 - keeperofdakeys:proc_macro, r=jseyfried

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

This PR allows proc macro functions to declare attribute names that should be marked as used when attached to the deriving item. There are a few questions for this PR.

- Currently this uses a separate attribute named `#[proc_macro_attributes(..)]`, is this the best choice?
- In order to make this work, the `check_attribute` function had to be modified to not error on attributes marked as used. This is a pretty large change in semantics, is there a better way to do this?
- I've got a few clones where I don't know if I need them (like turning `item` into a `TokenStream`), can these be avoided?
- Is switching to `MultiItemDecorator` the right thing here?

Also fixes https://github.com/rust-lang/rust/issues/37563.

2 files changed, 87 insertions, 42 deletions

diff --git a/src/libsyntax_ext/deriving/custom.rs b/src/libsyntax_ext/deriving/custom.rs
index f8cb1294a66..3305c1eae2b 100644
--- a/src/libsyntax_ext/deriving/custom.rs
+++ b/src/libsyntax_ext/deriving/custom.rs
@@ -12,20 +12,35 @@ 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 }
     }
 }
 
@@ -33,7 +48,7 @@ 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,7 +62,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,23 +70,15 @@ impl MultiItemModifier for CustomDerive {
             }
         }
 
-        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
     }
 }
 
diff --git a/src/libsyntax_ext/proc_macro_registrar.rs b/src/libsyntax_ext/proc_macro_registrar.rs
index a8accd63dcf..d6d31200a99 100644
--- a/src/libsyntax_ext/proc_macro_registrar.rs
+++ b/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> {
@@ -144,7 +145,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 => {
@@ -154,38 +156,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 \
@@ -219,8 +252,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"]);
 //              // ...
 //          }
 //      }
@@ -249,14 +282,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<_>>();