rustdoc: Pre-calculate traits that are in scope for doc links

This eliminates one more late use of resolver

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1373,11 +1373,15 @@ fn get_traits(self) -> impl Iterator<Item = DefId> + 'a {
         self.root.traits.decode(self).map(move |index| self.local_def_id(index))
     }
 
-    fn get_trait_impls(self) -> impl Iterator<Item = (DefId, Option<SimplifiedType>)> + 'a {
-        self.cdata.trait_impls.values().flat_map(move |impls| {
-            impls
-                .decode(self)
-                .map(move |(idx, simplified_self_ty)| (self.local_def_id(idx), simplified_self_ty))
+    fn get_trait_impls(self) -> impl Iterator<Item = (DefId, DefId, Option<SimplifiedType>)> + 'a {
+        self.cdata.trait_impls.iter().flat_map(move |((trait_cnum_raw, trait_index), impls)| {
+            let trait_def_id = DefId {
+                krate: self.cnum_map[CrateNum::from_u32(*trait_cnum_raw)],
+                index: *trait_index,
+            };
+            impls.decode(self).map(move |(impl_index, simplified_self_ty)| {
+                (trait_def_id, self.local_def_id(impl_index), simplified_self_ty)
+            })
         })
     }
 

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -493,7 +493,7 @@ pub fn traits_in_crate_untracked(&self, cnum: CrateNum) -> impl Iterator<Item =
     pub fn trait_impls_in_crate_untracked(
         &self,
         cnum: CrateNum,
-    ) -> impl Iterator<Item = (DefId, Option<SimplifiedType>)> + '_ {
+    ) -> impl Iterator<Item = (DefId, DefId, Option<SimplifiedType>)> + '_ {
         self.get_crate_data(cnum).get_trait_impls()
     }
 }

compiler/rustc_resolve/src/build_reduced_graph.rs

