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.
rust-lang · bors · Nov 9, 2016 · Nov 7, 2016 · Nov 8, 2016 · Nov 8, 2016
commit 31a508e1185713c6e570bb963dd3e097a228957c
diff --git a/src/libproc_macro/lib.rs b/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

diff --git a/src/librustc_metadata/creader.rs b/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));
             }
         }

diff --git 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,14 +64,17 @@ 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()
             }
         }
 
+        // 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
     }
 }
 
diff --git 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> {
@@ -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<_>>();

diff --git a/src/test/compile-fail-fulldeps/proc-macro/attribute.rs b/src/test/compile-fail-fulldeps/proc-macro/attribute.rs
@@ -33,8 +33,8 @@ pub fn foo3(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     input
 }
 
-#[proc_macro_derive(b, c)]
-//~^ ERROR: attribute must only have one argument
+#[proc_macro_derive(b, c, d)]
+//~^ ERROR: attribute must have either one or two arguments
 pub fn foo4(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     input
 }
@@ -44,3 +44,21 @@ pub fn foo4(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
 pub fn foo5(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     input
 }
+
+#[proc_macro_derive(f, attributes(g = "h"))]
+//~^ ERROR: must only be one word
+pub fn foo6(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    input
+}
+
+#[proc_macro_derive(i, attributes(j(k)))]
+//~^ ERROR: must only be one word
+pub fn foo7(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    input
+}
+
+#[proc_macro_derive(l, attributes(m), n)]
+//~^ ERROR: attribute must have either one or two arguments
+pub fn foo8(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
+    input
+}
diff --git a/src/test/compile-fail-fulldeps/proc-macro/auxiliary/derive-b.rs b/src/test/compile-fail-fulldeps/proc-macro/auxiliary/derive-b.rs
@@ -0,0 +1,25 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// force-host
+// no-prefer-dynamic
+
+#![feature(proc_macro)]
+#![feature(proc_macro_lib)]
+#![crate_type = "proc-macro"]
+
+extern crate proc_macro;
+
+use proc_macro::TokenStream;
+
+#[proc_macro_derive(B, attributes(B))]
+pub fn derive_b(input: TokenStream) -> TokenStream {
+    "".parse().unwrap()
+}
diff --git a/src/test/compile-fail-fulldeps/proc-macro/item-error.rs b/src/test/compile-fail-fulldeps/proc-macro/item-error.rs
@@ -0,0 +1,26 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// aux-build:derive-b.rs
+
+#![feature(proc_macro)]
+#![allow(warnings)]
+
+#[macro_use]
+extern crate derive_b;
+
+#[derive(B)]
+struct A {
+    a: &u64
+//~^ ERROR: missing lifetime specifier
+}
+
+fn main() {
+}
diff --git a/src/test/compile-fail-fulldeps/proc-macro/proc-macro-attributes.rs b/src/test/compile-fail-fulldeps/proc-macro/proc-macro-attributes.rs
@@ -0,0 +1,26 @@
+// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// aux-build:derive-b.rs
+
+#![feature(proc_macro)]
+#![allow(warnings)]
+
+#[macro_use]
+extern crate derive_b;
+
+#[derive(B)]
+#[B]
+#[C] //~ ERROR: The attribute `C` is currently unknown to the compiler
+#[B(D)]
+#[B(E = "foo")]
+struct B;
+
+fn main() {}
diff --git a/src/test/run-pass-fulldeps/proc-macro/auxiliary/add-impl.rs b/src/test/run-pass-fulldeps/proc-macro/auxiliary/add-impl.rs
@@ -21,13 +21,12 @@ use proc_macro::TokenStream;
 #[proc_macro_derive(AddImpl)]
 // #[cfg(proc_macro)]
 pub fn derive(input: TokenStream) -> TokenStream {
-    (input.to_string() + "
-        impl B {
+    "impl B {
             fn foo(&self) {}
         }
 
         fn foo() {}
 
         mod bar { pub fn foo() {} }
-    ").parse().unwrap()
+    ".parse().unwrap()
 }
diff --git a/src/test/run-pass-fulldeps/proc-macro/auxiliary/append-impl.rs b/src/test/run-pass-fulldeps/proc-macro/auxiliary/append-impl.rs
@@ -21,11 +21,8 @@ use proc_macro::TokenStream;
 
 #[proc_macro_derive(Append)]
 pub fn derive_a(input: TokenStream) -> TokenStream {
-    let mut input = input.to_string();
-    input.push_str("
-        impl Append for A {
-            fn foo(&self) {}
-        }
-    ");
-    input.parse().unwrap()
+    "impl Append for A {
+         fn foo(&self) {}
+     }
+    ".parse().unwrap()
 }