Implement `macro_rules!` placeholders and the macro scope map

5 files changed, 127 insertions, 24 deletions

diff --git a/src/librustc_resolve/assign_ids.rs b/src/librustc_resolve/assign_ids.rs
index f3aa78bd4e2..d4465822229 100644
--- a/src/librustc_resolve/assign_ids.rs
+++ b/src/librustc_resolve/assign_ids.rs
@@ -11,20 +11,27 @@
 use Resolver;
 use rustc::session::Session;
 use syntax::ast;
-use syntax::fold::Folder;
+use syntax::ext::mtwt;
+use syntax::fold::{self, Folder};
 use syntax::ptr::P;
 use syntax::util::move_map::MoveMap;
+use syntax::util::small_vector::SmallVector;
+
+use std::collections::HashMap;
+use std::mem;
 
 impl<'a> Resolver<'a> {
     pub fn assign_node_ids(&mut self, krate: ast::Crate) -> ast::Crate {
         NodeIdAssigner {
             sess: self.session,
+            macros_at_scope: &mut self.macros_at_scope,
         }.fold_crate(krate)
     }
 }
 
 struct NodeIdAssigner<'a> {
     sess: &'a Session,
+    macros_at_scope: &'a mut HashMap<ast::NodeId, Vec<ast::Mrk>>,
 }
 
 impl<'a> Folder for NodeIdAssigner<'a> {
@@ -38,22 +45,48 @@ impl<'a> Folder for NodeIdAssigner<'a> {
             block.id = self.new_id(block.id);
 
             let stmt = block.stmts.pop();
-            block.stmts = block.stmts.move_flat_map(|s| self.fold_stmt(s).into_iter());
-            if let Some(ast::Stmt { node: ast::StmtKind::Expr(expr), span, .. }) = stmt {
+            let mut macros = Vec::new();
+            block.stmts = block.stmts.move_flat_map(|stmt| {
+                if let ast::StmtKind::Item(ref item) = stmt.node {
+                    if let ast::ItemKind::Mac(..) = item.node {
+                        macros.push(mtwt::outer_mark(item.ident.ctxt));
+                        return None;
+                    }
+                }
+
+                let stmt = self.fold_stmt(stmt).pop().unwrap();
+                if !macros.is_empty() {
+                    self.macros_at_scope.insert(stmt.id, mem::replace(&mut macros, Vec::new()));
+                }
+                Some(stmt)
+            });
+
+            stmt.and_then(|mut stmt| {
                 // Avoid wasting a node id on a trailing expression statement,
                 // which shares a HIR node with the expression itself.
-                let expr = self.fold_expr(expr);
-                block.stmts.push(ast::Stmt {
-                    id: expr.id,
-                    node: ast::StmtKind::Expr(expr),
-                    span: span,
-                });
-            } else if let Some(stmt) = stmt {
-                block.stmts.extend(self.fold_stmt(stmt));
-            }
+                if let ast::StmtKind::Expr(expr) = stmt.node {
+                    let expr = self.fold_expr(expr);
+                    stmt.id = expr.id;
+                    stmt.node = ast::StmtKind::Expr(expr);
+                    Some(stmt)
+                } else {
+                    self.fold_stmt(stmt).pop()
+                }
+            }).map(|stmt| {
+                if !macros.is_empty() {
+                    self.macros_at_scope.insert(stmt.id, mem::replace(&mut macros, Vec::new()));
+                }
+                block.stmts.push(stmt);
+            });
 
             block
         })
     }
-}
 
+    fn fold_item(&mut self, item: P<ast::Item>) -> SmallVector<P<ast::Item>> {
+        match item.node {
+            ast::ItemKind::Mac(..) => SmallVector::zero(),
+            _ => fold::noop_fold_item(item, self),
+        }
+    }
+}
diff --git a/src/librustc_resolve/lib.rs b/src/librustc_resolve/lib.rs
index f62eeb41a98..d94a22eb94d 100644
--- a/src/librustc_resolve/lib.rs
+++ b/src/librustc_resolve/lib.rs
@@ -53,6 +53,7 @@ use rustc::ty::subst::{ParamSpace, FnSpace, TypeSpace};
 use rustc::hir::{Freevar, FreevarMap, TraitCandidate, TraitMap, GlobMap};
 use rustc::util::nodemap::{NodeMap, NodeSet, FnvHashMap, FnvHashSet};
 
+use syntax::ext::mtwt;
 use syntax::ast::{self, FloatTy};
 use syntax::ast::{CRATE_NODE_ID, Name, NodeId, CrateNum, IntTy, UintTy};
 use syntax::parse::token::{self, keywords};
@@ -651,6 +652,9 @@ enum RibKind<'a> {
 
     // We passed through a module.
     ModuleRibKind(Module<'a>),
+
+    // We passed through a `macro_rules!` statement with the given expansion
+    MacroDefinition(ast::Mrk),
 }
 
 #[derive(Copy, Clone)]
@@ -927,6 +931,10 @@ pub struct Resolver<'a> {
 
     pub definitions: Definitions,
 
