Perform node id assignment and `macros_at_scope` construction during

the `InvocationCollector` and `PlaceholderExpander` folds.

3 files changed, 79 insertions, 9 deletions

diff --git a/src/libsyntax/ext/base.rs b/src/libsyntax/ext/base.rs
index d46e2a9872e..1e8f8ef9ddd 100644
--- a/src/libsyntax/ext/base.rs
+++ b/src/libsyntax/ext/base.rs
@@ -464,6 +464,7 @@ pub type NamedSyntaxExtension = (Name, SyntaxExtension);
 
 pub trait Resolver {
     fn load_crate(&mut self, extern_crate: &ast::Item, allows_macros: bool) -> Vec<LoadedMacro>;
+    fn next_node_id(&mut self) -> ast::NodeId;
 
     fn visit_expansion(&mut self, mark: Mark, expansion: &Expansion);
     fn add_macro(&mut self, scope: Mark, ident: ast::Ident, ext: Rc<SyntaxExtension>);
@@ -479,10 +480,12 @@ pub enum LoadedMacro {
 }
 
 pub struct DummyResolver;
+
 impl Resolver for DummyResolver {
     fn load_crate(&mut self, _extern_crate: &ast::Item, _allows_macros: bool) -> Vec<LoadedMacro> {
         Vec::new()
     }
+    fn next_node_id(&mut self) -> ast::NodeId { ast::DUMMY_NODE_ID }
 
     fn visit_expansion(&mut self, _invoc: Mark, _expansion: &Expansion) {}
     fn add_macro(&mut self, _scope: Mark, _ident: ast::Ident, _ext: Rc<SyntaxExtension>) {}
diff --git a/src/libsyntax/ext/expand.rs b/src/libsyntax/ext/expand.rs
index beb1687dfdc..eef38ea28e0 100644
--- a/src/libsyntax/ext/expand.rs
+++ b/src/libsyntax/ext/expand.rs
@@ -195,6 +195,10 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
         krate.module.items = self.expand(items).make_items().into();
         krate.exported_macros = mem::replace(&mut self.cx.exported_macros, Vec::new());
 
+        for def in &mut krate.exported_macros {
+            def.id = self.cx.resolver.next_node_id()
+        }
+
         if self.cx.parse_sess.span_diagnostic.err_count() > err_count {
             self.cx.parse_sess.span_diagnostic.abort_if_errors();
         }
@@ -234,7 +238,7 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
 
         self.cx.current_expansion = orig_expansion_data;
 
-        let mut placeholder_expander = PlaceholderExpander::new();
+        let mut placeholder_expander = PlaceholderExpander::new(self.cx);
         while let Some(expansions) = expansions.pop() {
             for (mark, expansion) in expansions.into_iter().rev() {
                 let expansion = expansion.fold_with(&mut placeholder_expander);
@@ -530,9 +534,14 @@ impl<'a, 'b> Folder for InvocationCollector<'a, 'b> {
             None => return SmallVector::zero(),
         };
 
-        let (mac, style, attrs) = match stmt.node {
-            StmtKind::Mac(mac) => mac.unwrap(),
-            _ => return noop_fold_stmt(stmt, self),
+        let (mac, style, attrs) = if let StmtKind::Mac(mac) = stmt.node {
+            mac.unwrap()
+        } else {
+            // The placeholder expander gives ids to statements, so we avoid folding the id here.
+            let ast::Stmt { id, node, span } = stmt;
+            return noop_fold_stmt_kind(node, self).into_iter().map(|node| {
+                ast::Stmt { id: id, node: node, span: span }
+            }).collect()
         };
 
         let mut placeholder =
@@ -624,7 +633,7 @@ impl<'a, 'b> Folder for InvocationCollector<'a, 'b> {
                         }
                     }
                 }
-                SmallVector::one(item)
+                noop_fold_item(item, self)
             },
             _ => noop_fold_item(item, self),
         }
@@ -687,6 +696,11 @@ impl<'a, 'b> Folder for InvocationCollector<'a, 'b> {
     fn fold_item_kind(&mut self, item: ast::ItemKind) -> ast::ItemKind {
         noop_fold_item_kind(self.cfg.configure_item_kind(item), self)
     }
+
+    fn new_id(&mut self, id: ast::NodeId) -> ast::NodeId {
+        assert_eq!(id, ast::DUMMY_NODE_ID);
+        self.cx.resolver.next_node_id()
+    }
 }
 
 pub struct ExpansionConfig<'feat> {
diff --git a/src/libsyntax/ext/placeholders.rs b/src/libsyntax/ext/placeholders.rs
index abadcf867b1..7635705daa0 100644
--- a/src/libsyntax/ext/placeholders.rs
+++ b/src/libsyntax/ext/placeholders.rs
@@ -10,13 +10,16 @@
 
 use ast;
 use codemap::{DUMMY_SP, dummy_spanned};
+use ext::base::ExtCtxt;
 use ext::expand::{Expansion, ExpansionKind};
 use fold::*;
 use parse::token::keywords;
 use ptr::P;
+use util::move_map::MoveMap;
 use util::small_vector::SmallVector;
 
 use std::collections::HashMap;
+use std::mem;
 
 pub fn placeholder(kind: ExpansionKind, id: ast::NodeId) -> Expansion {
     fn mac_placeholder() -> ast::Mac {
@@ -69,13 +72,15 @@ pub fn macro_scope_placeholder() -> Expansion {
     placeholder(ExpansionKind::Items, ast::DUMMY_NODE_ID)
 }
 
-pub struct PlaceholderExpander {
+pub struct PlaceholderExpander<'a, 'b: 'a> {
     expansions: HashMap<ast::NodeId, Expansion>,
+    cx: &'a mut ExtCtxt<'b>,
 }
 
-impl PlaceholderExpander {
-    pub fn new() -> Self {
+impl<'a, 'b> PlaceholderExpander<'a, 'b> {
+    pub fn new(cx: &'a mut ExtCtxt<'b>) -> Self {
         PlaceholderExpander {
+            cx: cx,
             expansions: HashMap::new(),
         }
     }
@@ -89,7 +94,7 @@ impl PlaceholderExpander {
     }
 }
 
-impl Folder for PlaceholderExpander {
+impl<'a, 'b> Folder for PlaceholderExpander<'a, 'b> {
     fn fold_item(&mut self, item: P<ast::Item>) -> SmallVector<P<ast::Item>> {
         match item.node {
             // Scope placeholder
@@ -155,6 +160,54 @@ impl Folder for PlaceholderExpander {
             _ => noop_fold_ty(ty, self),
         }
     }
+
+    fn fold_block(&mut self, block: P<ast::Block>) -> P<ast::Block> {
+        noop_fold_block(block, self).map(|mut block| {
+            let mut macros = Vec::new();
+            let mut remaining_stmts = block.stmts.len();
+
+            block.stmts = block.stmts.move_flat_map(|mut stmt| {
+                remaining_stmts -= 1;
+
+                // Scope placeholder
+                if let ast::StmtKind::Item(ref item) = stmt.node {
+                    if let ast::ItemKind::Mac(..) = item.node {
+                        macros.push(item.ident.ctxt.data().outer_mark);
+                        return None;
+                    }
+                }
+
+                match stmt.node {
+                    // Avoid wasting a node id on a trailing expression statement,
+                    // which shares a HIR node with the expression itself.
+                    ast::StmtKind::Expr(ref expr) if remaining_stmts == 0 => stmt.id = expr.id,
+
+                    _ => {
+                        assert_eq!(stmt.id, ast::DUMMY_NODE_ID);
+                        stmt.id = self.cx.resolver.next_node_id();
+                    }
+                }
+
+                if !macros.is_empty() {
+                    let macros = mem::replace(&mut macros, Vec::new());
+                    self.cx.resolver.add_expansions_at_stmt(stmt.id, macros);
+                }
+
+                Some(stmt)
+            });
+
+            block
+        })
+    }
+
+    fn fold_mod(&mut self, module: ast::Mod) -> ast::Mod {
+        let mut module = noop_fold_mod(module, self);
+        module.items = module.items.move_flat_map(|item| match item.node {
+            ast::ItemKind::Mac(_) => None, // remove scope placeholders from modules
+            _ => Some(item),
+        });
+        module
+    }
 }
 
 pub fn reconstructed_macro_rules(def: &ast::MacroDef, path: &ast::Path) -> Expansion {