Only collect mono items from reachable blocks

Author: saethlin

Cargo.lock

@@ -4374,6 +4374,7 @@ dependencies = [
  "rustc_errors",
  "rustc_fluent_macro",
  "rustc_hir",
+ "rustc_index",
  "rustc_macros",
  "rustc_middle",
  "rustc_session",

compiler/rustc_codegen_ssa/src/mir/mod.rs

@@ -257,20 +257,19 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
     // Apply debuginfo to the newly allocated locals.
     fx.debug_introduce_locals(&mut start_bx);
 
-    let reachable_blocks = mir.reachable_blocks_in_mono(cx.tcx(), instance);
-
     // The builders will be created separately for each basic block at `codegen_block`.
     // So drop the builder of `start_llbb` to avoid having two at the same time.
     drop(start_bx);
 
+    let reachable_blocks = cx.tcx().reachable_blocks(instance);
+
     // Codegen the body of each block using reverse postorder
     for (bb, _) in traversal::reverse_postorder(mir) {
         if reachable_blocks.contains(bb) {
             fx.codegen_block(bb);
         } else {
-            // This may have references to things we didn't monomorphize, so we
-            // don't actually codegen the body. We still create the block so
-            // terminators in other blocks can reference it without worry.
+            // We want to skip this block, because it's not reachable. But we still create
+            // the block so terminators in other blocks can reference it.
             fx.codegen_block_as_unreachable(bb);
         }
     }

compiler/rustc_interface/src/passes.rs

@@ -631,6 +631,7 @@ pub static DEFAULT_QUERY_PROVIDERS: LazyLock<Providers> = LazyLock::new(|| {
     rustc_lint::provide(providers);
     rustc_symbol_mangling::provide(providers);
     rustc_codegen_ssa::provide(providers);
+    rustc_middle::mir::traversal::provide(providers);
     *providers
 });
 

compiler/rustc_middle/src/mir/mod.rs

@@ -30,7 +30,6 @@ pub use rustc_ast::Mutability;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::graph::dominators::Dominators;
-use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_index::bit_set::BitSet;
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_serialize::{Decodable, Encodable};
@@ -687,57 +686,6 @@ impl<'tcx> Body<'tcx> {
         self.injection_phase.is_some()
     }
 
-    /// Finds which basic blocks are actually reachable for a specific
-    /// monomorphization of this body.
-    ///
-    /// This is allowed to have false positives; just because this says a block
-    /// is reachable doesn't mean that's necessarily true. It's thus always
-    /// legal for this to return a filled set.
-    ///
-    /// Regardless, the [`BitSet::domain_size`] of the returned set will always
-    /// exactly match the number of blocks in the body so that `contains`
-    /// checks can be done without worrying about panicking.
-    ///
-    /// This is mostly useful because it lets us skip lowering the `false` side
-    /// of `if <T as Trait>::CONST`, as well as `intrinsics::debug_assertions`.
-    pub fn reachable_blocks_in_mono(
-        &self,
-        tcx: TyCtxt<'tcx>,
-        instance: Instance<'tcx>,
-    ) -> BitSet<BasicBlock> {
-        let mut set = BitSet::new_empty(self.basic_blocks.len());
-        self.reachable_blocks_in_mono_from(tcx, instance, &mut set, START_BLOCK);
-        set
-    }
-
-    fn reachable_blocks_in_mono_from(
-        &self,
-        tcx: TyCtxt<'tcx>,
-        instance: Instance<'tcx>,
-        set: &mut BitSet<BasicBlock>,
-        bb: BasicBlock,
-    ) {
-        if !set.insert(bb) {
-            return;
-        }
-
-        let data = &self.basic_blocks[bb];
-
-        if let Some((bits, targets)) = Self::try_const_mono_switchint(tcx, instance, data) {
-            let target = targets.target_for_value(bits);
-            ensure_sufficient_stack(|| {
-                self.reachable_blocks_in_mono_from(tcx, instance, set, target)
-            });
-            return;
-        }
-
-        for target in data.terminator().successors() {
-            ensure_sufficient_stack(|| {
-                self.reachable_blocks_in_mono_from(tcx, instance, set, target)
-            });
-        }
-    }
-
     /// If this basic block ends with a [`TerminatorKind::SwitchInt`] for which we can evaluate the
     /// dimscriminant in monomorphization, we return the discriminant bits and the
     /// [`SwitchTargets`], just so the caller doesn't also have to match on the terminator.

compiler/rustc_middle/src/mir/traversal.rs

@@ -1,3 +1,5 @@
+use crate::query::Providers;
+
 use super::*;
 
 /// Preorder traversal of a graph.
@@ -279,3 +281,74 @@ pub fn reverse_postorder<'a, 'tcx>(
 {
     body.basic_blocks.reverse_postorder().iter().map(|&bb| (bb, &body.basic_blocks[bb]))
 }
+
+/// Finds which basic blocks are actually reachable for a monomorphized [`Instance`].
+///
+/// This is allowed to have false positives; just because this says a block
+/// is reachable doesn't mean that's necessarily true. It's thus always
+/// legal for this to return a filled set.
+///
+/// Regardless, the [`BitSet::domain_size`] of the returned set will always
+/// exactly match the number of blocks in the body so that `contains`
+/// checks can be done without worrying about panicking.
+///
+/// This is mostly useful because it lets us skip lowering the `false` side
+/// of `if <T as Trait>::CONST`, as well as [`NullOp::UbChecks`].
+///
+/// [`NullOp::UbChecks`]: rustc_middle::mir::NullOp::UbChecks
+fn reachable_blocks<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> BitSet<BasicBlock> {
+    let body = tcx.instance_mir(instance.def);
+    let mut visitor = MonoReachable {
+        body,
+        tcx,
+        instance,
+        worklist: Vec::new(),
+        visited: BitSet::new_empty(body.basic_blocks.len()),
+    };
+
+    visitor.visited.insert(START_BLOCK);
+    visitor.visit(START_BLOCK);
+
+    while let Some(bb) = visitor.worklist.pop() {
+        if visitor.visited.insert(bb) {
+            visitor.visit(bb);
+        }
+    }
+
+    visitor.visited
+}
+
+struct MonoReachable<'a, 'tcx> {
+    body: &'a Body<'tcx>,
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    worklist: Vec<BasicBlock>,
+    visited: BitSet<BasicBlock>,
+}
+
+impl<'a, 'tcx> MonoReachable<'a, 'tcx> {
+    fn visit(&mut self, bb: BasicBlock) {
+        let block = &self.body.basic_blocks[bb];
+
+        if let Some((bits, targets)) =
+            Body::try_const_mono_switchint(self.tcx, self.instance, block)
+        {
+            let target = targets.target_for_value(bits);
+            self.push(target);
+        } else {
+            for target in block.terminator().successors() {
+                self.push(target);
+            }
+        }
+    }
+
+    fn push(&mut self, bb: BasicBlock) {
+        if !self.visited.contains(bb) {
+            self.worklist.push(bb);
+        }
+    }
+}
+
+pub fn provide(providers: &mut Providers) {
+    providers.reachable_blocks = reachable_blocks;
+}

compiler/rustc_middle/src/query/keys.rs

@@ -87,6 +87,15 @@ impl<'tcx> Key for ty::Instance<'tcx> {
     }
 }
 
+impl<'tcx> AsLocalKey for ty::Instance<'tcx> {
+    type LocalKey = Self;
+
+    #[inline(always)]
+    fn as_local_key(&self) -> Option<Self::LocalKey> {
+        self.def_id().is_local().then(|| *self)
+    }
+}
+
 impl<'tcx> Key for mir::interpret::GlobalId<'tcx> {
     type Cache<V> = DefaultCache<Self, V>;
 
@@ -534,6 +543,14 @@ impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx ty::List<Ty<'tcx>>) {
     }
 }
 
