Introduce and provide a hook for `should_codegen_locally`

4 files changed, 22 insertions, 4 deletions

diff --git a/compiler/rustc_middle/src/hooks/mod.rs b/compiler/rustc_middle/src/hooks/mod.rs
index bf10a71dbae..75dc685a16a 100644
--- a/compiler/rustc_middle/src/hooks/mod.rs
+++ b/compiler/rustc_middle/src/hooks/mod.rs
@@ -103,6 +103,10 @@ declare_hooks! {
 
     /// Create a list-like THIR representation for debugging.
     hook thir_flat(key: LocalDefId) -> String;
+
+    /// Returns `true` if we should codegen an instance in the local crate, or returns `false` if we
+    /// can just link to the upstream crate and therefore don't need a mono item.
+    hook should_codegen_locally(instance: crate::ty::Instance<'tcx>) -> bool;
 }
 
 #[cold]
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index bfd505c0672..27d27885ea8 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -228,6 +228,7 @@ use rustc_middle::ty::{
     self, AssocKind, GenericParamDefKind, Instance, InstanceKind, Ty, TyCtxt, TypeFoldable,
     TypeVisitableExt, VtblEntry,
 };
+use rustc_middle::util::Providers;
 use rustc_middle::{bug, span_bug};
 use rustc_session::config::EntryFnType;
 use rustc_session::Limit;
@@ -930,7 +931,7 @@ fn visit_instance_use<'tcx>(
 
 /// Returns `true` if we should codegen an instance in the local crate, or returns `false` if we
 /// can just link to the upstream crate and therefore don't need a mono item.
-pub(crate) fn should_codegen_locally<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> bool {
+pub(crate) fn should_codegen_locally_hook<'tcx>(tcx: TyCtxtAt<'tcx>, instance: Instance<'tcx>) -> bool {
     let Some(def_id) = instance.def.def_id_if_not_guaranteed_local_codegen() else {
         return true;
     };
@@ -946,7 +947,7 @@ pub(crate) fn should_codegen_locally<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance
     }
 
     if tcx.is_reachable_non_generic(def_id)
-        || instance.polymorphize(tcx).upstream_monomorphization(tcx).is_some()
+        || instance.polymorphize(*tcx).upstream_monomorphization(*tcx).is_some()
     {
         // We can link to the item in question, no instance needed in this crate.
         return false;
@@ -967,6 +968,12 @@ pub(crate) fn should_codegen_locally<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance
     true
 }
 
+/// Returns `true` if we should codegen an instance in the local crate, or returns `false` if we
+/// can just link to the upstream crate and therefore don't need a mono item.
+pub(crate) fn should_codegen_locally<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) -> bool {
+    tcx.should_codegen_locally(instance)
+}
+
 /// For a given pair of source and target type that occur in an unsizing coercion,
 /// this function finds the pair of types that determines the vtable linking
 /// them.
@@ -1613,3 +1620,7 @@ pub(crate) fn collect_crate_mono_items<'tcx>(
 
     (mono_items, state.usage_map.into_inner())
 }
+
+pub fn provide(providers: &mut Providers) {
+    providers.hooks.should_codegen_locally = should_codegen_locally_hook;
+}
diff --git a/compiler/rustc_monomorphize/src/lib.rs b/compiler/rustc_monomorphize/src/lib.rs
index aa3b4cd5b67..d169f086066 100644
--- a/compiler/rustc_monomorphize/src/lib.rs
+++ b/compiler/rustc_monomorphize/src/lib.rs
@@ -5,12 +5,13 @@
 
 use rustc_hir::lang_items::LangItem;
 use rustc_middle::bug;
-use rustc_middle::query::{Providers, TyCtxtAt};
+use rustc_middle::query::TyCtxtAt;
 use rustc_middle::traits;
 use rustc_middle::ty::adjustment::CustomCoerceUnsized;
 use rustc_middle::ty::Instance;
 use rustc_middle::ty::TyCtxt;
 use rustc_middle::ty::{self, Ty};
+use rustc_middle::util::Providers;
 use rustc_span::def_id::DefId;
 use rustc_span::def_id::LOCAL_CRATE;
 use rustc_span::ErrorGuaranteed;
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
index 9a7c488833a..8c7c5e0074a 100644
--- a/compiler/rustc_monomorphize/src/partitioning.rs
+++ b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -112,9 +112,9 @@ use rustc_middle::mir::mono::{
     CodegenUnit, CodegenUnitNameBuilder, InstantiationMode, Linkage, MonoItem, MonoItemData,
     Visibility,
 };
-use rustc_middle::query::Providers;
 use rustc_middle::ty::print::{characteristic_def_id_of_type, with_no_trimmed_paths};
 use rustc_middle::ty::{self, visit::TypeVisitableExt, InstanceKind, TyCtxt};
+use rustc_middle::util::Providers;
 use rustc_session::config::{DumpMonoStatsFormat, SwitchWithOptPath};
 use rustc_session::CodegenUnits;
 use rustc_span::symbol::Symbol;
@@ -1314,4 +1314,6 @@ pub fn provide(providers: &mut Providers) {
             .find(|cgu| cgu.name() == name)
             .unwrap_or_else(|| panic!("failed to find cgu with name {name:?}"))
     };
+
+    collector::provide(providers);
 }