Resolve `$crate` in all hygienic contexts for pretty-pringing

Stop visiting AST to discover those contexts, just iterate through hygiene data instead

src/librustc_resolve/build_reduced_graph.rs

@@ -758,7 +758,10 @@ pub fn get_module(&mut self, def_id: DefId) -> Module<'a> {
     }
 
     pub fn macro_def_scope(&mut self, expansion: Mark) -> Module<'a> {
-        let def_id = self.macro_defs[&expansion];
+        let def_id = match self.macro_defs.get(&expansion) {
+            Some(def_id) => *def_id,
+            None => return self.graph_root,
+        };
         if let Some(id) = self.definitions.as_local_node_id(def_id) {
             self.local_macro_def_scopes[&id]
         } else if def_id.krate == CrateNum::BuiltinMacros {

src/librustc_resolve/macros.rs

@@ -17,12 +17,11 @@
 use syntax::ext::base::{self, Determinacy};
 use syntax::ext::base::{MacroKind, SyntaxExtension};
 use syntax::ext::expand::{AstFragment, Invocation, InvocationKind};
-use syntax::ext::hygiene::Mark;
+use syntax::ext::hygiene::{self, Mark};
 use syntax::ext::tt::macro_rules;
 use syntax::feature_gate::{feature_err, emit_feature_err, is_builtin_attr_name};
 use syntax::feature_gate::{AttributeGate, GateIssue, Stability, BUILTIN_ATTRIBUTES};
 use syntax::symbol::{Symbol, kw, sym};
-use syntax::visit::Visitor;
 use syntax::util::lev_distance::find_best_match_for_name;
 use syntax_pos::{Span, DUMMY_SP};
 use errors::Applicability;
@@ -146,24 +145,14 @@ fn get_module_scope(&mut self, id: ast::NodeId) -> Mark {
         mark
     }
 
-    fn resolve_dollar_crates(&mut self, fragment: &AstFragment) {
-        struct ResolveDollarCrates<'a, 'b> {
-            resolver: &'a mut Resolver<'b>
-        }
-        impl<'a> Visitor<'a> for ResolveDollarCrates<'a, '_> {
-            fn visit_ident(&mut self, ident: Ident) {
-                if ident.name == kw::DollarCrate {
-                    let name = match self.resolver.resolve_crate_root(ident).kind {
-                        ModuleKind::Def(.., name) if name != kw::Invalid => name,
-                        _ => kw::Crate,
-                    };
-                    ident.span.ctxt().set_dollar_crate_name(name);
-                }
+    fn resolve_dollar_crates(&mut self) {
+        hygiene::update_dollar_crate_names(|ctxt| {
+            let ident = Ident::new(kw::DollarCrate, DUMMY_SP.with_ctxt(ctxt));
+            match self.resolve_crate_root(ident).kind {
+                ModuleKind::Def(.., name) if name != kw::Invalid => name,
+                _ => kw::Crate,
             }
-            fn visit_mac(&mut self, _: &ast::Mac) {}
-        }
-
-        fragment.visit_with(&mut ResolveDollarCrates { resolver: self });
+        });
     }
 
     fn visit_ast_fragment_with_placeholders(&mut self, mark: Mark, fragment: &AstFragment,

src/libsyntax/ext/base.rs

@@ -701,7 +701,7 @@ pub trait Resolver {
 
     fn get_module_scope(&mut self, id: ast::NodeId) -> Mark;
 
-    fn resolve_dollar_crates(&mut self, fragment: &AstFragment);
+    fn resolve_dollar_crates(&mut self);
     fn visit_ast_fragment_with_placeholders(&mut self, mark: Mark, fragment: &AstFragment,
                                             derives: &[Mark]);
     fn add_builtin(&mut self, ident: ast::Ident, ext: Lrc<SyntaxExtension>);

src/libsyntax/ext/expand.rs

@@ -429,7 +429,7 @@ fn resolve_imports(&mut self) {
     fn collect_invocations(&mut self, mut fragment: AstFragment, derives: &[Mark])
                            -> (AstFragment, Vec<Invocation>) {
         // Resolve `$crate`s in the fragment for pretty-printing.
-        self.cx.resolver.resolve_dollar_crates(&fragment);
+        self.cx.resolver.resolve_dollar_crates();
 
         let invocations = {
             let mut collector = InvocationCollector {

src/libsyntax_pos/hygiene.rs

@@ -33,7 +33,7 @@
 use serialize::{Encodable, Decodable, Encoder, Decoder};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::sync::Lrc;
-use std::{fmt, mem};
+use std::fmt;
 
 /// A SyntaxContext represents a chain of macro expansions (represented by marks).
 #[derive(Clone, Copy, PartialEq, Eq, Default, PartialOrd, Ord, Hash)]
@@ -387,6 +387,23 @@ pub fn walk_chain(span: Span, to: SyntaxContext) -> Span {
     HygieneData::with(|data| data.walk_chain(span, to))
 }
 
+pub fn update_dollar_crate_names(mut get_name: impl FnMut(SyntaxContext) -> Symbol) {
+    // The new contexts that need updating are at the end of the list and have `$crate` as a name.
+    let (len, to_update) = HygieneData::with(|data| (
+        data.syntax_contexts.len(),
+        data.syntax_contexts.iter().rev()
+            .take_while(|scdata| scdata.dollar_crate_name == kw::DollarCrate).count()
+    ));
+    // The callback must be called from outside of the `HygieneData` lock,
+    // since it will try to acquire it too.
+    let range_to_update = len - to_update .. len;
+    let names: Vec<_> =
+        range_to_update.clone().map(|idx| get_name(SyntaxContext::from_u32(idx as u32))).collect();
+    HygieneData::with(|data| range_to_update.zip(names.into_iter()).for_each(|(idx, name)| {
+        data.syntax_contexts[idx].dollar_crate_name = name;
+    }))
+}
+
 impl SyntaxContext {
     #[inline]
     pub const fn empty() -> Self {
@@ -614,17 +631,6 @@ pub fn outer_and_expn_info(self) -> (Mark, Option<ExpnInfo>) {
     pub fn dollar_crate_name(self) -> Symbol {
         HygieneData::with(|data| data.syntax_contexts[self.0 as usize].dollar_crate_name)
     }
-
-    pub fn set_dollar_crate_name(self, dollar_crate_name: Symbol) {
-        HygieneData::with(|data| {
-            let prev_dollar_crate_name = mem::replace(
-                &mut data.syntax_contexts[self.0 as usize].dollar_crate_name, dollar_crate_name
-            );
-            assert!(dollar_crate_name == prev_dollar_crate_name ||
-                    prev_dollar_crate_name == kw::DollarCrate,
-                    "$crate name is reset for a syntax context");
-        })
-    }
 }
 
 impl fmt::Debug for SyntaxContext {

src/test/ui/proc-macro/dollar-crate-issue-62325.rs

-//~ ERROR expected type, found `$`
-
+// check-pass
 // edition:2018
 // aux-build:test-macros.rs
 

src/test/ui/proc-macro/dollar-crate-issue-62325.stderr

-error: expected type, found `$`
-
-error: aborting due to previous error
-

src/test/ui/proc-macro/dollar-crate-issue-62325.stdout

-PRINT-ATTR INPUT (DISPLAY): struct A(identity!($crate :: S));
-PRINT-ATTR RE-COLLECTED (DISPLAY): struct A ( identity ! ( $ crate :: S ) ) ;
+PRINT-ATTR INPUT (DISPLAY): struct A(identity!(crate :: S));
+PRINT-ATTR RE-COLLECTED (DISPLAY): struct A ( identity ! ( crate :: S ) ) ;
 PRINT-ATTR INPUT (DEBUG): TokenStream [
     Ident {
         ident: "struct",
-        span: #0 bytes(0..0),
+        span: #2 bytes(LO..HI),
     },
     Ident {
         ident: "A",
-        span: #0 bytes(0..0),
+        span: #2 bytes(LO..HI),
     },
     Group {
         delimiter: Parenthesis,
         stream: TokenStream [
             Ident {
                 ident: "identity",
-                span: #0 bytes(0..0),
+                span: #2 bytes(LO..HI),
             },
             Punct {
                 ch: '!',
                 spacing: Alone,
-                span: #0 bytes(0..0),
+                span: #2 bytes(LO..HI),
             },
             Group {
                 delimiter: Parenthesis,
                 stream: TokenStream [
-                    Punct {
-                        ch: '$',
-                        spacing: Alone,
-                        span: #0 bytes(0..0),
-                    },
                     Ident {
-                        ident: "crate",
-                        span: #0 bytes(0..0),
+                        ident: "$crate",
+                        span: #2 bytes(LO..HI),
                     },
                     Punct {
                         ch: ':',
                         spacing: Joint,
-                        span: #0 bytes(0..0),
+                        span: #2 bytes(LO..HI),
                     },
                     Punct {
                         ch: ':',
                         spacing: Alone,
-                        span: #0 bytes(0..0),
+                        span: #2 bytes(LO..HI),
                     },
                     Ident {
                         ident: "S",
-                        span: #0 bytes(0..0),
+                        span: #2 bytes(LO..HI),
                     },
                 ],
-                span: #0 bytes(0..0),
+                span: #2 bytes(LO..HI),
             },
         ],
-        span: #0 bytes(0..0),
+        span: #2 bytes(LO..HI),
     },
     Punct {
         ch: ';',
         spacing: Alone,
-        span: #0 bytes(0..0),
+        span: #2 bytes(LO..HI),
     },
 ]