+impl<'tcx> Key for (ty::Instance<'tcx>, &'tcx mir::Body<'tcx>) {
+    type Cache<V> = DefaultCache<Self, V>;
+
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.default_span(tcx)
+    }
+}
+
 impl<'tcx> Key for (Ty<'tcx>, ty::ValTree<'tcx>) {
     type Cache<V> = DefaultCache<Self, V>;
 

compiler/rustc_middle/src/query/mod.rs

@@ -71,6 +71,7 @@ use rustc_hir::def_id::{
 };
 use rustc_hir::lang_items::{LangItem, LanguageItems};
 use rustc_hir::{Crate, ItemLocalId, ItemLocalMap, TraitCandidate};
+use rustc_index::bit_set::BitSet;
 use rustc_index::IndexVec;
 use rustc_query_system::ich::StableHashingContext;
 use rustc_query_system::query::{try_get_cached, QueryCache, QueryMode, QueryState};
@@ -270,7 +271,7 @@ rustc_queries! {
         feedable
     }
 
-    query unsizing_params_for_adt(key: DefId) -> &'tcx rustc_index::bit_set::BitSet<u32>
+    query unsizing_params_for_adt(key: DefId) -> &'tcx BitSet<u32>
     {
         arena_cache
         desc { |tcx|
@@ -460,7 +461,7 @@ rustc_queries! {
     }
 
     /// Set of param indexes for type params that are in the type's representation
-    query params_in_repr(key: DefId) -> &'tcx rustc_index::bit_set::BitSet<u32> {
+    query params_in_repr(key: DefId) -> &'tcx BitSet<u32> {
         desc { "finding type parameters in the representation" }
         arena_cache
         no_hash
@@ -2251,6 +2252,11 @@ rustc_queries! {
     query find_field((def_id, ident): (DefId, rustc_span::symbol::Ident)) -> Option<rustc_target::abi::FieldIdx> {
         desc { |tcx| "find the index of maybe nested field `{ident}` in `{}`", tcx.def_path_str(def_id) }
     }
+
+    query reachable_blocks(instance: ty::Instance<'tcx>) -> &'tcx BitSet<mir::BasicBlock> {
+        arena_cache
+        desc { |tcx| "determining reachable blocks in `{}`", tcx.def_path_str(instance.def_id()) }
+    }
 }
 
 rustc_query_append! { define_callbacks! }

compiler/rustc_monomorphize/Cargo.toml

@@ -9,6 +9,7 @@ rustc_data_structures = { path = "../rustc_data_structures" }
 rustc_errors = { path = "../rustc_errors" }
 rustc_fluent_macro = { path = "../rustc_fluent_macro" }
 rustc_hir = { path = "../rustc_hir" }
+rustc_index = { path = "../rustc_index" }
 rustc_macros = { path = "../rustc_macros" }
 rustc_middle = { path = "../rustc_middle" }
 rustc_session = { path = "../rustc_session" }

compiler/rustc_monomorphize/src/collector.rs

@@ -211,6 +211,7 @@ use rustc_hir as hir;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::lang_items::LangItem;
+use rustc_index::bit_set::BitSet;
 use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir::interpret::{AllocId, ErrorHandled, GlobalAlloc, Scalar};
 use rustc_middle::mir::mono::{InstantiationMode, MonoItem};
@@ -671,6 +672,7 @@ struct MirUsedCollector<'a, 'tcx> {
     visiting_call_terminator: bool,
     /// Set of functions for which it is OK to move large data into.
     skip_move_check_fns: Option<Vec<DefId>>,
+    reachable_blocks: Option<&'tcx BitSet<mir::BasicBlock>>,
 }
 
 impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
@@ -831,6 +833,16 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
 }
 
 impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
+    fn visit_basic_block_data(&mut self, block: mir::BasicBlock, data: &mir::BasicBlockData<'tcx>) {
+        if self
+            .reachable_blocks
+            .expect("we should only walk blocks with CollectionMode::UsedItems")
+            .contains(block)
+        {
+            self.super_basic_block_data(block, data)
+        }
+    }
+
     fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: Location) {
         debug!("visiting rvalue {:?}", *rvalue);
 
@@ -1411,6 +1423,11 @@ fn collect_items_of_instance<'tcx>(
         move_size_spans: vec![],
         visiting_call_terminator: false,
         skip_move_check_fns: None,
+        reachable_blocks: if mode == CollectionMode::UsedItems {
+            Some(tcx.reachable_blocks(instance))
+        } else {
+            None
+        },
     };
 
     if mode == CollectionMode::UsedItems {