resolve: Collapse `macro_rules` scope chains on the fly

5 files changed, 69 insertions, 29 deletions

diff --git a/compiler/rustc_resolve/src/build_reduced_graph.rs b/compiler/rustc_resolve/src/build_reduced_graph.rs
index 83016ed36a7..34145c3c138 100644
--- a/compiler/rustc_resolve/src/build_reduced_graph.rs
+++ b/compiler/rustc_resolve/src/build_reduced_graph.rs
@@ -7,7 +7,7 @@
 
 use crate::def_collector::collect_definitions;
 use crate::imports::{Import, ImportKind};
-use crate::macros::{MacroRulesBinding, MacroRulesScope};
+use crate::macros::{MacroRulesBinding, MacroRulesScope, MacroRulesScopeRef};
 use crate::Namespace::{self, MacroNS, TypeNS, ValueNS};
 use crate::{CrateLint, Determinacy, PathResult, ResolutionError, VisResolutionError};
 use crate::{
@@ -209,7 +209,7 @@ impl<'a> Resolver<'a> {
         &mut self,
         fragment: &AstFragment,
         parent_scope: ParentScope<'a>,
-    ) -> MacroRulesScope<'a> {
+    ) -> MacroRulesScopeRef<'a> {
         collect_definitions(self, fragment, parent_scope.expansion);
         let mut visitor = BuildReducedGraphVisitor { r: self, parent_scope };
         fragment.visit_with(&mut visitor);
@@ -220,7 +220,8 @@ impl<'a> Resolver<'a> {
         let def_id = module.def_id().expect("unpopulated module without a def-id");
         for child in self.cstore().item_children_untracked(def_id, self.session) {
             let child = child.map_id(|_| panic!("unexpected id"));
-            BuildReducedGraphVisitor { r: self, parent_scope: ParentScope::module(module) }
+            let parent_scope = ParentScope::module(module, self);
+            BuildReducedGraphVisitor { r: self, parent_scope }
                 .build_reduced_graph_for_external_crate_res(child);
         }
     }
@@ -1154,7 +1155,7 @@ impl<'a, 'b> BuildReducedGraphVisitor<'a, 'b> {
         false
     }
 
-    fn visit_invoc(&mut self, id: NodeId) -> MacroRulesScope<'a> {
+    fn visit_invoc(&mut self, id: NodeId) -> MacroRulesScopeRef<'a> {
         let invoc_id = id.placeholder_to_expn_id();
 
         self.parent_scope.module.unexpanded_invocations.borrow_mut().insert(invoc_id);
@@ -1162,7 +1163,9 @@ impl<'a, 'b> BuildReducedGraphVisitor<'a, 'b> {
         let old_parent_scope = self.r.invocation_parent_scopes.insert(invoc_id, self.parent_scope);
         assert!(old_parent_scope.is_none(), "invocation data is reset for an invocation");
 
-        MacroRulesScope::Invocation(invoc_id)
+        let scope = self.r.arenas.alloc_macro_rules_scope(MacroRulesScope::Invocation(invoc_id));
+        self.r.invocation_macro_rules_scopes.entry(invoc_id).or_default().insert(scope);
+        scope
     }
 
     fn proc_macro_stub(&self, item: &ast::Item) -> Option<(MacroKind, Ident, Span)> {
@@ -1196,7 +1199,7 @@ impl<'a, 'b> BuildReducedGraphVisitor<'a, 'b> {
         }
     }
 
-    fn define_macro(&mut self, item: &ast::Item) -> MacroRulesScope<'a> {
+    fn define_macro(&mut self, item: &ast::Item) -> MacroRulesScopeRef<'a> {
         let parent_scope = self.parent_scope;
         let expansion = parent_scope.expansion;
         let def_id = self.r.local_def_id(item.id);
@@ -1239,11 +1242,13 @@ impl<'a, 'b> BuildReducedGraphVisitor<'a, 'b> {
                 self.insert_unused_macro(ident, def_id, item.id, span);
             }
             self.r.visibilities.insert(def_id, vis);
