querify getting fake doc items

Cargo.lock

@@ -4604,6 +4604,7 @@ dependencies = [
  "rustc_abi",
  "rustc_ast",
  "rustc_ast_lowering",
+ "rustc_attr_ir",
  "rustc_attr_parsing",
  "rustc_crate_store",
  "rustc_data_structures",

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1308,6 +1308,17 @@ impl CrateMetadata {
         canonical_symbols
     }
 
+    /// Iterates over the fake_doc_items in the given crate.
+    fn get_fake_doc_items(&self, tcx: TyCtxt<'_>) -> Vec<DefId> {
+        let mut fake_doc_items = Vec::new();
+
+        for def_index in self.root.fake_doc_items.decode((self, tcx)) {
+            let id = self.local_def_id(def_index);
+            fake_doc_items.push(id);
+        }
+
+        fake_doc_items
+    }
     fn get_mod_child(&self, tcx: TyCtxt<'_>, id: DefIndex) -> ModChild {
         let ident = self.item_ident(tcx, id);
         let res = Res::Def(self.def_kind(id), self.local_def_id(id));

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

@@ -402,6 +402,7 @@ provide! { tcx, def_id, other, cdata,
     defined_lang_items => { cdata.get_lang_items(tcx) }
     diagnostic_items => { cdata.get_diagnostic_items(tcx) }
     canonical_symbols => { cdata.get_canonical_symbols(tcx) }
+    fake_doc_items => { cdata.get_fake_doc_items(tcx) }
     missing_lang_items => { cdata.get_missing_lang_items(tcx) }
 
     missing_extern_crate_item => {

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -642,6 +642,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
 
         let canonical_symbols = stat!("canonical-symbols", || self.encode_canonical_symbols());
 
+        let fake_doc_items = stat!("fake-doc-items", || self.encode_fake_doc_items());
+
         let native_libraries = stat!("native-libs", || self.encode_native_libraries());
 
         let foreign_modules = stat!("foreign-modules", || self.encode_foreign_modules());
@@ -761,6 +763,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 lang_items,
                 diagnostic_items,
                 canonical_symbols,
+                fake_doc_items,
                 lang_items_missing,
                 stripped_cfg_items,
                 native_libraries,
@@ -2169,6 +2172,13 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         self.lazy_array(diagnostic_items.iter().map(|(&name, def_id)| (name, def_id.index)))
     }
 
+    fn encode_fake_doc_items(&mut self) -> LazyArray<DefIndex> {
+        empty_proc_macro!(self);
+        let tcx = self.tcx;
+        let fake_doc_items = &tcx.fake_doc_items(LOCAL_CRATE);
+        self.lazy_array(fake_doc_items.iter().map(|cs| cs.index))
+    }
+
     fn encode_lang_items(&mut self) -> LazyArray<(DefIndex, LangItem)> {
         empty_proc_macro!(self);
         let lang_items = self.tcx.lang_items().iter();

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -270,6 +270,7 @@ pub(crate) struct CrateRoot {
     stripped_cfg_items: LazyArray<StrippedCfgItem<DefIndex>>,
     diagnostic_items: LazyArray<(Symbol, DefIndex)>,
     canonical_symbols: LazyArray<(Symbol, DefIndex)>,
+    fake_doc_items: LazyArray<DefIndex>,
     native_libraries: LazyArray<NativeLib>,
     foreign_modules: LazyArray<ForeignModule>,
     traits: LazyArray<DefIndex>,

compiler/rustc_middle/src/queries.rs

@@ -2835,6 +2835,20 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    /// Returns the fake doc items defined in a crate's root.
+    query fake_doc_items(_: CrateNum) -> &'tcx Vec<DefId> {
+        arena_cache
+        desc { "calculating the fake doc items" }
+        separate_provide_extern
+    }
+
+    /// Returns all fake doc items defined in all crates' roots.
+    query all_fake_doc_items(_: ()) -> &'tcx Vec<DefId> {
+        arena_cache
+        eval_always
+        desc { "calculating all fake doc items" }
+    }
+
     //-----------------------------------------------------------------------------
     // "Non-queries" are special dep kinds that are not queries.
     //-----------------------------------------------------------------------------

compiler/rustc_passes/Cargo.toml

@@ -8,6 +8,7 @@ edition = "2024"
 rustc_abi = { path = "../rustc_abi" }
 rustc_ast = { path = "../rustc_ast" }
 rustc_ast_lowering = { path = "../rustc_ast_lowering" }
+rustc_attr_ir = { path = "../rustc_attr_ir" }
 rustc_attr_parsing = { path = "../rustc_attr_parsing" }
 rustc_crate_store = { path = "../rustc_crate_store" }
 rustc_data_structures = { path = "../rustc_data_structures" }

compiler/rustc_passes/src/fake_doc_items.rs

@@ -0,0 +1,55 @@
+//! Collecting fake doc items.
+//!
+use rustc_attr_ir::{DocAttribute, find_attr};
+use rustc_middle::query::{LocalCrate, Providers};
+use rustc_middle::ty::TyCtxt;
+use rustc_span::def_id::{DefId, LOCAL_CRATE};
+use rustc_span::sym;
+
+/// Traverse and collect the fake doc items in the current crate
+fn fake_doc_items(tcx: TyCtxt<'_>, _: LocalCrate) -> Vec<DefId> {
+    let mut fake_doc_items = Vec::new();
+
+    // Optimization: can this crate even define fake doc items?
+    let features = tcx.features().enabled_features();
+    if features.contains(&sym::rustc_attrs) || features.contains(&sym::rustdoc_internals) {
+        // Collect fake doc items in this crate.
+        for id in tcx.hir_root_module().item_ids {
+            let id = id.hir_id();
+            if find_attr!(
+                tcx,
+                id,
+                RustcDocPrimitive(..)
+                    | Doc(DocAttribute { keyword: Some(..), .. })
+                    | Doc(DocAttribute { attribute: Some(..), .. })
+            ) {
+                fake_doc_items.push(id.expect_owner().to_def_id());
+            }
+        }
+    }
+
+    fake_doc_items
+}
+
+/// Traverse and collect all the fake doc items in all crates.
+fn all_fake_doc_items(tcx: TyCtxt<'_>, (): ()) -> Vec<DefId> {
+    let mut fake_doc_items = Vec::new();
+
+    // Collect fake doc items in visible crates.
+    for cnum in tcx
+        .crates(())
+        .iter()
+        .copied()
+        .filter(|cnum| tcx.is_user_visible_dep(*cnum))
+        .chain(std::iter::once(LOCAL_CRATE))
+    {
+        fake_doc_items.extend_from_slice(tcx.fake_doc_items(cnum))
+    }
+
+    fake_doc_items
+}
+
+pub(crate) fn provide(providers: &mut Providers) {
+    providers.fake_doc_items = fake_doc_items;
+    providers.all_fake_doc_items = all_fake_doc_items;
+}

compiler/rustc_passes/src/lib.rs

@@ -5,6 +5,7 @@
 //! This API is completely unstable and subject to change.
 
 // tidy-alphabetical-start
+#![feature(deref_patterns)]
 #![feature(option_into_flat_iter)]
 // tidy-alphabetical-end
 
@@ -21,6 +22,7 @@ mod diagnostic_items;
 mod diagnostics;
 mod eii;
 pub mod entry;
+mod fake_doc_items;
 pub mod hir_id_validator;
 pub mod input_stats;
 mod lang_items;
@@ -44,5 +46,6 @@ pub fn provide(providers: &mut Providers) {
     stability::provide(providers);
     upvars::provide(providers);
     check_export::provide(providers);
+    fake_doc_items::provide(providers);
     providers.check_externally_implementable_items = eii::check_externally_implementable_items;
 }

src/librustdoc/clean/types.rs

@@ -236,35 +236,12 @@ impl ExternalCrate {
             .unwrap_or(Unknown) // Well, at least we tried.
     }
 
-    fn mapped_root_anon_consts<T>(
+    fn fake_doc_items<T>(
         &self,
         tcx: TyCtxt<'_>,
         f: impl Fn(DefId, TyCtxt<'_>) -> Option<(DefId, T)>,
     ) -> impl Iterator<Item = (DefId, T)> {
-        let root = self.def_id();
-
-        if root.is_local() {
-            Either::Left(
-                tcx.hir_root_module()
-                    .item_ids
-                    .iter()
-                    .filter(move |&&id| matches!(tcx.hir_item(id).kind, hir::ItemKind::Const(..)))
-                    .filter_map(move |&id| f(id.owner_id.into(), tcx)),
-            )
-        } else {
-            Either::Right(
-                tcx.module_children(root)
-                    .iter()
-                    .filter_map(|item| {
-                        if let Res::Def(DefKind::Const { is_type_const: false }, did) = item.res {
-                            Some(did)
-                        } else {
-                            None
-                        }
-                    })
-                    .filter_map(move |did| f(did, tcx)),
-            )
-        }
+        tcx.fake_doc_items(self.crate_num).into_iter().filter_map(move |did| f(*did, tcx))
     }
 
     pub(crate) fn keywords(&self, tcx: TyCtxt<'_>) -> impl Iterator<Item = (DefId, Symbol)> {
@@ -285,7 +262,7 @@ impl ExternalCrate {
         let as_target = move |did: DefId, tcx: TyCtxt<'_>| -> Option<(DefId, Symbol)> {
             find_attr!(tcx, did, Doc(d) => callback(d)).flatten().map(|value| (did, value))
         };
-        self.mapped_root_anon_consts(tcx, as_target)
+        self.fake_doc_items(tcx, as_target)
     }
 
     pub(crate) fn primitives(
@@ -320,7 +297,7 @@ impl ExternalCrate {
             Some((def_id, prim))
         }
 
-        self.mapped_root_anon_consts(tcx, as_primitive)
+        self.fake_doc_items(tcx, as_primitive)
     }
 }
 
@@ -1895,27 +1872,35 @@ impl PrimitiveType {
     /// `rustc_doc_primitive`, then it's entirely random whether `std` or the other crate is picked.
     /// (no_std crates are usually fine unless multiple dependencies define a primitive.)
     pub(crate) fn primitive_locations(tcx: TyCtxt<'_>) -> &FxIndexMap<PrimitiveType, DefId> {
+        fn as_primitive(def_id: DefId, tcx: TyCtxt<'_>) -> Option<PrimitiveType> {
+            let (attr_span, prim_sym) = find_attr!(
+                tcx, def_id,
+                RustcDocPrimitive(span, prim) => (*span, *prim)
+            )?;
+            let Some(prim) = PrimitiveType::from_symbol(prim_sym) else {
+                span_bug!(attr_span, "primitive `{prim_sym}` is not a member of `PrimitiveType`");
+            };
+            Some(prim)
+        }
+
         static PRIMITIVE_LOCATIONS: OnceCell<FxIndexMap<PrimitiveType, DefId>> = OnceCell::new();
         PRIMITIVE_LOCATIONS.get_or_init(|| {
             let mut primitive_locations = FxIndexMap::default();
             // NOTE: technically this misses crates that are only passed with `--extern` and not loaded when checking the crate.
             // This is a degenerate case that I don't plan to support.
-            for &crate_num in tcx.crates(()) {
-                let e = ExternalCrate { crate_num };
-                let crate_name = e.name(tcx);
-                debug!(?crate_num, ?crate_name);
-                for (def_id, prim) in e.primitives(tcx) {
-                    // HACK: try to link to std instead where possible
-                    if crate_name == sym::core && primitive_locations.contains_key(&prim) {
-                        continue;
-                    }
+
+            let mut ids = tcx.all_fake_doc_items(()).clone();
+
+            // Primitives are unhygienically duplicated by `include!`.
+            // Sort them with core first, so that if std is present in the crate graph,
+            // core's items are overridden and we link to std preferentially.
+            ids.iter_mut().partition_in_place(|id| tcx.crate_name(id.krate) == sym::core);
+            for def_id in ids {
+                if let Some(prim) = as_primitive(def_id, tcx) {
                     primitive_locations.insert(prim, def_id);
                 }
             }
-            let local_primitives = ExternalCrate { crate_num: LOCAL_CRATE }.primitives(tcx);
-            for (def_id, prim) in local_primitives {
-                primitive_locations.insert(prim, def_id);
-            }
+
             primitive_locations
         })
     }

src/librustdoc/lib.rs

@@ -10,6 +10,7 @@
 #![feature(formatting_options)]
 #![feature(iter_intersperse)]
 #![feature(iter_order_by)]
+#![feature(iter_partition_in_place)]
 #![feature(rustc_private)]
 #![feature(test)]
 #![feature(trim_prefix_suffix)]