@@ -108,7 +108,7 @@ impl<'a> Resolver<'a> {
     /// Reachable macros with block module parents exist due to `#[macro_export] macro_rules!`,
     /// but they cannot use def-site hygiene, so the assumption holds
     /// (<https://github.com/rust-lang/rust/pull/77984#issuecomment-712445508>).
-    crate fn get_nearest_non_block_module(&mut self, mut def_id: DefId) -> Module<'a> {
+    pub fn get_nearest_non_block_module(&mut self, mut def_id: DefId) -> Module<'a> {
         loop {
             match self.get_module(def_id) {
                 Some(module) => return module,

compiler/rustc_resolve/src/lib.rs

@@ -614,7 +614,8 @@ fn res(&self) -> Option<Res> {
         }
     }
 
-    fn def_id(&self) -> DefId {
+    // Public for rustdoc.
+    pub fn def_id(&self) -> DefId {
         self.opt_def_id().expect("`ModuleData::def_id` is called on a block module")
     }
 
@@ -3407,6 +3408,16 @@ pub fn all_macros(&self) -> &FxHashMap<Symbol, Res> {
         &self.all_macros
     }
 
+    /// For rustdoc.
+    /// For local modules returns only reexports, for external modules returns all children.
+    pub fn module_children_or_reexports(&self, def_id: DefId) -> Vec<ModChild> {
+        if let Some(def_id) = def_id.as_local() {
+            self.reexport_map.get(&def_id).cloned().unwrap_or_default()
+        } else {
+            self.cstore().module_children_untracked(def_id, self.session)
+        }
+    }
+
     /// Retrieves the span of the given `DefId` if `DefId` is in the local crate.
     #[inline]
     pub fn opt_span(&self, def_id: DefId) -> Option<Span> {

src/librustdoc/core.rs

@@ -4,9 +4,9 @@
 use rustc_errors::json::JsonEmitter;
 use rustc_feature::UnstableFeatures;
 use rustc_hir::def::Res;
-use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::intravisit::{self, Visitor};
-use rustc_hir::{HirId, Path};
+use rustc_hir::{HirId, Path, TraitCandidate};
 use rustc_interface::interface;
 use rustc_middle::hir::nested_filter;
 use rustc_middle::middle::privacy::AccessLevels;
@@ -33,6 +33,9 @@
 crate use rustc_session::config::{DebuggingOptions, Input, Options};
 
 crate struct ResolverCaches {
+    /// Traits in scope for a given module.
+    /// See `collect_intra_doc_links::traits_implemented_by` for more details.
+    crate traits_in_scope: DefIdMap<Vec<TraitCandidate>>,
     crate all_traits: Option<Vec<DefId>>,
     crate all_trait_impls: Option<Vec<DefId>>,
 }
@@ -67,11 +70,6 @@
     crate auto_traits: Vec<DefId>,
     /// The options given to rustdoc that could be relevant to a pass.
     crate render_options: RenderOptions,
-    /// The traits in scope for a given module.
-    ///
-    /// See `collect_intra_doc_links::traits_implemented_by` for more details.
-    /// `map<module, set<trait>>`
-    crate module_trait_cache: FxHashMap<DefId, FxHashSet<DefId>>,
     /// This same cache is used throughout rustdoc, including in [`crate::html::render`].
     crate cache: Cache,
     /// Used by [`clean::inline`] to tell if an item has already been inlined.
@@ -350,7 +348,6 @@ impl<'tcx> DocContext<'tcx> {
         impl_trait_bounds: Default::default(),
         generated_synthetics: Default::default(),
         auto_traits,
-        module_trait_cache: FxHashMap::default(),
         cache: Cache::new(access_levels, render_options.document_private),
         inlined: FxHashSet::default(),
         output_format,

src/librustdoc/passes/collect_intra_doc_links.rs

@@ -16,7 +16,7 @@
 use rustc_middle::{bug, span_bug, ty};
 use rustc_resolve::ParentScope;
 use rustc_session::lint::Lint;
-use rustc_span::hygiene::{MacroKind, SyntaxContext};
+use rustc_span::hygiene::MacroKind;
 use rustc_span::symbol::{sym, Ident, Symbol};
 use rustc_span::{BytePos, DUMMY_SP};
 use smallvec::{smallvec, SmallVec};
@@ -925,20 +925,9 @@ fn trait_impls_for<'a>(
     ty: Ty<'a>,
     module: DefId,
 ) -> FxHashSet<(DefId, DefId)> {
-    let mut resolver = cx.resolver.borrow_mut();
-    let in_scope_traits = cx.module_trait_cache.entry(module).or_insert_with(|| {
-        resolver.access(|resolver| {
-            let parent_scope = &ParentScope::module(resolver.expect_module(module), resolver);
-            resolver
-                .traits_in_scope(None, parent_scope, SyntaxContext::root(), None)
-                .into_iter()
-                .map(|candidate| candidate.def_id)
-                .collect()
-        })
-    });
-
     let tcx = cx.tcx;
-    let iter = in_scope_traits.iter().flat_map(|&trait_| {
+    let iter = cx.resolver_caches.traits_in_scope[&module].iter().flat_map(|trait_candidate| {
+        let trait_ = trait_candidate.def_id;
         trace!("considering explicit impl for trait {:?}", trait_);
 
         // Look at each trait implementation to see if it's an impl for `did`

src/librustdoc/passes/collect_intra_doc_links/early.rs

@@ -7,11 +7,15 @@
 use rustc_ast::{self as ast, ItemKind};
 use rustc_ast_lowering::ResolverAstLowering;
 use rustc_hir::def::Namespace::TypeNS;
-use rustc_hir::def_id::{DefId, LocalDefId, CRATE_DEF_ID};
-use rustc_resolve::Resolver;
+use rustc_hir::def::{DefKind, Res};
+use rustc_hir::def_id::{DefId, DefIdMap, DefIdSet, LocalDefId, CRATE_DEF_ID};
+use rustc_hir::TraitCandidate;
+use rustc_middle::ty::{DefIdTree, Visibility};
+use rustc_resolve::{ParentScope, Resolver};
 use rustc_session::config::Externs;
-use rustc_span::{Span, DUMMY_SP};
+use rustc_span::{Span, SyntaxContext, DUMMY_SP};
 
+use std::collections::hash_map::Entry;
 use std::mem;
 
 crate fn early_resolve_intra_doc_links(
@@ -22,15 +26,18 @@
     let mut loader = IntraLinkCrateLoader {
         resolver,
         current_mod: CRATE_DEF_ID,
+        visited_mods: Default::default(),
+        traits_in_scope: Default::default(),
         all_traits: Default::default(),
         all_trait_impls: Default::default(),
     };
 
     // Overridden `visit_item` below doesn't apply to the crate root,
-    // so we have to visit its attributes and exports separately.
+    // so we have to visit its attributes and reexports separately.
     loader.load_links_in_attrs(&krate.attrs, krate.span);
+    loader.process_module_children_or_reexports(CRATE_DEF_ID.to_def_id());
     visit::walk_crate(&mut loader, krate);
-    loader.fill_resolver_caches();
+    loader.add_foreign_traits_in_scope();
 
     // FIXME: somehow rustdoc is still missing crates even though we loaded all
     // the known necessary crates. Load them all unconditionally until we find a way to fix this.
@@ -46,6 +53,7 @@
     }
 
     ResolverCaches {
+        traits_in_scope: loader.traits_in_scope,
         all_traits: Some(loader.all_traits),
         all_trait_impls: Some(loader.all_trait_impls),
     }
@@ -54,27 +62,87 @@
 struct IntraLinkCrateLoader<'r, 'ra> {
     resolver: &'r mut Resolver<'ra>,
     current_mod: LocalDefId,
+    visited_mods: DefIdSet,
+    traits_in_scope: DefIdMap<Vec<TraitCandidate>>,
     all_traits: Vec<DefId>,
     all_trait_impls: Vec<DefId>,
 }
 
 impl IntraLinkCrateLoader<'_, '_> {
-    fn fill_resolver_caches(&mut self) {
-        for cnum in self.resolver.cstore().crates_untracked() {
-            let all_traits = self.resolver.cstore().traits_in_crate_untracked(cnum);
-            let all_trait_impls = self.resolver.cstore().trait_impls_in_crate_untracked(cnum);
+    fn add_traits_in_scope(&mut self, def_id: DefId) {
+        // Calls to `traits_in_scope` are expensive, so try to avoid them if only possible.
+        // Keys in the `traits_in_scope` cache are always module IDs.
+        if let Entry::Vacant(entry) = self.traits_in_scope.entry(def_id) {
+            let module = self.resolver.get_nearest_non_block_module(def_id);
+            let module_id = module.def_id();
+            let entry = if module_id == def_id {
+                Some(entry)
+            } else if let Entry::Vacant(entry) = self.traits_in_scope.entry(module_id) {
+                Some(entry)
+            } else {
+                None
+            };
+            if let Some(entry) = entry {
+                entry.insert(self.resolver.traits_in_scope(
+                    None,
+                    &ParentScope::module(module, self.resolver),
+                    SyntaxContext::root(),
+                    None,
+                ));
+            }
+        }
+    }
+
+    fn add_traits_in_parent_scope(&mut self, def_id: DefId) {
+        if let Some(module_id) = self.resolver.parent(def_id) {
+            self.add_traits_in_scope(module_id);
+        }
+    }
+
+    /// Add traits in scope for links in impls collected by the `collect-intra-doc-links` pass.
+    /// That pass filters impls using type-based information, but we don't yet have such
+    /// information here, so we just conservatively calculate traits in scope for *all* modules
+    /// having impls in them.
+    fn add_foreign_traits_in_scope(&mut self) {
+        for cnum in Vec::from_iter(self.resolver.cstore().crates_untracked()) {
+            // FIXME: Due to #78696 rustdoc can query traits in scope for any crate root.
+            self.add_traits_in_scope(cnum.as_def_id());
+
+            let all_traits = Vec::from_iter(self.resolver.cstore().traits_in_crate_untracked(cnum));
+            let all_trait_impls =
+                Vec::from_iter(self.resolver.cstore().trait_impls_in_crate_untracked(cnum));
+
+            // Querying traits in scope is expensive so we try to prune the impl and traits lists
+            // using privacy, private traits and impls from other crates are never documented in
+            // the current crate, and links in their doc comments are not resolved.
+            for &def_id in &all_traits {
+                if self.resolver.cstore().visibility_untracked(def_id) == Visibility::Public {
+                    self.add_traits_in_parent_scope(def_id);
+                }
+            }
+            for &(trait_def_id, impl_def_id, simplified_self_ty) in &all_trait_impls {
+                if self.resolver.cstore().visibility_untracked(trait_def_id) == Visibility::Public
+                    && simplified_self_ty.and_then(|ty| ty.def()).map_or(true, |ty_def_id| {
+                        self.resolver.cstore().visibility_untracked(ty_def_id) == Visibility::Public
+                    })
+                {
+                    self.add_traits_in_parent_scope(impl_def_id);
+                }
+            }
 
             self.all_traits.extend(all_traits);
-            self.all_trait_impls.extend(all_trait_impls.into_iter().map(|(def_id, _)| def_id));
+            self.all_trait_impls.extend(all_trait_impls.into_iter().map(|(_, def_id, _)| def_id));
         }
     }
 
     fn load_links_in_attrs(&mut self, attrs: &[ast::Attribute], span: Span) {
-        // FIXME: this needs to consider export inlining.
+        // FIXME: this needs to consider reexport inlining.
         let attrs = clean::Attributes::from_ast(attrs, None);
         for (parent_module, doc) in attrs.collapsed_doc_value_by_module_level() {
             let module_id = parent_module.unwrap_or(self.current_mod.to_def_id());
 
+            self.add_traits_in_scope(module_id);
+
             for link in markdown_links(&doc.as_str()) {
                 let path_str = if let Some(Ok(x)) = preprocess_link(&link) {
                     x.path_str
@@ -85,6 +153,26 @@ fn load_links_in_attrs(&mut self, attrs: &[ast::Attribute], span: Span) {
             }
         }
     }
+
+    /// When reexports are inlined, they are replaced with item which they refer to, those items
+    /// may have links in their doc comments, those links are resolved at the item definition site,
+    /// so we need to know traits in scope at that definition site.
+    fn process_module_children_or_reexports(&mut self, module_id: DefId) {
+        if !self.visited_mods.insert(module_id) {
+            return; // avoid infinite recursion
+        }
+
+        for child in self.resolver.module_children_or_reexports(module_id) {
+            if child.vis == Visibility::Public {
+                if let Some(def_id) = child.res.opt_def_id() {
+                    self.add_traits_in_parent_scope(def_id);
+                }
+                if let Res::Def(DefKind::Mod, module_id) = child.res {
+                    self.process_module_children_or_reexports(module_id);
+                }
+            }
+        }
+    }
 }
 
 impl Visitor<'_> for IntraLinkCrateLoader<'_, '_> {
@@ -93,6 +181,7 @@ fn visit_item(&mut self, item: &ast::Item) {
             let old_mod = mem::replace(&mut self.current_mod, self.resolver.local_def_id(item.id));
 
             self.load_links_in_attrs(&item.attrs, item.span);
+            self.process_module_children_or_reexports(self.current_mod.to_def_id());
             visit::walk_item(self, item);
 
             self.current_mod = old_mod;