-            MacroRulesScope::Binding(self.r.arenas.alloc_macro_rules_binding(MacroRulesBinding {
-                parent_macro_rules_scope: parent_scope.macro_rules,
-                binding,
-                ident,
-            }))
+            self.r.arenas.alloc_macro_rules_scope(MacroRulesScope::Binding(
+                self.r.arenas.alloc_macro_rules_binding(MacroRulesBinding {
+                    parent_macro_rules_scope: parent_scope.macro_rules,
+                    binding,
+                    ident,
+                }),
+            ))
         } else {
             let module = parent_scope.module;
             let vis = match item.kind {
diff --git a/compiler/rustc_resolve/src/diagnostics.rs b/compiler/rustc_resolve/src/diagnostics.rs
index 4e115c62c9e..59d5f104acf 100644
--- a/compiler/rustc_resolve/src/diagnostics.rs
+++ b/compiler/rustc_resolve/src/diagnostics.rs
@@ -630,7 +630,7 @@ impl<'a> Resolver<'a> {
                     }
                 }
                 Scope::MacroRules(macro_rules_scope) => {
-                    if let MacroRulesScope::Binding(macro_rules_binding) = macro_rules_scope {
+                    if let MacroRulesScope::Binding(macro_rules_binding) = macro_rules_scope.get() {
                         let res = macro_rules_binding.binding.res();
                         if filter_fn(res) {
                             suggestions
diff --git a/compiler/rustc_resolve/src/late.rs b/compiler/rustc_resolve/src/late.rs
index 2337f0d09ab..d2a1e7f84cc 100644
--- a/compiler/rustc_resolve/src/late.rs
+++ b/compiler/rustc_resolve/src/late.rs
@@ -677,7 +677,7 @@ impl<'a: 'ast, 'b, 'ast> LateResolutionVisitor<'a, 'b, 'ast> {
         // During late resolution we only track the module component of the parent scope,
         // although it may be useful to track other components as well for diagnostics.
         let graph_root = resolver.graph_root;
-        let parent_scope = ParentScope::module(graph_root);
+        let parent_scope = ParentScope::module(graph_root, resolver);
         let start_rib_kind = ModuleRibKind(graph_root);
         LateResolutionVisitor {
             r: resolver,
diff --git a/compiler/rustc_resolve/src/lib.rs b/compiler/rustc_resolve/src/lib.rs
index 30cd9944b1a..32db64b3669 100644
--- a/compiler/rustc_resolve/src/lib.rs
+++ b/compiler/rustc_resolve/src/lib.rs
@@ -64,7 +64,7 @@ use diagnostics::{extend_span_to_previous_binding, find_span_of_binding_until_ne
 use diagnostics::{ImportSuggestion, LabelSuggestion, Suggestion};
 use imports::{Import, ImportKind, ImportResolver, NameResolution};
 use late::{HasGenericParams, PathSource, Rib, RibKind::*};
-use macros::{MacroRulesBinding, MacroRulesScope};
+use macros::{MacroRulesBinding, MacroRulesScope, MacroRulesScopeRef};
 
 type Res = def::Res<NodeId>;
 
@@ -100,7 +100,7 @@ impl Determinacy {
 enum Scope<'a> {
     DeriveHelpers(ExpnId),
     DeriveHelpersCompat,
-    MacroRules(MacroRulesScope<'a>),
+    MacroRules(MacroRulesScopeRef<'a>),
     CrateRoot,
     Module(Module<'a>),
     RegisteredAttrs,
@@ -133,18 +133,18 @@ enum ScopeSet {
 pub struct ParentScope<'a> {
     module: Module<'a>,
     expansion: ExpnId,
-    macro_rules: MacroRulesScope<'a>,
+    macro_rules: MacroRulesScopeRef<'a>,
     derives: &'a [ast::Path],
 }
 
 impl<'a> ParentScope<'a> {
     /// Creates a parent scope with the passed argument used as the module scope component,
     /// and other scope components set to default empty values.
-    pub fn module(module: Module<'a>) -> ParentScope<'a> {
+    pub fn module(module: Module<'a>, resolver: &Resolver<'a>) -> ParentScope<'a> {
         ParentScope {
             module,
             expansion: ExpnId::root(),
-            macro_rules: MacroRulesScope::Empty,
+            macro_rules: resolver.arenas.alloc_macro_rules_scope(MacroRulesScope::Empty),
             derives: &[],
         }
     }
@@ -974,7 +974,10 @@ pub struct Resolver<'a> {
     invocation_parent_scopes: FxHashMap<ExpnId, ParentScope<'a>>,
     /// `macro_rules` scopes *produced* by expanding the macro invocations,
     /// include all the `macro_rules` items and other invocations generated by them.
-    output_macro_rules_scopes: FxHashMap<ExpnId, MacroRulesScope<'a>>,
+    output_macro_rules_scopes: FxHashMap<ExpnId, MacroRulesScopeRef<'a>>,
+    /// References to all `MacroRulesScope::Invocation(invoc_id)`s, used to update such scopes
+    /// when their corresponding `invoc_id`s get expanded.
+    invocation_macro_rules_scopes: FxHashMap<ExpnId, FxHashSet<MacroRulesScopeRef<'a>>>,
     /// Helper attributes that are in scope for the given expansion.
     helper_attrs: FxHashMap<ExpnId, Vec<Ident>>,
 
@@ -1043,6 +1046,9 @@ impl<'a> ResolverArenas<'a> {
     fn alloc_name_resolution(&'a self) -> &'a RefCell<NameResolution<'a>> {
         self.name_resolutions.alloc(Default::default())
     }
+    fn alloc_macro_rules_scope(&'a self, scope: MacroRulesScope<'a>) -> MacroRulesScopeRef<'a> {
+        PtrKey(self.dropless.alloc(Cell::new(scope)))
+    }
     fn alloc_macro_rules_binding(
         &'a self,
         binding: MacroRulesBinding<'a>,
@@ -1230,14 +1236,11 @@ impl<'a> Resolver<'a> {
         let (registered_attrs, registered_tools) =
             macros::registered_attrs_and_tools(session, &krate.attrs);
 
-        let mut invocation_parent_scopes = FxHashMap::default();
-        invocation_parent_scopes.insert(ExpnId::root(), ParentScope::module(graph_root));
-
         let features = session.features_untracked();
         let non_macro_attr =
             |mark_used| Lrc::new(SyntaxExtension::non_macro_attr(mark_used, session.edition()));
 
-        Resolver {
+        let mut resolver = Resolver {
             session,
 
             definitions,
@@ -1304,8 +1307,9 @@ impl<'a> Resolver<'a> {
             dummy_ext_bang: Lrc::new(SyntaxExtension::dummy_bang(session.edition())),
             dummy_ext_derive: Lrc::new(SyntaxExtension::dummy_derive(session.edition())),
             non_macro_attrs: [non_macro_attr(false), non_macro_attr(true)],
-            invocation_parent_scopes,
+            invocation_parent_scopes: Default::default(),
             output_macro_rules_scopes: Default::default(),
+            invocation_macro_rules_scopes: Default::default(),
             helper_attrs: Default::default(),
             local_macro_def_scopes: FxHashMap::default(),
             name_already_seen: FxHashMap::default(),
@@ -1332,7 +1336,12 @@ impl<'a> Resolver<'a> {
             invocation_parents,
             next_disambiguator: Default::default(),
             trait_impl_items: Default::default(),
-        }
+        };
+
+        let root_parent_scope = ParentScope::module(graph_root, &resolver);
+        resolver.invocation_parent_scopes.insert(ExpnId::root(), root_parent_scope);
+
+        resolver
     }
 
     pub fn next_node_id(&mut self) -> NodeId {
@@ -1702,7 +1711,7 @@ impl<'a> Resolver<'a> {
                 }
                 Scope::DeriveHelpers(..) => Scope::DeriveHelpersCompat,
                 Scope::DeriveHelpersCompat => Scope::MacroRules(parent_scope.macro_rules),
-                Scope::MacroRules(macro_rules_scope) => match macro_rules_scope {
+                Scope::MacroRules(macro_rules_scope) => match macro_rules_scope.get() {
                     MacroRulesScope::Binding(binding) => {
                         Scope::MacroRules(binding.parent_macro_rules_scope)
                     }
@@ -3199,7 +3208,7 @@ impl<'a> Resolver<'a> {
             }
         };
         let module = self.get_module(module_id);
-        let parent_scope = &ParentScope::module(module);
+        let parent_scope = &ParentScope::module(module, self);
         let res = self.resolve_ast_path(&path, ns, parent_scope).map_err(|_| ())?;
         Ok((path, res))
     }
diff --git a/compiler/rustc_resolve/src/macros.rs b/compiler/rustc_resolve/src/macros.rs
index b5b281b93bc..6bc9419ea84 100644
--- a/compiler/rustc_resolve/src/macros.rs
+++ b/compiler/rustc_resolve/src/macros.rs
@@ -11,6 +11,7 @@ use rustc_ast_lowering::ResolverAstLowering;
 use rustc_ast_pretty::pprust;
 use rustc_attr::StabilityLevel;
 use rustc_data_structures::fx::FxHashSet;
+use rustc_data_structures::ptr_key::PtrKey;
 use rustc_errors::struct_span_err;
 use rustc_expand::base::{Indeterminate, InvocationRes, ResolverExpand, SyntaxExtension};
 use rustc_expand::compile_declarative_macro;
@@ -29,6 +30,7 @@ use rustc_span::{Span, DUMMY_SP};
 
 use rustc_data_structures::sync::Lrc;
 use rustc_span::hygiene::{AstPass, MacroKind};
+use std::cell::Cell;
 use std::{mem, ptr};
 
 type Res = def::Res<NodeId>;
@@ -39,7 +41,7 @@ type Res = def::Res<NodeId>;
 pub struct MacroRulesBinding<'a> {
     crate binding: &'a NameBinding<'a>,
     /// `macro_rules` scope into which the `macro_rules` item was planted.
-    crate parent_macro_rules_scope: MacroRulesScope<'a>,
+    crate parent_macro_rules_scope: MacroRulesScopeRef<'a>,
     crate ident: Ident,
 }
 
@@ -59,6 +61,14 @@ pub enum MacroRulesScope<'a> {
     Invocation(ExpnId),
 }
 
+/// `macro_rules!` scopes are always kept by reference and inside a cell.
+/// The reason is that we update all scopes with value `MacroRulesScope::Invocation(invoc_id)`
+/// in-place immediately after `invoc_id` gets expanded.
+/// This helps to avoid uncontrollable growth of `macro_rules!` scope chains,
+/// which usually grow lineraly with the number of macro invocations
+/// in a module (including derives) and hurt performance.
+pub(crate) type MacroRulesScopeRef<'a> = PtrKey<'a, Cell<MacroRulesScope<'a>>>;
+
 // Macro namespace is separated into two sub-namespaces, one for bang macros and
 // one for attribute-like macros (attributes, derives).
 // We ignore resolutions from one sub-namespace when searching names in scope for another.
@@ -163,6 +173,22 @@ impl<'a> ResolverExpand for Resolver<'a> {
         let output_macro_rules_scope = self.build_reduced_graph(fragment, parent_scope);
         self.output_macro_rules_scopes.insert(expansion, output_macro_rules_scope);
 
+        // Update all `macro_rules` scopes referring to this invocation. This is an optimization
+        // used to avoid long scope chains, see the comments on `MacroRulesScopeRef`.
+        if let Some(invocation_scopes) = self.invocation_macro_rules_scopes.remove(&expansion) {
+            for invocation_scope in &invocation_scopes {
+                invocation_scope.set(output_macro_rules_scope.get());
+            }
+            // All `macro_rules` scopes that previously referred to `expansion`
+            // are now rerouted to its output scope, if it's also an invocation.
+            if let MacroRulesScope::Invocation(invoc_id) = output_macro_rules_scope.get() {
+                self.invocation_macro_rules_scopes
+                    .entry(invoc_id)
+                    .or_default()
+                    .extend(invocation_scopes);
+            }
+        }
+
         parent_scope.module.unexpanded_invocations.borrow_mut().remove(&expansion);
     }
 
@@ -655,7 +681,7 @@ impl<'a> Resolver<'a> {
                         }
                         result
                     }
-                    Scope::MacroRules(macro_rules_scope) => match macro_rules_scope {
+                    Scope::MacroRules(macro_rules_scope) => match macro_rules_scope.get() {
                         MacroRulesScope::Binding(macro_rules_binding)
                             if ident == macro_rules_binding.ident =>
                         {