Perform node id assignment and macros_at_scope construction during

the `InvocationCollector` and `PlaceholderExpander` folds.

src/librustc_driver/driver.rs

@@ -707,8 +707,6 @@ pub fn phase_2_configure_and_expand<'a, F>(sess: &Session,
                                                   &sess.features.borrow())
     });
 
-    let krate = time(sess.time_passes(), "assigning node ids", || resolver.assign_node_ids(krate));
-
     if sess.opts.debugging_opts.input_stats {
         println!("Post-expansion node count: {}", count_nodes(&krate));
     }

src/librustc_resolve/assign_ids.rs

@@ -1,92 +0,0 @@
-// 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 Resolver;
-use rustc::session::Session;
-use rustc::util::nodemap::FnvHashMap;
-use syntax::ast;
-use syntax::ext::hygiene::Mark;
-use syntax::fold::{self, Folder};
-use syntax::ptr::P;
-use syntax::util::move_map::MoveMap;
-use syntax::util::small_vector::SmallVector;
-
-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 FnvHashMap<ast::NodeId, Vec<Mark>>,
-}
-
-impl<'a> Folder for NodeIdAssigner<'a> {
-    fn new_id(&mut self, old_id: ast::NodeId) -> ast::NodeId {
-        assert_eq!(old_id, ast::DUMMY_NODE_ID);
-        self.sess.next_node_id()
-    }
-
-    fn fold_block(&mut self, block: P<ast::Block>) -> P<ast::Block> {
-        block.map(|mut block| {
-            block.id = self.new_id(block.id);
-
-            let stmt = block.stmts.pop();
-            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(item.ident.ctxt.data().outer_mark);
-                        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.
-                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),
-        }
-    }
-}

src/librustc_resolve/lib.rs

@@ -84,7 +84,6 @@ mod macros;
 mod check_unused;
 mod build_reduced_graph;
 mod resolve_imports;
-mod assign_ids;
 
 enum SuggestionType {
     Macro(String),

src/librustc_resolve/macros.rs

@@ -41,6 +41,10 @@ impl<'a> base::Resolver for Resolver<'a> {
         self.macro_loader.load_crate(extern_crate, allows_macros)
     }
 
+    fn next_node_id(&mut self) -> ast::NodeId {
+        self.session.next_node_id()
+    }
+
     fn visit_expansion(&mut self, mark: Mark, expansion: &Expansion) {
         expansion.visit_with(&mut ExpansionVisitor {
             current_module: self.expansion_data[mark.as_u32() as usize].module.clone(),

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>) {}

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 {
+        let (mac, style, attrs) = if let StmtKind::Mac(mac) = stmt.node {
-            StmtKind::Mac(mac) => mac.unwrap(),
+            mac.unwrap()
-            _ => return noop_fold_stmt(stmt, self),
+        } 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> {

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 {
+impl<'a, 'b> PlaceholderExpander<'a, 'b> {
-    pub fn new() -> Self {
+    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 {