Implement function-like procedural macros ( `#[proc_macro]`)

2 files changed, 92 insertions, 9 deletions

diff --git a/src/libsyntax_ext/proc_macro_impl.rs b/src/libsyntax_ext/proc_macro_impl.rs
index b454628acb1..f60e5824db9 100644
--- a/src/libsyntax_ext/proc_macro_impl.rs
+++ b/src/libsyntax_ext/proc_macro_impl.rs
@@ -56,3 +56,38 @@ impl base::AttrProcMacro for AttrProcMacro {
         }
     }
 }
+
+pub struct BangProcMacro {
+    pub inner: fn(TsShim) -> TsShim,
+}
+
+impl base::ProcMacro for BangProcMacro {
+    fn expand<'cx>(&self,
+                   ecx: &'cx mut ExtCtxt,
+                   span: Span,
+                   input: TokenStream)
+                   -> TokenStream {
+        let input = __internal::token_stream_wrap(input);
+
+        let res = __internal::set_parse_sess(&ecx.parse_sess, || {
+            panic::catch_unwind(panic::AssertUnwindSafe(|| (self.inner)(input)))
+        });
+
+        match res {
+            Ok(stream) => __internal::token_stream_inner(stream),
+            Err(e) => {
+                let msg = "proc macro 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 325f09a83dd..9c96ad547e1 100644
--- a/src/libsyntax_ext/proc_macro_registrar.rs
+++ b/src/libsyntax_ext/proc_macro_registrar.rs
@@ -27,6 +27,9 @@ use syntax_pos::{Span, DUMMY_SP};
 
 use deriving;
 
+const PROC_MACRO_KINDS: [&'static str; 3] =
+    ["proc_macro_derive", "proc_macro_attribute", "proc_macro"];
+
 struct ProcMacroDerive {
     trait_name: ast::Name,
     function_name: Ident,
@@ -34,14 +37,15 @@ struct ProcMacroDerive {
     attrs: Vec<ast::Name>,
 }
 
-struct AttrProcMacro {
+struct ProcMacroDef {
     function_name: Ident,
     span: Span,
 }
 
 struct CollectProcMacros<'a> {
     derives: Vec<ProcMacroDerive>,
-    attr_macros: Vec<AttrProcMacro>,
+    attr_macros: Vec<ProcMacroDef>,
+    bang_macros: Vec<ProcMacroDef>,
     in_root: bool,
     handler: &'a errors::Handler,
     is_proc_macro_crate: bool,
@@ -58,17 +62,18 @@ pub fn modify(sess: &ParseSess,
     let ecfg = ExpansionConfig::default("proc_macro".to_string());
     let mut cx = ExtCtxt::new(sess, ecfg, resolver);
 
-    let (derives, attr_macros) = {
+    let (derives, attr_macros, bang_macros) = {
         let mut collect = CollectProcMacros {
             derives: Vec::new(),
             attr_macros: Vec::new(),
+            bang_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.attr_macros)
+        (collect.derives, collect.attr_macros, collect.bang_macros)
     };
 
     if !is_proc_macro_crate {
@@ -83,7 +88,7 @@ pub fn modify(sess: &ParseSess,
         return krate;
     }
 
-    krate.module.items.push(mk_registrar(&mut cx, &derives, &attr_macros));
+    krate.module.items.push(mk_registrar(&mut cx, &derives, &attr_macros, &bang_macros));
 
     if krate.exported_macros.len() > 0 {
         handler.err("cannot export macro_rules! macros from a `proc-macro` \
@@ -93,6 +98,10 @@ pub fn modify(sess: &ParseSess,
     return krate
 }
 
+fn is_proc_macro_attr(attr: &ast::Attribute) -> bool {
+    PROC_MACRO_KINDS.iter().any(|kind| attr.check_name(kind))
+}
+
 impl<'a> CollectProcMacros<'a> {
     fn check_not_pub_in_root(&self, vis: &ast::Visibility, sp: Span) {
         if self.is_proc_macro_crate &&
@@ -196,12 +205,12 @@ impl<'a> CollectProcMacros<'a> {
     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");
+                does not take any arguments");
             return;
         }
 
         if self.in_root && item.vis == ast::Visibility::Public {
-            self.attr_macros.push(AttrProcMacro {
+            self.attr_macros.push(ProcMacroDef {
                 span: item.span,
                 function_name: item.ident,
             });
@@ -215,6 +224,29 @@ impl<'a> CollectProcMacros<'a> {
             self.handler.span_err(item.span, msg);
         }
     }
+
+    fn collect_bang_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
+                does not take any arguments");
+            return;
+        }
+
+        if self.in_root && item.vis == ast::Visibility::Public {
+            self.bang_macros.push(ProcMacroDef {
+                span: item.span,
+                function_name: item.ident,
+            });
+        } else {
+            let msg = if !self.in_root {
+                "functions tagged with `#[proc_macro]` must \
+                 currently reside in the root of the crate"
+            } else {
+                "functions tagged with `#[proc_macro]` must be `pub`"
+            };
+            self.handler.span_err(item.span, msg);
+        }
+    }
 }
 
 impl<'a> Visitor<'a> for CollectProcMacros<'a> {
@@ -232,7 +264,7 @@ impl<'a> Visitor<'a> for CollectProcMacros<'a> {
         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 is_proc_macro_attr(&attr) {
                 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",
@@ -285,6 +317,8 @@ impl<'a> Visitor<'a> for CollectProcMacros<'a> {
             self.collect_custom_derive(item, attr);
         } else if attr.check_name("proc_macro_attribute") {
             self.collect_attr_proc_macro(item, attr);
+        } else if attr.check_name("proc_macro") {
+            self.collect_bang_proc_macro(item, attr);
         };
 
         visit::walk_item(self, item);
@@ -320,7 +354,8 @@ impl<'a> Visitor<'a> for CollectProcMacros<'a> {
 //      }
 fn mk_registrar(cx: &mut ExtCtxt,
                 custom_derives: &[ProcMacroDerive],
-                custom_attrs: &[AttrProcMacro]) -> P<ast::Item> {
+                custom_attrs: &[ProcMacroDef],
+                custom_macros: &[ProcMacroDef]) -> P<ast::Item> {
     let eid = cx.codemap().record_expansion(ExpnInfo {
         call_site: DUMMY_SP,
         callee: NameAndSpan {
@@ -342,6 +377,7 @@ fn mk_registrar(cx: &mut ExtCtxt,
     let registrar = Ident::from_str("registrar");
     let register_custom_derive = Ident::from_str("register_custom_derive");
     let register_attr_proc_macro = Ident::from_str("register_attr_proc_macro");
+    let register_bang_proc_macro = Ident::from_str("register_bang_proc_macro");
 
     let mut stmts = custom_derives.iter().map(|cd| {
         let path = cx.path_global(cd.span, vec![cd.function_name]);
@@ -371,6 +407,18 @@ fn mk_registrar(cx: &mut ExtCtxt,
                                   vec![registrar, name, cx.expr_path(path)]))
     }));
 
+    stmts.extend(custom_macros.iter().map(|cm| {
+        let name = cx.expr_str(cm.span, cm.function_name.name);
+        let path = cx.path_global(cm.span, vec![cm.function_name]);
+        let registrar = cx.expr_ident(cm.span, registrar);
+
+        let ufcs_path = cx.path(span,
+                                vec![proc_macro, __internal, registry, register_bang_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);