mono collector: Reduce \# of locking while walking the graph

While profiling Zed's dev build I've noticed that while most of the time `upstream_monomorphizations` takes a lot of time in monomorpization_collector, in some cases (e.g. build of `editor` itself)
the rest of monomorphization_collector_graph_walk dominates it. Most of the time is spent in collect_items_rec.

This PR aims to reduce the number of locks taking place; instead of locking output MonoItems once per children of current node, we do so once per *parent*. We also get to reuse locks for mentioned and used items.
While this commit does not reduce Wall time of Zed's build, it does shave off `cargo build -j1` from 43s to 41.5s.

2 files changed, 38 insertions, 27 deletions

diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 76dad6b3571..1e3744e19f5 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -205,6 +205,7 @@
 //! this is not implemented however: a mono item will be produced
 //! regardless of whether it is actually needed or not.
 
+use std::cell::OnceCell;
 use std::path::PathBuf;
 
 use rustc_attr_parsing::InlineAttr;
@@ -348,6 +349,27 @@ impl<'tcx> Extend<Spanned<MonoItem<'tcx>>> for MonoItems<'tcx> {
     }
 }
 
+fn collect_items_root<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    starting_item: Spanned<MonoItem<'tcx>>,
+    state: &SharedState<'tcx>,
+    recursion_limit: Limit,
+) {
+    if !state.visited.lock_mut().insert(starting_item.node) {
+        // We've been here already, no need to search again.
+        return;
+    }
+    let mut recursion_depths = DefIdMap::default();
+    collect_items_rec(
+        tcx,
+        starting_item,
+        state,
+        &mut recursion_depths,
+        recursion_limit,
+        CollectionMode::UsedItems,
+    );
+}
+
 /// Collect all monomorphized items reachable from `starting_point`, and emit a note diagnostic if a
 /// post-monomorphization error is encountered during a collection step.
 ///
@@ -362,24 +384,6 @@ fn collect_items_rec<'tcx>(
     recursion_limit: Limit,
     mode: CollectionMode,
 ) {
-    if mode == CollectionMode::UsedItems {
-        if !state.visited.lock_mut().insert(starting_item.node) {
-            // We've been here already, no need to search again.
-            return;
-        }
-    } else {
-        if state.visited.lock().contains(&starting_item.node) {
-            // We've already done a *full* visit on this one, no need to do the "mention" visit.
-            return;
-        }
-        if !state.mentioned.lock_mut().insert(starting_item.node) {
-            // We've been here already, no need to search again.
-            return;
-        }
-        // There's some risk that we first do a 'mention' visit and then a full visit. But there's no
-        // harm in that, the mention visit will trigger all the queries and the results are cached.
-    }
-
     let mut used_items = MonoItems::new();
     let mut mentioned_items = MonoItems::new();
     let recursion_depth_reset;
@@ -536,6 +540,20 @@ fn collect_items_rec<'tcx>(
         state.usage_map.lock_mut().record_used(starting_item.node, &used_items);
     }
 
+    {
+        let mut visited = OnceCell::default();
+        if mode == CollectionMode::UsedItems {
+            used_items
+                .items
+                .retain(|k, _| visited.get_mut_or_init(|| state.visited.lock_mut()).insert(*k));
+        }
+
+        let mut mentioned = OnceCell::default();
+        mentioned_items.items.retain(|k, _| {
+            !visited.get_or_init(|| state.visited.lock()).contains(k)
+                && mentioned.get_mut_or_init(|| state.mentioned.lock_mut()).insert(*k)
+        });
+    }
     if mode == CollectionMode::MentionedItems {
         assert!(used_items.is_empty(), "'mentioned' collection should never encounter used items");
     } else {
@@ -1689,15 +1707,7 @@ pub(crate) fn collect_crate_mono_items<'tcx>(
 
     tcx.sess.time("monomorphization_collector_graph_walk", || {
         par_for_each_in(roots, |root| {
-            let mut recursion_depths = DefIdMap::default();
-            collect_items_rec(
-                tcx,
-                dummy_spanned(*root),
-                &state,
-                &mut recursion_depths,
-                recursion_limit,
-                CollectionMode::UsedItems,
-            );
+            collect_items_root(tcx, dummy_spanned(*root), &state, recursion_limit);
         });
     });
 
diff --git a/compiler/rustc_monomorphize/src/lib.rs b/compiler/rustc_monomorphize/src/lib.rs
index 8469e0f17a6..1b484da698a 100644
--- a/compiler/rustc_monomorphize/src/lib.rs
+++ b/compiler/rustc_monomorphize/src/lib.rs
@@ -4,6 +4,7 @@
 #![feature(file_buffered)]
 #![feature(if_let_guard)]
 #![feature(impl_trait_in_assoc_type)]
+#![feature(once_cell_get_mut)]
 // tidy-alphabetical-end
 
 use rustc_hir::lang_items::LangItem;