Implement `#[proc_macro_attribute]`

* Add support for `#[proc_macro]`

* Reactivate `proc_macro` feature and gate `#[proc_macro_attribute]` under it

* Have `#![feature(proc_macro)]` imply `#![feature(use_extern_macros)]`,
error on legacy import of proc macros via `#[macro_use]`

3 files changed, 195 insertions, 69 deletions

diff --git a/src/libsyntax_ext/lib.rs b/src/libsyntax_ext/lib.rs
index bdec86158a4..ebec23d0901 100644
--- a/src/libsyntax_ext/lib.rs
+++ b/src/libsyntax_ext/lib.rs
@@ -47,6 +47,8 @@ pub mod proc_macro_registrar;
 // for custom_derive
 pub mod deriving;
 
+pub mod proc_macro_impl;
+
 use std::rc::Rc;
 use syntax::ast;
 use syntax::ext::base::{MacroExpanderFn, NormalTT, MultiModifier, NamedSyntaxExtension};
diff --git a/src/libsyntax_ext/proc_macro_impl.rs b/src/libsyntax_ext/proc_macro_impl.rs
new file mode 100644
index 00000000000..b454628acb1
--- /dev/null
+++ b/src/libsyntax_ext/proc_macro_impl.rs
@@ -0,0 +1,58 @@
+// 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.
+
+use std::panic;
+
+use errors::FatalError;
+
+use syntax::codemap::Span;
+use syntax::ext::base::*;
+use syntax::tokenstream::TokenStream;
+use syntax::ext::base;
+
+use proc_macro::TokenStream as TsShim;
+use proc_macro::__internal;
+
+pub struct AttrProcMacro {
+    pub inner: fn(TsShim, TsShim) -> TsShim,
+}
+
+impl base::AttrProcMacro for AttrProcMacro {
+    fn expand<'cx>(&self,
+                   ecx: &'cx mut ExtCtxt,
+                   span: Span,
+                   annotation: TokenStream,
+                   annotated: TokenStream)
+                   -> TokenStream {
+        let annotation = __internal::token_stream_wrap(annotation);
+        let annotated = __internal::token_stream_wrap(annotated);
+
+        let res = __internal::set_parse_sess(&ecx.parse_sess, || {
+            panic::catch_unwind(panic::AssertUnwindSafe(|| (self.inner)(annotation, annotated)))
+        });
+
+        match res {
+            Ok(stream) => __internal::token_stream_inner(stream),
+            Err(e) => {
+                let msg = "custom attribute panicked";
+                let mut err = ecx.struct_span_fatal(span, msg);
+                if let Some(s) = e.downcast_ref::<String>() {
+                    err.help(&format!("message: {}", s));
+                }
+                if let Some(s) = e.downcast_ref::<&'static str>() {
+                    err.help(&format!("message: {}", s));
+                }
+
+                err.emit();
+                panic!(FatalError);
+            }
+        }
+    }
+}
diff --git a/src/libsyntax_ext/proc_macro_registrar.rs b/src/libsyntax_ext/proc_macro_registrar.rs
index c93e2c054d2..c8af16e9242 100644
--- a/src/libsyntax_ext/proc_macro_registrar.rs
+++ b/src/libsyntax_ext/proc_macro_registrar.rs
@@ -11,18 +11,20 @@
 use std::mem;
 
 use errors;
+
 use syntax::ast::{self, Ident, NodeId};
 use syntax::codemap::{ExpnInfo, NameAndSpan, MacroAttribute};
 use syntax::ext::base::ExtCtxt;
 use syntax::ext::build::AstBuilder;
 use syntax::ext::expand::ExpansionConfig;
-use syntax::parse::ParseSess;
 use syntax::fold::Folder;
+use syntax::parse::ParseSess;
 use syntax::ptr::P;
 use syntax::symbol::Symbol;
-use syntax_pos::{Span, DUMMY_SP};
 use syntax::visit::{self, Visitor};
 
