From aca0688a1dd16e0bc41ac41c5b861384c418427e Mon Sep 17 00:00:00 2001 From: Yotam Ofek Date: Thu, 19 Jun 2025 13:50:54 +0000 Subject: [PATCH] De-dup common code from `ExternalCrate` methods --- src/librustdoc/clean/types.rs | 118 ++++++++----------- src/librustdoc/clean/utils.rs | 4 +- src/librustdoc/passes/collect_trait_impls.rs | 2 +- 3 files changed, 53 insertions(+), 71 deletions(-) diff --git a/src/librustdoc/clean/types.rs b/src/librustdoc/clean/types.rs index 2d9670a3d10cd..408ef611ee527 100644 --- a/src/librustdoc/clean/types.rs +++ b/src/librustdoc/clean/types.rs @@ -4,6 +4,7 @@ use std::sync::{Arc, OnceLock as OnceCell}; use std::{fmt, iter}; use arrayvec::ArrayVec; +use itertools::Either; use rustc_abi::{ExternAbi, VariantIdx}; use rustc_attr_data_structures::{ AttributeKind, ConstStability, Deprecation, Stability, StableSince, @@ -199,49 +200,49 @@ impl ExternalCrate { .unwrap_or(Unknown) // Well, at least we tried. } - pub(crate) fn keywords(&self, tcx: TyCtxt<'_>) -> ThinVec<(DefId, Symbol)> { + fn mapped_root_modules( + &self, + tcx: TyCtxt<'_>, + f: impl Fn(DefId, TyCtxt<'_>) -> Option<(DefId, T)>, + ) -> impl Iterator { let root = self.def_id(); - let as_keyword = |res: Res| { - if let Res::Def(DefKind::Mod, def_id) = res { - let mut keyword = None; - let meta_items = tcx - .get_attrs(def_id, sym::doc) - .flat_map(|attr| attr.meta_item_list().unwrap_or_default()); - for meta in meta_items { - if meta.has_name(sym::keyword) - && let Some(v) = meta.value_str() - { - keyword = Some(v); - break; - } - } - return keyword.map(|p| (def_id, p)); - } - None - }; if root.is_local() { - tcx.hir_root_module() - .item_ids - .iter() - .filter_map(|&id| { - let item = tcx.hir_item(id); - match item.kind { - hir::ItemKind::Mod(..) => { - as_keyword(Res::Def(DefKind::Mod, id.owner_id.to_def_id())) - } - _ => None, - } - }) - .collect() + Either::Left( + tcx.hir_root_module() + .item_ids + .iter() + .filter(move |&&id| matches!(tcx.hir_item(id).kind, hir::ItemKind::Mod(..))) + .filter_map(move |&id| f(id.owner_id.into(), tcx)), + ) } else { - tcx.module_children(root).iter().map(|item| item.res).filter_map(as_keyword).collect() + Either::Right( + tcx.module_children(root) + .iter() + .filter_map(|item| { + if let Res::Def(DefKind::Mod, did) = item.res { Some(did) } else { None } + }) + .filter_map(move |did| f(did, tcx)), + ) } } - pub(crate) fn primitives(&self, tcx: TyCtxt<'_>) -> ThinVec<(DefId, PrimitiveType)> { - let root = self.def_id(); + pub(crate) fn keywords(&self, tcx: TyCtxt<'_>) -> impl Iterator { + fn as_keyword(did: DefId, tcx: TyCtxt<'_>) -> Option<(DefId, Symbol)> { + tcx.get_attrs(did, sym::doc) + .flat_map(|attr| attr.meta_item_list().unwrap_or_default()) + .filter(|meta| meta.has_name(sym::keyword)) + .find_map(|meta| meta.value_str()) + .map(|value| (did, value)) + } + self.mapped_root_modules(tcx, as_keyword) + } + + pub(crate) fn primitives( + &self, + tcx: TyCtxt<'_>, + ) -> impl Iterator { // Collect all inner modules which are tagged as implementations of // primitives. // @@ -259,40 +260,21 @@ impl ExternalCrate { // Also note that this does not attempt to deal with modules tagged // duplicately for the same primitive. This is handled later on when // rendering by delegating everything to a hash map. - let as_primitive = |res: Res| { - let Res::Def(DefKind::Mod, def_id) = res else { return None }; - tcx.get_attrs(def_id, sym::rustc_doc_primitive) - .map(|attr| { - let attr_value = attr.value_str().expect("syntax should already be validated"); - let Some(prim) = PrimitiveType::from_symbol(attr_value) else { - span_bug!( - attr.span(), - "primitive `{attr_value}` is not a member of `PrimitiveType`" - ); - }; - - (def_id, prim) - }) - .next() - }; + fn as_primitive(def_id: DefId, tcx: TyCtxt<'_>) -> Option<(DefId, PrimitiveType)> { + tcx.get_attrs(def_id, sym::rustc_doc_primitive).next().map(|attr| { + let attr_value = attr.value_str().expect("syntax should already be validated"); + let Some(prim) = PrimitiveType::from_symbol(attr_value) else { + span_bug!( + attr.span(), + "primitive `{attr_value}` is not a member of `PrimitiveType`" + ); + }; - if root.is_local() { - tcx.hir_root_module() - .item_ids - .iter() - .filter_map(|&id| { - let item = tcx.hir_item(id); - match item.kind { - hir::ItemKind::Mod(..) => { - as_primitive(Res::Def(DefKind::Mod, id.owner_id.to_def_id())) - } - _ => None, - } - }) - .collect() - } else { - tcx.module_children(root).iter().map(|item| item.res).filter_map(as_primitive).collect() + (def_id, prim) + }) } + + self.mapped_root_modules(tcx, as_primitive) } } @@ -1966,7 +1948,7 @@ impl PrimitiveType { let e = ExternalCrate { crate_num }; let crate_name = e.name(tcx); debug!(?crate_num, ?crate_name); - for &(def_id, prim) in &e.primitives(tcx) { + 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; diff --git a/src/librustdoc/clean/utils.rs b/src/librustdoc/clean/utils.rs index c58b07a5b6731..2c9878636abf4 100644 --- a/src/librustdoc/clean/utils.rs +++ b/src/librustdoc/clean/utils.rs @@ -61,7 +61,7 @@ pub(crate) fn krate(cx: &mut DocContext<'_>) -> Crate { let keywords = local_crate.keywords(cx.tcx); { let ItemKind::ModuleItem(m) = &mut module.inner.kind else { unreachable!() }; - m.items.extend(primitives.iter().map(|&(def_id, prim)| { + m.items.extend(primitives.map(|(def_id, prim)| { Item::from_def_id_and_parts( def_id, Some(prim.as_sym()), @@ -69,7 +69,7 @@ pub(crate) fn krate(cx: &mut DocContext<'_>) -> Crate { cx, ) })); - m.items.extend(keywords.into_iter().map(|(def_id, kw)| { + m.items.extend(keywords.map(|(def_id, kw)| { Item::from_def_id_and_parts(def_id, Some(kw), ItemKind::KeywordItem, cx) })); } diff --git a/src/librustdoc/passes/collect_trait_impls.rs b/src/librustdoc/passes/collect_trait_impls.rs index f4e4cd924f7ff..2339a6b69cd8a 100644 --- a/src/librustdoc/passes/collect_trait_impls.rs +++ b/src/librustdoc/passes/collect_trait_impls.rs @@ -35,7 +35,7 @@ pub(crate) fn collect_trait_impls(mut krate: Crate, cx: &mut DocContext<'_>) -> }); let local_crate = ExternalCrate { crate_num: LOCAL_CRATE }; - let prims: FxHashSet = local_crate.primitives(tcx).iter().map(|p| p.1).collect(); + let prims: FxHashSet = local_crate.primitives(tcx).map(|(_, p)| p).collect(); let crate_items = { let mut coll = ItemAndAliasCollector::new(&cx.cache);