De-dup common code from `ExternalCrate` methods

3 files changed, 53 insertions, 71 deletions

diff --git a/src/librustdoc/clean/types.rs b/src/librustdoc/clean/types.rs
index 2d9670a3d10..408ef611ee5 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<T>(
+        &self,
+        tcx: TyCtxt<'_>,
+        f: impl Fn(DefId, TyCtxt<'_>) -> Option<(DefId, T)>,
+    ) -> impl Iterator<Item = (DefId, T)> {
         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<Item = (DefId, Symbol)> {
+        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<Item = (DefId, PrimitiveType)> {
         // 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 c58b07a5b67..2c9878636ab 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 f4e4cd924f7..2339a6b69cd 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<PrimitiveType> = local_crate.primitives(tcx).iter().map(|p| p.1).collect();
+    let prims: FxHashSet<PrimitiveType> = local_crate.primitives(tcx).map(|(_, p)| p).collect();
 
     let crate_items = {
         let mut coll = ItemAndAliasCollector::new(&cx.cache);