+use syntax_pos::{Span, DUMMY_SP};
+
 use deriving;
 
 struct CustomDerive {
@@ -32,8 +34,14 @@ struct CustomDerive {
     attrs: Vec<ast::Name>,
 }
 
-struct CollectCustomDerives<'a> {
+struct AttrProcMacro {
+    function_name: Ident,
+    span: Span,
+}
+
+struct CollectProcMacros<'a> {
     derives: Vec<CustomDerive>,
+    attr_macros: Vec<AttrProcMacro>,
     in_root: bool,
     handler: &'a errors::Handler,
     is_proc_macro_crate: bool,
@@ -50,16 +58,17 @@ pub fn modify(sess: &ParseSess,
     let ecfg = ExpansionConfig::default("proc_macro".to_string());
     let mut cx = ExtCtxt::new(sess, ecfg, resolver);
 
-    let derives = {
-        let mut collect = CollectCustomDerives {
+    let (derives, attr_macros) = {
+        let mut collect = CollectProcMacros {
             derives: Vec::new(),
+            attr_macros: Vec::new(),
             in_root: true,
             handler: handler,
             is_proc_macro_crate: is_proc_macro_crate,
             is_test_crate: is_test_crate,
         };
         visit::walk_crate(&mut collect, &krate);
-        collect.derives
+        (collect.derives, collect.attr_macros)
     };
 
     if !is_proc_macro_crate {
@@ -74,7 +83,7 @@ pub fn modify(sess: &ParseSess,
         return krate;
     }
 
-    krate.module.items.push(mk_registrar(&mut cx, &derives));
+    krate.module.items.push(mk_registrar(&mut cx, &derives, &attr_macros));
 
     if krate.exported_macros.len() > 0 {
         handler.err("cannot export macro_rules! macros from a `proc-macro` \
@@ -84,7 +93,7 @@ pub fn modify(sess: &ParseSess,
     return krate
 }
 
-impl<'a> CollectCustomDerives<'a> {
+impl<'a> CollectProcMacros<'a> {
     fn check_not_pub_in_root(&self, vis: &ast::Visibility, sp: Span) {
         if self.is_proc_macro_crate &&
            self.in_root &&
@@ -92,61 +101,11 @@ impl<'a> CollectCustomDerives<'a> {
             self.handler.span_err(sp,
                                   "`proc-macro` crate types cannot \
                                    export any items other than functions \
-                                   tagged with `#[proc_macro_derive]` \
-                                   currently");
+                                   tagged with `#[proc_macro_derive]` currently");
         }
     }
-}
-
-impl<'a> Visitor<'a> for CollectCustomDerives<'a> {
-    fn visit_item(&mut self, item: &'a ast::Item) {
-        let mut attrs = item.attrs.iter().filter(|a| a.check_name("proc_macro_derive"));
-
-        // First up, make sure we're checking a bare function. If we're not then
-        // we're just not interested in this item.
-        //
-        // If we find one, try to locate a `#[proc_macro_derive]` attribute on
-        // it.
-        match item.node {
-            ast::ItemKind::Fn(..) => {}
-            _ => {
-                // Check for invalid use of proc_macro_derive
-                if let Some(attr) = attrs.next() {
-                    self.handler.span_err(attr.span(),
-                                          "the `#[proc_macro_derive]` \
-                                          attribute may only be used \
-                                          on bare functions");
-                    return;
-                }
-                self.check_not_pub_in_root(&item.vis, item.span);
-                return visit::walk_item(self, item)
-            }
-        }
-
-        let attr = match attrs.next() {
-            Some(attr) => attr,
-            None => {
-                self.check_not_pub_in_root(&item.vis, item.span);
-                return visit::walk_item(self, item)
-            }
-        };
-
-        if let Some(a) = attrs.next() {
-            self.handler.span_err(a.span(), "multiple `#[proc_macro_derive]` \
-                                             attributes found");
-        }
-
-        if self.is_test_crate {
-            return;
-        }
-
-        if !self.is_proc_macro_crate {
-            self.handler.span_err(attr.span(),
-                                  "the `#[proc_macro_derive]` attribute is \
-                                   only usable with crates of the `proc-macro` \
-                                   crate type");
-        }
 
+    fn collect_custom_derive(&mut self, item: &'a ast::Item, attr: &'a ast::Attribute) {
         // Once we've located the `#[proc_macro_derive]` attribute, verify
         // that it's of the form `#[proc_macro_derive(Foo)]` or
         // `#[proc_macro_derive(Foo, attributes(A, ..))]`
@@ -232,6 +191,101 @@ impl<'a> Visitor<'a> for CollectCustomDerives<'a> {
             };
             self.handler.span_err(item.span, msg);
         }
+    }
+
+    fn collect_attr_proc_macro(&mut self, item: &'a ast::Item, attr: &'a ast::Attribute) {
+        if let Some(_) = attr.meta_item_list() {
+            self.handler.span_err(attr.span, "`#[proc_macro_attribute]` attribute
+                cannot contain any meta items");
+            return;
+        }
+
+        if self.in_root && item.vis == ast::Visibility::Public {
+            self.attr_macros.push(AttrProcMacro {
+                span: item.span,
+                function_name: item.ident,
+            });
+        } else {
+            let msg = if !self.in_root {
+                "functions tagged with `#[proc_macro_attribute]` must \
+                 currently reside in the root of the crate"
+            } else {
+                "functions tagged with `#[proc_macro_attribute]` must be `pub`"
+            };
+            self.handler.span_err(item.span, msg);
+        }
+    }
+}
+
+impl<'a> Visitor<'a> for CollectProcMacros<'a> {
+    fn visit_item(&mut self, item: &'a ast::Item) {
+        // First up, make sure we're checking a bare function. If we're not then
+        // we're just not interested in this item.
+        //
+        // If we find one, try to locate a `#[proc_macro_derive]` attribute on
+        // it.
+        let is_fn = match item.node {
+            ast::ItemKind::Fn(..) => true,
+            _ => false,
+        };
+
+        let mut found_attr: Option<&'a ast::Attribute> = None;
+
+        for attr in &item.attrs {
+            if attr.check_name("proc_macro_derive") || attr.check_name("proc_macro_attribute") {
+                if let Some(prev_attr) = found_attr {
+                    let msg = if attr.name() == prev_attr.name() {
+                        format!("Only one `#[{}]` attribute is allowed on any given function",
+                                attr.name())
+                    } else {
+                        format!("`#[{}]` and `#[{}]` attributes cannot both be applied \
+                                to the same function", attr.name(), prev_attr.name())
+                    };
+
+                    self.handler.struct_span_err(attr.span(), &msg)
+                        .span_note(prev_attr.span(), "Previous attribute here")
+                        .emit();
+
+                    return;
+                }
+
+                found_attr = Some(attr);
+            }
+        }
+
+        let attr = match found_attr {
+            None => {
+                self.check_not_pub_in_root(&item.vis, item.span);
+                return visit::walk_item(self, item);
+            },
+            Some(attr) => attr,
+        };
+
+        if !is_fn {
+            let msg = format!("the `#[{}]` attribute may only be used on bare functions",
+                              attr.name());
+
+            self.handler.span_err(attr.span(), &msg);
+            return;
+        }
+
+        if self.is_test_crate {
+            return;
+        }
+
+        if !self.is_proc_macro_crate {
+            let msg = format!("the `#[{}]` attribute is only usable with crates of the \
+                              `proc-macro` crate type", attr.name());
+
+            self.handler.span_err(attr.span(), &msg);
+            return;
+        }
+
+        if attr.check_name("proc_macro_derive") {
+            self.collect_custom_derive(item, attr);
+        } else if attr.check_name("proc_macro_attribute") {
+            self.collect_attr_proc_macro(item, attr);
+        };
 
         visit::walk_item(self, item);
     }
@@ -265,7 +319,8 @@ impl<'a> Visitor<'a> for CollectCustomDerives<'a> {
 //          }
 //      }
 fn mk_registrar(cx: &mut ExtCtxt,
-                custom_derives: &[CustomDerive]) -> P<ast::Item> {
+                custom_derives: &[CustomDerive],
+                custom_attrs: &[AttrProcMacro]) -> P<ast::Item> {
     let eid = cx.codemap().record_expansion(ExpnInfo {
         call_site: DUMMY_SP,
         callee: NameAndSpan {
@@ -286,25 +341,36 @@ fn mk_registrar(cx: &mut ExtCtxt,
     let registry = Ident::from_str("Registry");
     let registrar = Ident::from_str("registrar");
     let register_custom_derive = Ident::from_str("register_custom_derive");
-    let stmts = custom_derives.iter().map(|cd| {
+    let register_attr_proc_macro = Ident::from_str("register_attr_proc_macro");
+
+    let mut 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);
         let attrs = cx.expr_vec_slice(
             span,
             cd.attrs.iter().map(|&s| cx.expr_str(cd.span, s)).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), attrs])
-    }).map(|expr| {
-        cx.stmt_expr(expr)
+
+        cx.stmt_expr(cx.expr_call(span, cx.expr_path(ufcs_path),
+                                  vec![registrar, trait_name, cx.expr_path(path), attrs]))
+
     }).collect::<Vec<_>>();
 
+    stmts.extend(custom_attrs.iter().map(|ca| {
+        let name = cx.expr_str(ca.span, ca.function_name.name);
+        let path = cx.path_global(ca.span, vec![ca.function_name]);
+        let registrar = cx.expr_ident(ca.span, registrar);
+
+        let ufcs_path = cx.path(span,
+                                vec![proc_macro, __internal, registry, register_attr_proc_macro]);
+
+        cx.stmt_expr(cx.expr_call(span, cx.expr_path(ufcs_path),
+                                  vec![registrar, name, cx.expr_path(path)]))
+    }));
+
     let path = cx.path(span, vec![proc_macro, __internal, registry]);
     let registrar_path = cx.ty_path(path);
     let arg_ty = cx.ty_rptr(span, registrar_path, None, ast::Mutability::Mutable);