Auto merge of #92244 - petrochenkov:alltraits, r=cjgillot

rustc_metadata: Encode list of all crate's traits into metadata

While working on #88679 I noticed that rustdoc is casually doing something quite expensive, something that is used only for error reporting in rustc - collecting all traits from all crates in the dependency tree.

This PR trades some minor extra time spent by metadata encoder in rustc for major gains for rustdoc (and for rustc runs with errors, which execute the `all_traits` query for better diagnostics).

Author: bors

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1375,6 +1375,10 @@ impl<'a, 'tcx> CrateMetadataRef<'a> {
         )
     }
 
+    fn get_traits(&'a self) -> impl Iterator<Item = DefId> + 'a {
+        self.root.traits.decode(self).map(|index| self.local_def_id(index))
+    }
+
     fn get_implementations_for_trait(
         &self,
         tcx: TyCtxt<'tcx>,

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

@@ -4,9 +4,11 @@ use crate::native_libs;
 
 use rustc_ast as ast;
 use rustc_data_structures::stable_map::FxHashMap;
+use rustc_hir as hir;
 use rustc_hir::def::{CtorKind, DefKind};
 use rustc_hir::def_id::{CrateNum, DefId, DefIdMap, CRATE_DEF_INDEX, LOCAL_CRATE};
 use rustc_hir::definitions::{DefKey, DefPath, DefPathHash};
+use rustc_hir::itemlikevisit::ItemLikeVisitor;
 use rustc_middle::hir::exports::Export;
 use rustc_middle::middle::exported_symbols::ExportedSymbol;
 use rustc_middle::middle::stability::DeprecationEntry;
@@ -195,6 +197,8 @@ provide! { <'tcx> tcx, def_id, other, cdata,
 
     extra_filename => { cdata.root.extra_filename.clone() }
 
+    traits_in_crate => { tcx.arena.alloc_from_iter(cdata.get_traits()) }
+
     implementations_of_trait => {
         cdata.get_implementations_for_trait(tcx, Some(other))
     }
@@ -285,6 +289,28 @@ pub fn provide(providers: &mut Providers) {
                 foreign_modules::collect(tcx).into_iter().map(|m| (m.def_id, m)).collect();
             Lrc::new(modules)
         },
+        traits_in_crate: |tcx, cnum| {
+            assert_eq!(cnum, LOCAL_CRATE);
+
+            #[derive(Default)]
+            struct TraitsVisitor {
+                traits: Vec<DefId>,
+            }
+            impl ItemLikeVisitor<'_> for TraitsVisitor {
+                fn visit_item(&mut self, item: &hir::Item<'_>) {
+                    if let hir::ItemKind::Trait(..) | hir::ItemKind::TraitAlias(..) = item.kind {
+                        self.traits.push(item.def_id.to_def_id());
+                    }
+                }
+                fn visit_trait_item(&mut self, _trait_item: &hir::TraitItem<'_>) {}
+                fn visit_impl_item(&mut self, _impl_item: &hir::ImplItem<'_>) {}
+                fn visit_foreign_item(&mut self, _foreign_item: &hir::ForeignItem<'_>) {}
+            }
+
+            let mut visitor = TraitsVisitor::default();
+            tcx.hir().visit_all_item_likes(&mut visitor);
+            tcx.arena.alloc_slice(&visitor.traits)
+        },
 
         // Returns a map from a sufficiently visible external item (i.e., an
         // external item that is visible from at least one local module) to a

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -614,8 +614,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
 
         // Encode the def IDs of impls, for coherence checking.
         i = self.position();
-        let impls = self.encode_impls();
-        let impl_bytes = self.position() - i;
+        let (traits, impls) = self.encode_traits_and_impls();
+        let traits_and_impls_bytes = self.position() - i;
 
         let tcx = self.tcx;
 
@@ -727,6 +727,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             foreign_modules,
             source_map,
             impls,
+            traits,
             exported_symbols,
             interpret_alloc_index,
             tables,
@@ -753,7 +754,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             eprintln!(" diagnostic item bytes: {}", diagnostic_item_bytes);
             eprintln!("          native bytes: {}", native_lib_bytes);
             eprintln!("      source_map bytes: {}", source_map_bytes);
-            eprintln!("            impl bytes: {}", impl_bytes);
+            eprintln!("traits and impls bytes: {}", traits_and_impls_bytes);
             eprintln!("    exp. symbols bytes: {}", exported_symbols_bytes);
             eprintln!("  def-path table bytes: {}", def_path_table_bytes);
             eprintln!(" def-path hashes bytes: {}", def_path_hash_map_bytes);
@@ -1791,16 +1792,23 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
     }
 
     /// Encodes an index, mapping each trait to its (local) implementations.
-    fn encode_impls(&mut self) -> Lazy<[TraitImpls]> {
-        empty_proc_macro!(self);
-        debug!("EncodeContext::encode_impls()");
+    fn encode_traits_and_impls(&mut self) -> (Lazy<[DefIndex]>, Lazy<[TraitImpls]>) {
+        if self.is_proc_macro {
+            return (Lazy::empty(), Lazy::empty());
+        }
+        debug!("EncodeContext::encode_traits_and_impls()");
         let tcx = self.tcx;
-        let mut visitor = ImplVisitor { tcx, impls: FxHashMap::default() };
+        let mut visitor =
+            TraitsAndImplsVisitor { tcx, impls: FxHashMap::default(), traits: Default::default() };
         tcx.hir().visit_all_item_likes(&mut visitor);
 
+        let mut all_traits = visitor.traits;
         let mut all_impls: Vec<_> = visitor.impls.into_iter().collect();
 
         // Bring everything into deterministic order for hashing
+        all_traits.sort_by_cached_key(|&local_def_index| {
+            tcx.hir().def_path_hash(LocalDefId { local_def_index })
+        });
         all_impls.sort_by_cached_key(|&(trait_def_id, _)| tcx.def_path_hash(trait_def_id));
 
         let all_impls: Vec<_> = all_impls
@@ -1818,7 +1826,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             })
             .collect();
 
-        self.lazy(&all_impls)
+        (self.lazy(&all_traits), self.lazy(&all_impls))
     }
 
     // Encodes all symbols exported from this crate into the metadata.
@@ -2040,27 +2048,34 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
     }
 }
 
-struct ImplVisitor<'tcx> {
+struct TraitsAndImplsVisitor<'tcx> {
     tcx: TyCtxt<'tcx>,
+    traits: Vec<DefIndex>,
     impls: FxHashMap<DefId, Vec<(DefIndex, Option<fast_reject::SimplifiedType>)>>,
 }
 
