def_collector: Simplify tracking of macro invocation parents

Avoid the tricky scheme with callbacks and keep the invocation parent data where it logically belongs - in `Definitions`.

This also allows to create `InvocationData` entries in resolve when the data is actually ready, and remove cells and "uninitialized" variants from it.

6 files changed, 55 insertions, 87 deletions

diff --git a/src/librustc/hir/map/def_collector.rs b/src/librustc/hir/map/def_collector.rs
index 41073773e9f..f946ca2903a 100644
--- a/src/librustc/hir/map/def_collector.rs
+++ b/src/librustc/hir/map/def_collector.rs
@@ -1,6 +1,5 @@
 use crate::hir::map::definitions::*;
-use crate::hir::def_id::{CRATE_DEF_INDEX, DefIndex};
-use crate::session::CrateDisambiguator;
+use crate::hir::def_id::DefIndex;
 
 use syntax::ast::*;
 use syntax::ext::hygiene::Mark;
@@ -14,31 +13,12 @@ pub struct DefCollector<'a> {
     definitions: &'a mut Definitions,
     parent_def: Option<DefIndex>,
     expansion: Mark,
-    pub visit_macro_invoc: Option<&'a mut dyn FnMut(MacroInvocationData)>,
-}
-
-pub struct MacroInvocationData {
-    pub mark: Mark,
-    pub def_index: DefIndex,
 }
 
 impl<'a> DefCollector<'a> {
     pub fn new(definitions: &'a mut Definitions, expansion: Mark) -> Self {
-        DefCollector {
-            definitions,
-            expansion,
-            parent_def: None,
-            visit_macro_invoc: None,
-        }
-    }
-
-    pub fn collect_root(&mut self,
-                        crate_name: &str,
-                        crate_disambiguator: CrateDisambiguator) {
-        let root = self.definitions.create_root_def(crate_name,
-                                                    crate_disambiguator);
-        assert_eq!(root, CRATE_DEF_INDEX);
-        self.parent_def = Some(root);
+        let parent_def = Some(definitions.invocation_parent(expansion));
+        DefCollector { definitions, parent_def, expansion }
     }
 
     fn create_def(&mut self,
@@ -97,12 +77,7 @@ impl<'a> DefCollector<'a> {
     }
 
     fn visit_macro_invoc(&mut self, id: NodeId) {
-        if let Some(ref mut visit) = self.visit_macro_invoc {
-            visit(MacroInvocationData {
-                mark: id.placeholder_to_mark(),
-                def_index: self.parent_def.unwrap(),
-            })
-        }
+        self.definitions.set_invocation_parent(id.placeholder_to_mark(), self.parent_def.unwrap());
     }
 }
 
diff --git a/src/librustc/hir/map/definitions.rs b/src/librustc/hir/map/definitions.rs
index 9e7898e10b0..9188a78e489 100644
--- a/src/librustc/hir/map/definitions.rs
+++ b/src/librustc/hir/map/definitions.rs
@@ -100,6 +100,9 @@ pub struct Definitions {
     expansions_that_defined: FxHashMap<DefIndex, Mark>,
     next_disambiguator: FxHashMap<(DefIndex, DefPathData), u32>,
     def_index_to_span: FxHashMap<DefIndex, Span>,
+    /// When collecting definitions from an AST fragment produced by a macro invocation `Mark`
+    /// we know what parent node that fragment should be attached to thanks to this table.
+    invocation_parents: FxHashMap<Mark, DefIndex>,
 }
 
 /// A unique identifier that we can use to lookup a definition
@@ -434,6 +437,7 @@ impl Definitions {
         assert!(self.def_index_to_node.is_empty());
         self.def_index_to_node.push(ast::CRATE_NODE_ID);
         self.node_to_def_index.insert(ast::CRATE_NODE_ID, root_index);
+        self.set_invocation_parent(Mark::root(), root_index);
 
         // Allocate some other DefIndices that always must exist.
         GlobalMetaDataKind::allocate_def_indices(self);
@@ -526,6 +530,15 @@ impl Definitions {
     pub fn add_parent_module_of_macro_def(&mut self, mark: Mark, module: DefId) {
         self.parent_modules_of_macro_defs.insert(mark, module);
     }
+
+    pub fn invocation_parent(&self, invoc_id: Mark) -> DefIndex {
+        self.invocation_parents[&invoc_id]
+    }
+
+    pub fn set_invocation_parent(&mut self, invoc_id: Mark, parent: DefIndex) {
+        let old_parent = self.invocation_parents.insert(invoc_id, parent);
+        assert!(old_parent.is_none(), "parent def-index is reset for an invocation");
+    }
 }
 
 impl DefPathData {
diff --git a/src/librustc/hir/map/mod.rs b/src/librustc/hir/map/mod.rs
index 63f60d0ab95..f9092196ab3 100644
--- a/src/librustc/hir/map/mod.rs
+++ b/src/librustc/hir/map/mod.rs
@@ -1,5 +1,5 @@
 use self::collector::NodeCollector;
-pub use self::def_collector::{DefCollector, MacroInvocationData};
+pub use self::def_collector::DefCollector;
 pub use self::definitions::{
     Definitions, DefKey, DefPath, DefPathData, DisambiguatedDefPathData, DefPathHash
 };
diff --git a/src/librustc_resolve/build_reduced_graph.rs b/src/librustc_resolve/build_reduced_graph.rs
index 974487ab9d2..8515029193e 100644
--- a/src/librustc_resolve/build_reduced_graph.rs
+++ b/src/librustc_resolve/build_reduced_graph.rs
@@ -944,12 +944,19 @@ pub struct BuildReducedGraphVisitor<'a, 'b> {
 
 impl<'a, 'b> BuildReducedGraphVisitor<'a, 'b> {
     fn visit_invoc(&mut self, id: ast::NodeId) -> &'b InvocationData<'b> {
-        let mark = id.placeholder_to_mark();
-        self.resolver.current_module.unresolved_invocations.borrow_mut().insert(mark);
-        let invocation = self.resolver.invocations[&mark];
-        invocation.module.set(self.resolver.current_module);
-        invocation.parent_legacy_scope.set(self.current_legacy_scope);
-        invocation
+        let invoc_id = id.placeholder_to_mark();
+
+        self.resolver.current_module.unresolved_invocations.borrow_mut().insert(invoc_id);
+
+        let invocation_data = self.resolver.arenas.alloc_invocation_data(InvocationData {
+            module: self.resolver.current_module,
+            parent_legacy_scope: self.current_legacy_scope,
+            output_legacy_scope: Cell::new(None),
+        });
+        let old_invocation_data = self.resolver.invocations.insert(invoc_id, invocation_data);
+        assert!(old_invocation_data.is_none(), "invocation data is reset for an invocation");
+
+        invocation_data
     }
 }
 
diff --git a/src/librustc_resolve/lib.rs b/src/librustc_resolve/lib.rs
index 207b0b3754a..09d7d8ace6b 100644
--- a/src/librustc_resolve/lib.rs
+++ b/src/librustc_resolve/lib.rs
@@ -19,7 +19,7 @@ use GenericParameters::*;
 use RibKind::*;
 use smallvec::smallvec;
 
-use rustc::hir::map::{Definitions, DefCollector};
+use rustc::hir::map::Definitions;
 use rustc::hir::{self, PrimTy, Bool, Char, Float, Int, Uint, Str};
 use rustc::middle::cstore::CrateStore;
 use rustc::session::Session;
@@ -1901,8 +1901,7 @@ impl<'a> Resolver<'a> {
         module_map.insert(DefId::local(CRATE_DEF_INDEX), graph_root);
 
         let mut definitions = Definitions::default();
-        DefCollector::new(&mut definitions, Mark::root())
-            .collect_root(crate_name, session.local_crate_disambiguator());
+        definitions.create_root_def(crate_name, session.local_crate_disambiguator());
 
         let mut extern_prelude: FxHashMap<Ident, ExternPreludeEntry<'_>> =
             session.opts.externs.iter().map(|kv| (Ident::from_str(kv.0), Default::default()))
diff --git a/src/librustc_resolve/macros.rs b/src/librustc_resolve/macros.rs
index 7f57b0a5052..8361bbe45c4 100644
--- a/src/librustc_resolve/macros.rs
+++ b/src/librustc_resolve/macros.rs
@@ -8,7 +8,7 @@ use crate::build_reduced_graph::{BuildReducedGraphVisitor, IsMacroExport};
 use crate::resolve_imports::ImportResolver;
 use rustc::hir::def_id::{CrateNum, DefId, DefIndex, CRATE_DEF_INDEX};
 use rustc::hir::def::{self, DefKind, NonMacroAttrKind};
-use rustc::hir::map::{self, DefCollector};
+use rustc::hir::map::DefCollector;
 use rustc::middle::stability;
 use rustc::{ty, lint, span_bug};
 use syntax::ast::{self, Ident};
@@ -32,14 +32,14 @@ use rustc_data_structures::sync::Lrc;
 
 type Res = def::Res<ast::NodeId>;
 
+// FIXME: Merge this with `ParentScope`.
 #[derive(Clone, Debug)]
 pub struct InvocationData<'a> {
-    def_index: DefIndex,
     /// The module in which the macro was invoked.
-    crate module: Cell<Module<'a>>,
+    crate module: Module<'a>,
     /// The legacy scope in which the macro was invoked.
     /// The invocation path is resolved in this scope.
-    crate parent_legacy_scope: Cell<LegacyScope<'a>>,
+    crate parent_legacy_scope: LegacyScope<'a>,
     /// The legacy scope *produced* by expanding this macro invocation,
     /// includes all the macro_rules items, other invocations, etc generated by it.
     /// `None` if the macro is not expanded yet.
@@ -49,10 +49,9 @@ pub struct InvocationData<'a> {
 impl<'a> InvocationData<'a> {
     pub fn root(graph_root: Module<'a>) -> Self {
         InvocationData {
-            module: Cell::new(graph_root),
-            def_index: CRATE_DEF_INDEX,
-            parent_legacy_scope: Cell::new(LegacyScope::Empty),
-            output_legacy_scope: Cell::new(Some(LegacyScope::Empty)),
+            module: graph_root,
+            parent_legacy_scope: LegacyScope::Empty,
+            output_legacy_scope: Cell::new(None),
         }
     }
 }
@@ -74,9 +73,6 @@ pub struct LegacyBinding<'a> {
 /// can potentially expand into macro definitions.
 #[derive(Copy, Clone, Debug)]
 pub enum LegacyScope<'a> {
-    /// Created when invocation data is allocated in the arena;
-    /// must be replaced with a proper scope later.
-    Uninitialized,
     /// Empty "root" scope at the crate start containing no names.
     Empty,
     /// The scope introduced by a `macro_rules!` macro definition.
@@ -139,11 +135,11 @@ impl<'a> base::Resolver for Resolver<'a> {
     fn get_module_scope(&mut self, id: ast::NodeId) -> Mark {
         let mark = Mark::fresh(Mark::root());
         let module = self.module_map[&self.definitions.local_def_id(id)];
+        self.definitions.set_invocation_parent(mark, module.def_id().unwrap().index);
         self.invocations.insert(mark, self.arenas.alloc_invocation_data(InvocationData {
-            module: Cell::new(module),
-            def_index: module.def_id().unwrap().index,
-            parent_legacy_scope: Cell::new(LegacyScope::Empty),
-            output_legacy_scope: Cell::new(Some(LegacyScope::Empty)),
+            module,
+            parent_legacy_scope: LegacyScope::Empty,
+            output_legacy_scope: Cell::new(None),
         }));
         mark
     }
@@ -160,16 +156,20 @@ impl<'a> base::Resolver for Resolver<'a> {
 
     fn visit_ast_fragment_with_placeholders(&mut self, mark: Mark, fragment: &AstFragment,
                                             derives: &[Mark]) {
-        let invocation = self.invocations[&mark];
-        self.collect_def_ids(mark, invocation, fragment);
+        fragment.visit_with(&mut DefCollector::new(&mut self.definitions, mark));
 
-        self.current_module = invocation.module.get();
+        let invocation = self.invocations[&mark];
+        self.current_module = invocation.module;
         self.current_module.unresolved_invocations.borrow_mut().remove(&mark);
         self.current_module.unresolved_invocations.borrow_mut().extend(derives);
+        let parent_def = self.definitions.invocation_parent(mark);
+        for &derive_invoc_id in derives {
+            self.definitions.set_invocation_parent(derive_invoc_id, parent_def);
+        }
         self.invocations.extend(derives.iter().map(|&derive| (derive, invocation)));
         let mut visitor = BuildReducedGraphVisitor {
             resolver: self,
-            current_legacy_scope: invocation.parent_legacy_scope.get(),
+            current_legacy_scope: invocation.parent_legacy_scope,
             expansion: mark,
         };
         fragment.visit_with(&mut visitor);
@@ -259,9 +259,9 @@ impl<'a> Resolver<'a> {
     fn invoc_parent_scope(&self, invoc_id: Mark, derives: Vec<ast::Path>) -> ParentScope<'a> {
         let invoc = self.invocations[&invoc_id];
         ParentScope {
-            module: invoc.module.get().nearest_item_scope(),
+            module: invoc.module.nearest_item_scope(),
             expansion: invoc_id.parent(),
-            legacy: invoc.parent_legacy_scope.get(),
+            legacy: invoc.parent_legacy_scope,
             derives,
         }
     }
@@ -829,10 +829,9 @@ impl<'a> Resolver<'a> {
                         binding.parent_legacy_scope
                     ),
                     LegacyScope::Invocation(invoc) => WhereToResolve::MacroRules(
-                        invoc.output_legacy_scope.get().unwrap_or(invoc.parent_legacy_scope.get())
+                        invoc.output_legacy_scope.get().unwrap_or(invoc.parent_legacy_scope)
                     ),
                     LegacyScope::Empty => WhereToResolve::Module(parent_scope.module),
-                    LegacyScope::Uninitialized => unreachable!(),
                 }
                 WhereToResolve::CrateRoot => match ns {
                     TypeNS => {
@@ -1084,31 +1083,6 @@ impl<'a> Resolver<'a> {
         }
     }
 
-    fn collect_def_ids(&mut self,
-                       mark: Mark,
-                       invocation: &'a InvocationData<'a>,
-                       fragment: &AstFragment) {
-        let Resolver { ref mut invocations, arenas, graph_root, .. } = *self;
-        let InvocationData { def_index, .. } = *invocation;
-
-        let visit_macro_invoc = &mut |invoc: map::MacroInvocationData| {
-            invocations.entry(invoc.mark).or_insert_with(|| {
-                arenas.alloc_invocation_data(InvocationData {
-                    def_index: invoc.def_index,
-                    module: Cell::new(graph_root),
-                    parent_legacy_scope: Cell::new(LegacyScope::Uninitialized),
-                    output_legacy_scope: Cell::new(None),
-                })
-            });
-        };
-
-        let mut def_collector = DefCollector::new(&mut self.definitions, mark);
-        def_collector.visit_macro_invoc = Some(visit_macro_invoc);
-        def_collector.with_parent(def_index, |def_collector| {
-            fragment.visit_with(def_collector)
-        });
-    }
-
     crate fn check_reserved_macro_name(&mut self, ident: Ident, res: Res) {
         // Reserve some names that are not quite covered by the general check
         // performed on `Resolver::builtin_attrs`.