+    // Maps the node id of a statement to the expansions of the `macro_rules!`s
+    // immediately above the statement (if appropriate).
+    macros_at_scope: HashMap<NodeId, Vec<ast::Mrk>>,
+
     graph_root: Module<'a>,
 
     prelude: Option<Module<'a>>,
@@ -1113,6 +1121,7 @@ impl<'a> Resolver<'a> {
             session: session,
 
             definitions: Definitions::new(),
+            macros_at_scope: HashMap::new(),
 
             // The outermost module has def ID 0; this is not reflected in the
             // AST.
@@ -1421,6 +1430,16 @@ impl<'a> Resolver<'a> {
                     };
                 }
             }
+
+            if let MacroDefinition(mac) = self.get_ribs(ns)[i].kind {
+                // If an invocation of this macro created `ident`, give up on `ident`
+                // and switch to `ident`'s source from the macro definition.
+                if let Some((source_ident, source_macro)) = mtwt::source(ident) {
+                    if mac == source_macro {
+                        ident = source_ident;
+                    }
+                }
+            }
         }
 
         None
@@ -2069,6 +2088,7 @@ impl<'a> Resolver<'a> {
         let orig_module = self.current_module;
         let anonymous_module = self.module_map.get(&block.id).cloned(); // clones a reference
 
+        let mut num_value_ribs = 1;
         if let Some(anonymous_module) = anonymous_module {
             debug!("(resolving block) found anonymous module, moving down");
             self.value_ribs.push(Rib::new(ModuleRibKind(anonymous_module)));
@@ -2079,11 +2099,22 @@ impl<'a> Resolver<'a> {
         }
 
         // Descend into the block.
-        visit::walk_block(self, block);
+        for stmt in &block.stmts {
+            if let Some(marks) = self.macros_at_scope.remove(&stmt.id) {
+                num_value_ribs += marks.len() as u32;
+                for mark in marks {
+                    self.value_ribs.push(Rib::new(MacroDefinition(mark)));
+                }
+            }
+
+            self.visit_stmt(stmt);
+        }
 
         // Move back up.
         self.current_module = orig_module;
-        self.value_ribs.pop();
+        for _ in 0 .. num_value_ribs {
+            self.value_ribs.pop();
+        }
         if let Some(_) = anonymous_module {
             self.type_ribs.pop();
         }
@@ -2497,7 +2528,7 @@ impl<'a> Resolver<'a> {
             Def::Local(_, node_id) => {
                 for rib in ribs {
                     match rib.kind {
-                        NormalRibKind | ModuleRibKind(..) => {
+                        NormalRibKind | ModuleRibKind(..) | MacroDefinition(..) => {
                             // Nothing to do. Continue.
                         }
                         ClosureRibKind(function_id) => {
@@ -2546,7 +2577,7 @@ impl<'a> Resolver<'a> {
                 for rib in ribs {
                     match rib.kind {
                         NormalRibKind | MethodRibKind(_) | ClosureRibKind(..) |
-                        ModuleRibKind(..) => {
+                        ModuleRibKind(..) | MacroDefinition(..) => {
                             // Nothing to do. Continue.
                         }
                         ItemRibKind => {
diff --git a/src/libsyntax/ext/expand.rs b/src/libsyntax/ext/expand.rs
index 5b5a5d0e531..10b66d08955 100644
--- a/src/libsyntax/ext/expand.rs
+++ b/src/libsyntax/ext/expand.rs
@@ -15,7 +15,7 @@ use attr::HasAttrs;
 use ext::mtwt;
 use attr;
 use attr::AttrMetaMethods;
-use codemap::{Spanned, ExpnInfo, NameAndSpan, MacroBang, MacroAttribute};
+use codemap::{dummy_spanned, Spanned, ExpnInfo, NameAndSpan, MacroBang, MacroAttribute};
 use syntax_pos::{self, Span, ExpnId};
 use config::StripUnconfigured;
 use ext::base::*;
@@ -105,6 +105,23 @@ pub fn expand_expr(expr: ast::Expr, fld: &mut MacroExpander) -> P<ast::Expr> {
     }
 }
 
+struct MacroScopePlaceholder;
+impl MacResult for MacroScopePlaceholder {
+    fn make_items(self: Box<Self>) -> Option<SmallVector<P<ast::Item>>> {
+        Some(SmallVector::one(P(ast::Item {
+            ident: keywords::Invalid.ident(),
+            attrs: Vec::new(),
+            id: ast::DUMMY_NODE_ID,
+            node: ast::ItemKind::Mac(dummy_spanned(ast::Mac_ {
+                path: ast::Path { span: syntax_pos::DUMMY_SP, global: false, segments: Vec::new() },
+                tts: Vec::new(),
+            })),
+            vis: ast::Visibility::Inherited,
+            span: syntax_pos::DUMMY_SP,
+        })))
+    }
+}
+
 /// Expand a macro invocation. Returns the result of expansion.
 fn expand_mac_invoc<T>(mac: ast::Mac, ident: Option<Ident>, attrs: Vec<ast::Attribute>, span: Span,
                        fld: &mut MacroExpander) -> T
@@ -143,6 +160,7 @@ fn expand_mac_invoc<T>(mac: ast::Mac, ident: Option<Ident>, attrs: Vec<ast::Attr
         };
 
         let ident = ident.unwrap_or(keywords::Invalid.ident());
+        let marked_tts = mark_tts(&tts, mark);
         match *extension {
             NormalTT(ref expandfun, exp_span, allow_internal_unstable) => {
                 if ident.name != keywords::Invalid.name() {
@@ -161,7 +179,6 @@ fn expand_mac_invoc<T>(mac: ast::Mac, ident: Option<Ident>, attrs: Vec<ast::Attr
                     },
                 });
 
-                let marked_tts = mark_tts(&tts, mark);
                 Some(expandfun.expand(fld.cx, call_site, &marked_tts))
             }
 
@@ -181,7 +198,6 @@ fn expand_mac_invoc<T>(mac: ast::Mac, ident: Option<Ident>, attrs: Vec<ast::Attr
                     }
                 });
 
-                let marked_tts = mark_tts(&tts, mark);
                 Some(expander.expand(fld.cx, call_site, ident, marked_tts))
             }
 
@@ -210,15 +226,14 @@ fn expand_mac_invoc<T>(mac: ast::Mac, ident: Option<Ident>, attrs: Vec<ast::Attr
                     span: call_site,
                     imported_from: None,
                     use_locally: true,
-                    body: tts,
+                    body: marked_tts,
                     export: attr::contains_name(&attrs, "macro_export"),
                     allow_internal_unstable: attr::contains_name(&attrs, "allow_internal_unstable"),
                     attrs: attrs,
                 });
 
                 // macro_rules! has a side effect but expands to nothing.
-                fld.cx.bt_pop();
-                None
+                Some(Box::new(MacroScopePlaceholder))
             }
 
             MultiDecorator(..) | MultiModifier(..) => {
@@ -343,6 +358,12 @@ fn expand_multi_modified(a: Annotatable, fld: &mut MacroExpander) -> SmallVector
     match a {
         Annotatable::Item(it) => match it.node {
             ast::ItemKind::Mac(..) => {
+                if match it.node {
+                    ItemKind::Mac(ref mac) => mac.node.path.segments.is_empty(),
+                    _ => unreachable!(),
+                } {
+                    return SmallVector::one(Annotatable::Item(it));
+                }
                 it.and_then(|it| match it.node {
                     ItemKind::Mac(mac) =>
                         expand_mac_invoc(mac, Some(it.ident), it.attrs, it.span, fld),
diff --git a/src/libsyntax/ext/mtwt.rs b/src/libsyntax/ext/mtwt.rs
index a4c698a9226..ac5bac58d2a 100644
--- a/src/libsyntax/ext/mtwt.rs
+++ b/src/libsyntax/ext/mtwt.rs
@@ -17,7 +17,7 @@
 
 pub use self::SyntaxContext_::*;
 
-use ast::{Mrk, SyntaxContext};
+use ast::{Ident, Mrk, SyntaxContext};
 
 use std::cell::RefCell;
 use std::collections::HashMap;
@@ -112,6 +112,20 @@ pub fn outer_mark(ctxt: SyntaxContext) -> Mrk {
     })
 }
 
+/// If `ident` is macro expanded, return the source ident from the macro definition
+/// and the mark of the expansion that created the macro definition.
+pub fn source(ident: Ident) -> Option<(Ident /* source ident */, Mrk /* source macro */)> {
+    with_sctable(|sctable| {
+        let ctxts = sctable.table.borrow();
+        if let Mark(_expansion_mark, macro_ctxt) = ctxts[ident.ctxt.0 as usize] {
+            if let Mark(definition_mark, orig_ctxt) = ctxts[macro_ctxt.0 as usize] {
+                return Some((Ident::new(ident.name, orig_ctxt), definition_mark));
+            }
+        }
+        None
+    })
+}
+
 #[cfg(test)]
 mod tests {
     use ast::{EMPTY_CTXT, Ident, Mrk, Name, SyntaxContext};
diff --git a/src/libsyntax/test.rs b/src/libsyntax/test.rs
index 0a60b7fd430..327696e87b0 100644
--- a/src/libsyntax/test.rs
+++ b/src/libsyntax/test.rs
@@ -185,6 +185,8 @@ impl<'a> fold::Folder for TestHarnessGenerator<'a> {
 
         mod_folded
     }
+
+    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac { mac }
 }
 
 struct EntryPointCleaner {
@@ -234,6 +236,8 @@ impl fold::Folder for EntryPointCleaner {
 
         SmallVector::one(folded)
     }
+
+    fn fold_mac(&mut self, mac: ast::Mac) -> ast::Mac { mac }
 }
 
 fn mk_reexport_mod(cx: &mut TestCtxt, tests: Vec<ast::Ident>,