-impl<'tcx, 'v> ItemLikeVisitor<'v> for ImplVisitor<'tcx> {
+impl<'tcx, 'v> ItemLikeVisitor<'v> for TraitsAndImplsVisitor<'tcx> {
     fn visit_item(&mut self, item: &hir::Item<'_>) {
-        if let hir::ItemKind::Impl { .. } = item.kind {
-            if let Some(trait_ref) = self.tcx.impl_trait_ref(item.def_id.to_def_id()) {
-                let simplified_self_ty = fast_reject::simplify_type(
-                    self.tcx,
-                    trait_ref.self_ty(),
-                    SimplifyParams::No,
-                    StripReferences::No,
-                );
+        match item.kind {
+            hir::ItemKind::Trait(..) | hir::ItemKind::TraitAlias(..) => {
+                self.traits.push(item.def_id.local_def_index);
+            }
+            hir::ItemKind::Impl(..) => {
+                if let Some(trait_ref) = self.tcx.impl_trait_ref(item.def_id.to_def_id()) {
+                    let simplified_self_ty = fast_reject::simplify_type(
+                        self.tcx,
+                        trait_ref.self_ty(),
+                        SimplifyParams::No,
+                        StripReferences::No,
+                    );
 
-                self.impls
-                    .entry(trait_ref.def_id)
-                    .or_default()
-                    .push((item.def_id.local_def_index, simplified_self_ty));
+                    self.impls
+                        .entry(trait_ref.def_id)
+                        .or_default()
+                        .push((item.def_id.local_def_index, simplified_self_ty));
+                }
             }
+            _ => {}
         }
     }
 

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -222,6 +222,7 @@ crate struct CrateRoot<'tcx> {
     diagnostic_items: Lazy<[(Symbol, DefIndex)]>,
     native_libraries: Lazy<[NativeLib]>,
     foreign_modules: Lazy<[ForeignModule]>,
+    traits: Lazy<[DefIndex]>,
     impls: Lazy<[TraitImpls]>,
     interpret_alloc_index: Lazy<[u32]>,
     proc_macro_data: Option<ProcMacroData>,

compiler/rustc_middle/src/query/mod.rs

@@ -1609,11 +1609,11 @@ rustc_queries! {
         desc { "fetching all foreign CrateNum instances" }
     }
 
-    /// A vector of every trait accessible in the whole crate
-    /// (i.e., including those from subcrates). This is used only for
-    /// error reporting.
-    query all_traits(_: ()) -> &'tcx [DefId] {
-        desc { "fetching all foreign and local traits" }
+    /// A list of all traits in a crate, used by rustdoc and error reporting.
+    /// NOTE: Not named just `traits` due to a naming conflict.
+    query traits_in_crate(_: CrateNum) -> &'tcx [DefId] {
+        desc { "fetching all traits in a crate" }
+        separate_provide_extern
     }
 
     /// The list of symbols exported from the given crate.

compiler/rustc_middle/src/ty/context.rs

@@ -1577,6 +1577,12 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn const_eval_limit(self) -> Limit {
         self.limits(()).const_eval_limit
     }
+
+    pub fn all_traits(self) -> impl Iterator<Item = DefId> + 'tcx {
+        iter::once(LOCAL_CRATE)
+            .chain(self.crates(()).iter().copied())
+            .flat_map(move |cnum| self.traits_in_crate(cnum).iter().copied())
+    }
 }
 
 /// A trait implemented for all `X<'a>` types that can be safely and

compiler/rustc_trait_selection/src/traits/error_reporting/mod.rs

@@ -1567,14 +1567,14 @@ impl<'a, 'tcx> InferCtxtPrivExt<'a, 'tcx> for InferCtxt<'a, 'tcx> {
             self.tcx.find_map_relevant_impl(trait_def_id, trait_ref.skip_binder().self_ty(), Some)
         };
         let required_trait_path = self.tcx.def_path_str(trait_ref.def_id());
-        let all_traits = self.tcx.all_traits(());
-        let traits_with_same_path: std::collections::BTreeSet<_> = all_traits
-            .iter()
-            .filter(|trait_def_id| **trait_def_id != trait_ref.def_id())
-            .filter(|trait_def_id| self.tcx.def_path_str(**trait_def_id) == required_trait_path)
+        let traits_with_same_path: std::collections::BTreeSet<_> = self
+            .tcx
+            .all_traits()
+            .filter(|trait_def_id| *trait_def_id != trait_ref.def_id())
+            .filter(|trait_def_id| self.tcx.def_path_str(*trait_def_id) == required_trait_path)
             .collect();
         for trait_with_same_path in traits_with_same_path {
-            if let Some(impl_def_id) = get_trait_impl(*trait_with_same_path) {
+            if let Some(impl_def_id) = get_trait_impl(trait_with_same_path) {
                 let impl_span = self.tcx.def_span(impl_def_id);
                 err.span_help(impl_span, "trait impl with same name found");
                 let trait_crate = self.tcx.crate_name(trait_with_same_path.krate);

compiler/rustc_typeck/src/check/method/mod.rs

@@ -7,7 +7,7 @@ mod prelude2021;
 pub mod probe;
 mod suggest;
 
-pub use self::suggest::{SelfSource, TraitInfo};
+pub use self::suggest::SelfSource;
 pub use self::CandidateSource::*;
 pub use self::MethodError::*;
 
@@ -31,7 +31,6 @@ use rustc_trait_selection::traits::query::evaluate_obligation::InferCtxtExt;
 use self::probe::{IsSuggestion, ProbeScope};
 
 pub fn provide(providers: &mut ty::query::Providers) {
-    suggest::provide(providers);
     probe::provide(providers);
 }
 

compiler/rustc_typeck/src/check/method/suggest.rs

@@ -5,8 +5,8 @@ use crate::check::FnCtxt;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_errors::{pluralize, struct_span_err, Applicability, DiagnosticBuilder};
 use rustc_hir as hir;
-use rustc_hir::def::{DefKind, Namespace, Res};
-use rustc_hir::def_id::{DefId, LocalDefId, CRATE_DEF_INDEX};
+use rustc_hir::def::Namespace;
+use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::lang_items::LangItem;
 use rustc_hir::{ExprKind, Node, QPath};
 use rustc_infer::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
@@ -1922,76 +1922,10 @@ impl Ord for TraitInfo {
     }
 }
 
-/// Retrieves all traits in this crate and any dependent crates.
+/// Retrieves all traits in this crate and any dependent crates,
+/// and wraps them into `TraitInfo` for custom sorting.
 pub fn all_traits(tcx: TyCtxt<'_>) -> Vec<TraitInfo> {
-    tcx.all_traits(()).iter().map(|&def_id| TraitInfo { def_id }).collect()
-}
-
-/// Computes all traits in this crate and any dependent crates.
-fn compute_all_traits(tcx: TyCtxt<'_>, (): ()) -> &[DefId] {
-    use hir::itemlikevisit;
-
-    let mut traits = vec![];
-
-    // Crate-local:
-
-    struct Visitor<'a> {
-        traits: &'a mut Vec<DefId>,
-    }
-
-    impl<'v, 'a> itemlikevisit::ItemLikeVisitor<'v> for Visitor<'a> {
-        fn visit_item(&mut self, i: &'v hir::Item<'v>) {
-            match i.kind {
-                hir::ItemKind::Trait(..) | hir::ItemKind::TraitAlias(..) => {
-                    self.traits.push(i.def_id.to_def_id());
-                }
-                _ => (),
-            }
-        }
-
-        fn visit_trait_item(&mut self, _trait_item: &hir::TraitItem<'_>) {}
-
-        fn visit_impl_item(&mut self, _impl_item: &hir::ImplItem<'_>) {}
-
-        fn visit_foreign_item(&mut self, _foreign_item: &hir::ForeignItem<'_>) {}
-    }
-
-    tcx.hir().visit_all_item_likes(&mut Visitor { traits: &mut traits });
-
-    // Cross-crate:
-
-    let mut external_mods = FxHashSet::default();
-    fn handle_external_res(
-        tcx: TyCtxt<'_>,
-        traits: &mut Vec<DefId>,
-        external_mods: &mut FxHashSet<DefId>,
-        res: Res<!>,
-    ) {
-        match res {
-            Res::Def(DefKind::Trait | DefKind::TraitAlias, def_id) => {
-                traits.push(def_id);
-            }
-            Res::Def(DefKind::Mod, def_id) => {
-                if !external_mods.insert(def_id) {
-                    return;
-                }
-                for child in tcx.item_children(def_id).iter() {
-                    handle_external_res(tcx, traits, external_mods, child.res)
-                }
-            }
-            _ => {}
-        }
-    }
-    for &cnum in tcx.crates(()).iter() {
-        let def_id = DefId { krate: cnum, index: CRATE_DEF_INDEX };
-        handle_external_res(tcx, &mut traits, &mut external_mods, Res::Def(DefKind::Mod, def_id));
-    }
-
-    tcx.arena.alloc_from_iter(traits)
-}
-
-pub fn provide(providers: &mut ty::query::Providers) {
-    providers.all_traits = compute_all_traits;
+    tcx.all_traits().map(|def_id| TraitInfo { def_id }).collect()
 }
 
 fn find_use_placement<'tcx>(tcx: TyCtxt<'tcx>, target_module: LocalDefId) -> (Option<Span>, bool) {

src/librustdoc/clean/blanket_impl.rs

@@ -19,7 +19,7 @@ impl<'a, 'tcx> BlanketImplFinder<'a, 'tcx> {
 
         trace!("get_blanket_impls({:?})", ty);
         let mut impls = Vec::new();
-        for &trait_def_id in self.cx.tcx.all_traits(()).iter() {
+        for trait_def_id in self.cx.tcx.all_traits() {
             if !self.cx.cache.access_levels.is_public(trait_def_id)
                 || self.cx.generated_synthetics.get(&(ty, trait_def_id)).is_some()
             {

src/librustdoc/core.rs

@@ -369,10 +369,8 @@ crate fn run_global_ctxt(
         impl_trait_bounds: Default::default(),
         generated_synthetics: Default::default(),
         auto_traits: tcx
-            .all_traits(())
-            .iter()
-            .cloned()
-            .filter(|trait_def_id| tcx.trait_is_auto(*trait_def_id))
+            .all_traits()
+            .filter(|&trait_def_id| tcx.trait_is_auto(trait_def_id))
             .collect(),
         module_trait_cache: FxHashMap::default(),
         cache: Cache::new(access_levels, render_options.document_private),

src/librustdoc/passes/collect_trait_impls.rs

@@ -129,7 +129,7 @@ crate fn collect_trait_impls(mut krate: Crate, cx: &mut DocContext<'_>) -> Crate
     // `tcx.crates(())` doesn't include the local crate, and `tcx.all_trait_implementations`
     // doesn't work with it anyway, so pull them from the HIR map instead
     let mut extra_attrs = Vec::new();
-    for &trait_did in cx.tcx.all_traits(()).iter() {
+    for trait_did in cx.tcx.all_traits() {
         for &impl_did in cx.tcx.hir().trait_impls(trait_did) {
             let impl_did = impl_did.to_def_id();
             cx.tcx.sess.prof.generic_activity("build_local_trait_impl").run(|| {