Skip to content

Commit 31a508e

Browse files
keeperofdakeyskeeperofdakeys
committed
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 parent d377cf5 commit 31a508e

File tree

16 files changed

+260
-46
lines changed

16 files changed

+260
-46
lines changed

‎src/libproc_macro/lib.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,8 @@ pub mod __internal {
9595
pub trait Registry {
9696
fn register_custom_derive(&mut self,
9797
trait_name: &str,
98-
expand: fn(TokenStream) -> TokenStream);
98+
expand: fn(TokenStream) -> TokenStream,
99+
attributes: &[&'static str]);
99100
}
100101

101102
// Emulate scoped_thread_local!() here essentially

‎src/librustc_metadata/creader.rs‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -624,8 +624,12 @@ impl<'a> CrateLoader<'a> {
624624
impl Registry for MyRegistrar {
625625
fn register_custom_derive(&mut self,
626626
trait_name: &str,
627-
expand: fn(TokenStream) -> TokenStream) {
628-
let derive = SyntaxExtension::CustomDerive(Box::new(CustomDerive::new(expand)));
627+
expand: fn(TokenStream) -> TokenStream,
628+
attributes: &[&'static str]) {
629+
let attrs = attributes.iter().map(|s| InternedString::new(s)).collect();
630+
let derive = SyntaxExtension::CustomDerive(
631+
Box::new(CustomDerive::new(expand, attrs))
632+
);
629633
self.0.push((intern(trait_name), derive));
630634
}
631635
}

‎src/libsyntax_ext/deriving/custom.rs‎

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -12,20 +12,37 @@ use std::panic;
1212

1313
use errors::FatalError;
1414
use proc_macro::{TokenStream, __internal};
15-
use syntax::ast::{self, ItemKind};
15+
use syntax::ast::{self, ItemKind, Attribute};
16+
use syntax::attr::{mark_used, mark_known};
1617
use syntax::codemap::{ExpnInfo, MacroAttribute, NameAndSpan, Span};
1718
use syntax::ext::base::*;
1819
use syntax::fold::Folder;
20+
use syntax::parse::token::InternedString;
1921
use syntax::parse::token::intern;
2022
use syntax::print::pprust;
23+
use syntax::visit::Visitor;
24+
25+
struct MarkAttrs<'a>(&'a [InternedString]);
26+
27+
impl<'a> Visitor for MarkAttrs<'a> {
28+
fn visit_attribute(&mut self, attr: &Attribute) {
29+
if self.0.contains(&attr.name()) {
30+
mark_used(attr);
31+
mark_known(attr);
32+
}
33+
}
34+
}
2135

2236
pub struct CustomDerive {
2337
inner: fn(TokenStream) -> TokenStream,
38+
attrs: Vec<InternedString>,
2439
}
2540

2641
impl CustomDerive {
27-
pub fn new(inner: fn(TokenStream) -> TokenStream) -> CustomDerive {
28-
CustomDerive { inner: inner }
42+
pub fn new(inner: fn(TokenStream) -> TokenStream,
43+
attrs: Vec<InternedString>)
44+
-> CustomDerive {
45+
CustomDerive { inner: inner, attrs: attrs }
2946
}
3047
}
3148

@@ -47,14 +64,17 @@ impl MultiItemModifier for CustomDerive {
4764
};
4865
match item.node {
4966
ItemKind::Struct(..) |
50-
ItemKind::Enum(..) => {}
67+
ItemKind::Enum(..) => {},
5168
_ => {
5269
ecx.span_err(span, "custom derive attributes may only be \
5370
applied to struct/enum items");
5471
return Vec::new()
5572
}
5673
}
5774

75+
// Mark attributes as known, and used.
76+
MarkAttrs(&self.attrs).visit_item(&item);
77+
5878
let input_span = Span {
5979
expn_id: ecx.codemap().record_expansion(ExpnInfo {
6080
call_site: span,
@@ -66,12 +86,13 @@ impl MultiItemModifier for CustomDerive {
6686
}),
6787
..item.span
6888
};
69-
let input = __internal::new_token_stream(item);
89+
90+
let input = __internal::new_token_stream(item.clone());
7091
let res = __internal::set_parse_sess(&ecx.parse_sess, || {
7192
let inner = self.inner;
7293
panic::catch_unwind(panic::AssertUnwindSafe(|| inner(input)))
7394
});
74-
let item = match res {
95+
let new_items = match res {
7596
Ok(stream) => __internal::token_stream_items(stream),
7697
Err(e) => {
7798
let msg = "custom derive attribute panicked";
@@ -88,12 +109,13 @@ impl MultiItemModifier for CustomDerive {
88109
}
89110
};
90111

91-
// Right now we have no knowledge of spans at all in custom derive
92-
// macros, everything is just parsed as a string. Reassign all spans to
93-
// the input `item` for better errors here.
94-
item.into_iter().flat_map(|item| {
112+
let mut res = vec![Annotatable::Item(item)];
113+
// Reassign spans of all expanded items to the input `item`
114+
// for better errors here.
115+
res.extend(new_items.into_iter().flat_map(|item| {
95116
ChangeSpan { span: input_span }.fold_item(item)
96-
}).map(Annotatable::Item).collect()
117+
}).map(Annotatable::Item));
118+
res
97119
}
98120
}
99121

‎src/libsyntax_ext/proc_macro_registrar.rs‎

Lines changed: 52 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ struct CustomDerive {
3030
trait_name: InternedString,
3131
function_name: Ident,
3232
span: Span,
33+
attrs: Vec<InternedString>,
3334
}
3435

3536
struct CollectCustomDerives<'a> {
@@ -133,7 +134,8 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
133134
}
134135

135136
// Once we've located the `#[proc_macro_derive]` attribute, verify
136-
// that it's of the form `#[proc_macro_derive(Foo)]`
137+
// that it's of the form `#[proc_macro_derive(Foo)]` or
138+
// `#[proc_macro_derive(Foo, attributes(A, ..))]`
137139
let list = match attr.meta_item_list() {
138140
Some(list) => list,
139141
None => {
@@ -143,38 +145,69 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
143145
return
144146
}
145147
};
146-
if list.len() != 1 {
148+
if list.len() != 1 && list.len() != 2 {
147149
self.handler.span_err(attr.span(),
148-
"attribute must only have one argument");
150+
"attribute must have either one or two arguments");
149151
return
150152
}
151-
let attr = &list[0];
152-
let trait_name = match attr.name() {
153+
let trait_attr = &list[0];
154+
let attributes_attr = list.get(1);
155+
let trait_name = match trait_attr.name() {
153156
Some(name) => name,
154157
_ => {
155-
self.handler.span_err(attr.span(), "not a meta item");
158+
self.handler.span_err(trait_attr.span(), "not a meta item");
156159
return
157160
}
158161
};
159-
if !attr.is_word() {
160-
self.handler.span_err(attr.span(), "must only be one word");
162+
if !trait_attr.is_word() {
163+
self.handler.span_err(trait_attr.span(), "must only be one word");
161164
}
162165

163166
if deriving::is_builtin_trait(&trait_name) {
164-
self.handler.span_err(attr.span(),
167+
self.handler.span_err(trait_attr.span(),
165168
"cannot override a built-in #[derive] mode");
166169
}
167170

168171
if self.derives.iter().any(|d| d.trait_name == trait_name) {
169-
self.handler.span_err(attr.span(),
172+
self.handler.span_err(trait_attr.span(),
170173
"derive mode defined twice in this crate");
171174
}
172175

176+
let proc_attrs: Vec<_> = if let Some(attr) = attributes_attr {
177+
if !attr.check_name("attributes") {
178+
self.handler.span_err(attr.span(), "second argument must be `attributes`")
179+
}
180+
attr.meta_item_list().unwrap_or_else(|| {
181+
self.handler.span_err(attr.span(),
182+
"attribute must be of form: \
183+
`attributes(foo, bar)`");
184+
&[]
185+
}).into_iter().filter_map(|attr| {
186+
let name = match attr.name() {
187+
Some(name) => name,
188+
_ => {
189+
self.handler.span_err(attr.span(), "not a meta item");
190+
return None;
191+
},
192+
};
193+
194+
if !attr.is_word() {
195+
self.handler.span_err(attr.span(), "must only be one word");
196+
return None;
197+
}
198+
199+
Some(name)
200+
}).collect()
201+
} else {
202+
Vec::new()
203+
};
204+
173205
if self.in_root {
174206
self.derives.push(CustomDerive {
175207
span: item.span,
176208
trait_name: trait_name,
177209
function_name: item.ident,
210+
attrs: proc_attrs,
178211
});
179212
} else {
180213
let msg = "functions tagged with `#[proc_macro_derive]` must \
@@ -208,8 +241,8 @@ impl<'a> Visitor for CollectCustomDerives<'a> {
208241
//
209242
// #[plugin_registrar]
210243
// fn registrar(registrar: &mut Registry) {
211-
// registrar.register_custom_derive($name_trait1, ::$name1);
212-
// registrar.register_custom_derive($name_trait2, ::$name2);
244+
// registrar.register_custom_derive($name_trait1, ::$name1, &[]);
245+
// registrar.register_custom_derive($name_trait2, ::$name2, &["attribute_name"]);
213246
// // ...
214247
// }
215248
// }
@@ -238,14 +271,18 @@ fn mk_registrar(cx: &mut ExtCtxt,
238271
let stmts = custom_derives.iter().map(|cd| {
239272
let path = cx.path_global(cd.span, vec![cd.function_name]);
240273
let trait_name = cx.expr_str(cd.span, cd.trait_name.clone());
241-
(path, trait_name)
242-
}).map(|(path, trait_name)| {
274+
let attrs = cx.expr_vec_slice(
275+
span,
276+
cd.attrs.iter().map(|s| cx.expr_str(cd.span, s.clone())).collect::<Vec<_>>()
277+
);
278+
(path, trait_name, attrs)
279+
}).map(|(path, trait_name, attrs)| {
243280
let registrar = cx.expr_ident(span, registrar);
244281
let ufcs_path = cx.path(span, vec![proc_macro, __internal, registry,
245282
register_custom_derive]);
246283
cx.expr_call(span,
247284
cx.expr_path(ufcs_path),
248-
vec![registrar, trait_name, cx.expr_path(path)])
285+
vec![registrar, trait_name, cx.expr_path(path), attrs])
249286
}).map(|expr| {
250287
cx.stmt_expr(expr)
251288
}).collect::<Vec<_>>();

‎src/test/compile-fail-fulldeps/proc-macro/attribute.rs‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,8 +33,8 @@ pub fn foo3(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
3333
input
3434
}
3535

36-
#[proc_macro_derive(b, c)]
37-
//~^ ERROR: attribute must only have one argument
36+
#[proc_macro_derive(b, c, d)]
37+
//~^ ERROR: attribute must have either one or two arguments
3838
pub fn foo4(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
3939
input
4040
}
@@ -44,3 +44,21 @@ pub fn foo4(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
4444
pub fn foo5(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
4545
input
4646
}
47+
48+
#[proc_macro_derive(f, attributes(g = "h"))]
49+
//~^ ERROR: must only be one word
50+
pub fn foo6(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
51+
input
52+
}
53+
54+
#[proc_macro_derive(i, attributes(j(k)))]
55+
//~^ ERROR: must only be one word
56+
pub fn foo7(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
57+
input
58+
}
59+
60+
#[proc_macro_derive(l, attributes(m), n)]
61+
//~^ ERROR: attribute must have either one or two arguments
62+
pub fn foo8(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
63+
input
64+
}

‎src/test/compile-fail-fulldeps/proc-macro/auxiliary/derive-b.rs‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
2+
// file at the top-level directory of this distribution and at
3+
// http://rust-lang.org/COPYRIGHT.
4+
//
5+
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6+
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7+
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8+
// option. This file may not be copied, modified, or distributed
9+
// except according to those terms.
10+
11+
// force-host
12+
// no-prefer-dynamic
13+
14+
#![feature(proc_macro)]
15+
#![feature(proc_macro_lib)]
16+
#![crate_type = "proc-macro"]
17+
18+
extern crate proc_macro;
19+
20+
use proc_macro::TokenStream;
21+
22+
#[proc_macro_derive(B, attributes(B))]
23+
pub fn derive_b(input: TokenStream) -> TokenStream {
24+
"".parse().unwrap()
25+
}

‎src/test/compile-fail-fulldeps/proc-macro/item-error.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
2+
// file at the top-level directory of this distribution and at
3+
// http://rust-lang.org/COPYRIGHT.
4+
//
5+
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6+
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7+
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8+
// option. This file may not be copied, modified, or distributed
9+
// except according to those terms.
10+
11+
// aux-build:derive-b.rs
12+
13+
#![feature(proc_macro)]
14+
#![allow(warnings)]
15+
16+
#[macro_use]
17+
extern crate derive_b;
18+
19+
#[derive(B)]
20+
struct A {
21+
a: &u64
22+
//~^ ERROR: missing lifetime specifier
23+
}
24+
25+
fn main() {
26+
}

‎src/test/compile-fail-fulldeps/proc-macro/proc-macro-attributes.rs‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// Copyright 2016 The Rust Project Developers. See the COPYRIGHT
2+
// file at the top-level directory of this distribution and at
3+
// http://rust-lang.org/COPYRIGHT.
4+
//
5+
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
6+
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
7+
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
8+
// option. This file may not be copied, modified, or distributed
9+
// except according to those terms.
10+
11+
// aux-build:derive-b.rs
12+
13+
#![feature(proc_macro)]
14+
#![allow(warnings)]
15+
16+
#[macro_use]
17+
extern crate derive_b;
18+
19+
#[derive(B)]
20+
#[B]
21+
#[C] //~ ERROR: The attribute `C` is currently unknown to the compiler
22+
#[B(D)]
23+
#[B(E = "foo")]
24+
struct B;
25+
26+
fn main() {}

‎src/test/run-pass-fulldeps/proc-macro/auxiliary/add-impl.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,13 +21,12 @@ use proc_macro::TokenStream;
2121
#[proc_macro_derive(AddImpl)]
2222
// #[cfg(proc_macro)]
2323
pub fn derive(input: TokenStream) -> TokenStream {
24-
(input.to_string() + "
25-
impl B {
24+
"impl B {
2625
fn foo(&self) {}
2726
}
2827
2928
fn foo() {}
3029
3130
mod bar { pub fn foo() {} }
32-
").parse().unwrap()
31+
".parse().unwrap()
3332
}

‎src/test/run-pass-fulldeps/proc-macro/auxiliary/append-impl.rs‎

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,8 @@ use proc_macro::TokenStream;
2121

2222
#[proc_macro_derive(Append)]
2323
pub fn derive_a(input: TokenStream) -> TokenStream {
24-
let mut input = input.to_string();
25-
input.push_str("
26-
impl Append for A {
27-
fn foo(&self) {}
28-
}
29-
");
30-
input.parse().unwrap()
24+
"impl Append for A {
25+
fn foo(&self) {}
26+
}
27+
".parse().unwrap()
3128
}

0 commit comments

Comments
 (0)