From dd9cda6090c80915f98b9d1a383e6abcf1d86c65 Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Sun, 17 Mar 2024 09:37:52 +0100
Subject: [PATCH] avoid processing mentioned items that are also still used

---
 compiler/rustc_middle/src/mir/mod.rs          |  2 +-
 .../src/mentioned_items.rs                    |  3 ++-
 compiler/rustc_monomorphize/src/collector.rs  | 27 ++++++++++++++++---
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index 6c9d5c670825f..c1fba223cd2f3 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -314,7 +314,7 @@ impl<'tcx> CoroutineInfo<'tcx> {
 }
 
 /// Some item that needs to monomorphize successfully for a MIR body to be considered well-formed.
-#[derive(Copy, Clone, PartialEq, Debug, HashStable, TyEncodable, TyDecodable)]
+#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, HashStable, TyEncodable, TyDecodable)]
 #[derive(TypeFoldable, TypeVisitable)]
 pub enum MentionedItem<'tcx> {
     Fn(DefId, GenericArgsRef<'tcx>),
diff --git a/compiler/rustc_mir_transform/src/mentioned_items.rs b/compiler/rustc_mir_transform/src/mentioned_items.rs
index e05c925acd739..437037aa0661d 100644
--- a/compiler/rustc_mir_transform/src/mentioned_items.rs
+++ b/compiler/rustc_mir_transform/src/mentioned_items.rs
@@ -33,7 +33,8 @@ impl<'tcx> MirPass<'tcx> for MentionedItems {
 impl<'tcx> Visitor<'tcx> for MentionedItemsVisitor<'_, 'tcx> {
     fn visit_constant(&mut self, constant: &ConstOperand<'tcx>, _: Location) {
         let const_ = constant.const_;
-        // This is how function items get referenced: via constants of `FnDef` type.
+        // This is how function items get referenced: via constants of `FnDef` type. This handles
+        // both functions that are called anf those that are just turned to function pointers.
         if let ty::FnDef(def_id, args) = const_.ty().kind() {
             debug!("adding to required_items: {def_id:?}");
             self.mentioned_items
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 75460c9bbdc30..78d5e1f482e4f 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -718,6 +718,8 @@ struct MirUsedCollector<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     body: &'a mir::Body<'tcx>,
     used_items: &'a mut MonoItems<'tcx>,
+    /// See the comment in `collect_items_of_instance` for the purpose of this set.
+    used_mentioned_items: &'a mut FxHashSet<MentionedItem<'tcx>>,
     instance: Instance<'tcx>,
     /// Spans for move size lints already emitted. Helps avoid duplicate lints.
     move_size_spans: Vec<Span>,
@@ -988,6 +990,10 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
         match terminator.kind {
             mir::TerminatorKind::Call { ref func, ref args, ref fn_span, .. } => {
                 let callee_ty = func.ty(self.body, tcx);
+                // *Before* monomorphizing, record that we already handled this mention.
+                if let ty::FnDef(def_id, args) = callee_ty.kind() {
+                    self.used_mentioned_items.insert(MentionedItem::Fn(*def_id, args));
+                }
                 let callee_ty = self.monomorphize(callee_ty);
                 self.check_fn_args_move_size(callee_ty, args, *fn_span, location);
                 visit_fn_use(self.tcx, callee_ty, true, source, &mut self.used_items)
@@ -1633,10 +1639,22 @@ fn collect_items_of_instance<'tcx>(
     mode: CollectionMode,
 ) {
     let body = tcx.instance_mir(instance.def);
+    // Naively, in "used" collection mode, all functions get added to *both* `used_items` and
+    // `mentioned_items`. Mentioned items processing will then notice that they have already been
+    // visited, but at that point each mentioned item has been monomorphized, added to the
+    // `mentioned_items` worklist, and checked in the global set of visited items. To removes that
+    // overhead, we have a special optimization that avoids adding items to `mentioned_items` when
+    // they are already added in `used_items`. We could just scan `used_items`, but that's a linear
+    // scan and not very efficient. Furthermore we can only do that *after* monomorphizing the
+    // mentioned item. So instead we collect all pre-monomorphized `MentionedItem` that were already
+    // added to `used_items` in a hash set, which can efficiently query in the
+    // `body.mentioned_items` loop below.
+    let mut used_mentioned_items = FxHashSet::<MentionedItem<'tcx>>::default();
     let mut collector = MirUsedCollector {
         tcx,
         body,
         used_items,
+        used_mentioned_items: &mut used_mentioned_items,
         instance,
         move_size_spans: vec![],
         visiting_call_terminator: false,
@@ -1656,10 +1674,13 @@ fn collect_items_of_instance<'tcx>(
         }
     }
 
-    // Always gather mentioned items.
+    // Always gather mentioned items. We try to avoid processing items that we have already added to
+    // `used_items` above.
     for item in &body.mentioned_items {
-        let item_mono = collector.monomorphize(item.node);
-        visit_mentioned_item(tcx, &item_mono, item.span, mentioned_items);
+        if !collector.used_mentioned_items.contains(&item.node) {
+            let item_mono = collector.monomorphize(item.node);
+            visit_mentioned_item(tcx, &item_mono, item.span, mentioned_items);
+        }
     }
 }