move all mono-time checks into their own folder, and their own query

5 files changed, 124 insertions, 101 deletions

diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index 85151e5f093..429e31b2c88 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -205,13 +205,8 @@
 //! this is not implemented however: a mono item will be produced
 //! regardless of whether it is actually needed or not.
 
-mod abi_check;
-mod move_check;
-
 use std::path::PathBuf;
 
-use move_check::MoveCheckState;
-use rustc_abi::Size;
 use rustc_data_structures::sync::{LRef, MTLock, par_for_each_in};
 use rustc_data_structures::unord::{UnordMap, UnordSet};
 use rustc_hir as hir;
@@ -228,15 +223,15 @@ use rustc_middle::ty::adjustment::{CustomCoerceUnsized, PointerCoercion};
 use rustc_middle::ty::layout::ValidityRequirement;
 use rustc_middle::ty::print::{shrunk_instance_name, with_no_trimmed_paths};
 use rustc_middle::ty::{
-    self, AssocKind, GenericArgs, GenericParamDefKind, Instance, InstanceKind, Ty, TyCtxt,
-    TypeFoldable, TypeVisitableExt, VtblEntry,
+    self, GenericArgs, GenericParamDefKind, Instance, InstanceKind, Ty, TyCtxt, TypeFoldable,
+    TypeVisitableExt, VtblEntry,
 };
 use rustc_middle::util::Providers;
 use rustc_middle::{bug, span_bug};
 use rustc_session::Limit;
 use rustc_session::config::EntryFnType;
 use rustc_span::source_map::{Spanned, dummy_spanned, respan};
-use rustc_span::symbol::{Ident, sym};
+use rustc_span::symbol::sym;
 use rustc_span::{DUMMY_SP, Span};
 use tracing::{debug, instrument, trace};
 
@@ -612,8 +607,6 @@ struct MirUsedCollector<'a, 'tcx> {
     /// Note that this contains *not-monomorphized* items!
     used_mentioned_items: &'a mut UnordSet<MentionedItem<'tcx>>,
     instance: Instance<'tcx>,
-    visiting_call_terminator: bool,
-    move_check: move_check::MoveCheckState,
 }
 
 impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
@@ -760,13 +753,12 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
         };
 
         match terminator.kind {
-            mir::TerminatorKind::Call { ref func, ref args, ref fn_span, .. }
-            | mir::TerminatorKind::TailCall { ref func, ref args, ref fn_span } => {
+            mir::TerminatorKind::Call { ref func, .. }
+            | mir::TerminatorKind::TailCall { ref func, .. } => {
                 let callee_ty = func.ty(self.body, tcx);
                 // *Before* monomorphizing, record that we already handled this mention.
                 self.used_mentioned_items.insert(MentionedItem::Fn(callee_ty));
                 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)
             }
             mir::TerminatorKind::Drop { ref place, .. } => {
@@ -826,14 +818,7 @@ impl<'a, 'tcx> MirVisitor<'tcx> for MirUsedCollector<'a, 'tcx> {
             push_mono_lang_item(self, reason.lang_item());
         }
 
-        self.visiting_call_terminator = matches!(terminator.kind, mir::TerminatorKind::Call { .. });
         self.super_terminator(terminator, location);
-        self.visiting_call_terminator = false;
-    }
-
-    fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, location: Location) {
-        self.super_operand(operand, location);
-        self.check_operand_move_size(operand, location);
     }
 }
 
@@ -1183,20 +1168,6 @@ fn collect_alloc<'tcx>(tcx: TyCtxt<'tcx>, alloc_id: AllocId, output: &mut MonoIt
     }
 }
 
-fn assoc_fn_of_type<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, fn_ident: Ident) -> Option<DefId> {
-    for impl_def_id in tcx.inherent_impls(def_id) {
-        if let Some(new) = tcx.associated_items(impl_def_id).find_by_name_and_kind(
-            tcx,
-            fn_ident,
-            AssocKind::Fn,
-            def_id,
-        ) {
-            return Some(new.def_id);
-        }
-    }
-    None
-}
-
 /// Scans the MIR in order to find function calls, closures, and drop-glue.
 ///
 /// Anything that's found is added to `output`. Furthermore the "mentioned items" of the MIR are returned.
@@ -1208,7 +1179,8 @@ fn collect_items_of_instance<'tcx>(
     mentioned_items: &mut MonoItems<'tcx>,
     mode: CollectionMode,
 ) {
-    tcx.ensure().check_feature_dependent_abi(instance);
+    // This item is getting monomorphized, do mono-time checks.
+    tcx.ensure().check_mono_item(instance);
 
     let body = tcx.instance_mir(instance.def);
     // Naively, in "used" collection mode, all functions get added to *both* `used_items` and
@@ -1228,8 +1200,6 @@ fn collect_items_of_instance<'tcx>(
         used_items,
         used_mentioned_items: &mut used_mentioned_items,
         instance,
-        visiting_call_terminator: false,
-        move_check: MoveCheckState::new(),
     };
 
     if mode == CollectionMode::UsedItems {
@@ -1626,5 +1596,4 @@ pub(crate) fn collect_crate_mono_items<'tcx>(
 
 pub(crate) fn provide(providers: &mut Providers) {
     providers.hooks.should_codegen_locally = should_codegen_locally;
-    abi_check::provide(providers);
 }
diff --git a/compiler/rustc_monomorphize/src/lib.rs b/compiler/rustc_monomorphize/src/lib.rs
index e92e6978d0f..0cfc4371db5 100644
--- a/compiler/rustc_monomorphize/src/lib.rs
+++ b/compiler/rustc_monomorphize/src/lib.rs
@@ -16,6 +16,7 @@ use rustc_span::ErrorGuaranteed;
 
 mod collector;
 mod errors;
+mod mono_checks;
 mod partitioning;
 mod polymorphize;
 mod util;
@@ -47,4 +48,5 @@ fn custom_coerce_unsize_info<'tcx>(
 pub fn provide(providers: &mut Providers) {
     partitioning::provide(providers);
     polymorphize::provide(providers);
+    mono_checks::provide(providers);
 }
diff --git a/compiler/rustc_monomorphize/src/collector/abi_check.rs b/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs
index f1c57f1e997..221200e9497 100644
--- a/compiler/rustc_monomorphize/src/collector/abi_check.rs
+++ b/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs
@@ -1,9 +1,7 @@
 //! This module ensures that if a function's ABI requires a particular target feature,
 //! that target feature is enabled both on the callee and all callers.
 use rustc_hir::CRATE_HIR_ID;
-use rustc_middle::mir::visit::Visitor as MirVisitor;
-use rustc_middle::mir::{self, Location, traversal};
-use rustc_middle::query::Providers;
+use rustc_middle::mir::{self, traversal};
 use rustc_middle::ty::inherent::*;
 use rustc_middle::ty::{self, Instance, InstanceKind, ParamEnv, Ty, TyCtxt};
 use rustc_session::lint::builtin::ABI_UNSUPPORTED_VECTOR_TYPES;
@@ -120,43 +118,31 @@ fn check_call_site_abi<'tcx>(
     });
 }
 
-struct MirCallesAbiCheck<'a, 'tcx> {
-    tcx: TyCtxt<'tcx>,
-    body: &'a mir::Body<'tcx>,
-    instance: Instance<'tcx>,
-}
-
-impl<'a, 'tcx> MirVisitor<'tcx> for MirCallesAbiCheck<'a, 'tcx> {
-    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, _: Location) {
+fn check_callees_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>, body: &mir::Body<'tcx>) {
+    // Check all function call terminators.
+    for (bb, _data) in traversal::mono_reachable(body, tcx, instance) {
+        let terminator = body.basic_blocks[bb].terminator();
         match terminator.kind {
             mir::TerminatorKind::Call { ref func, ref fn_span, .. }
             | mir::TerminatorKind::TailCall { ref func, ref fn_span, .. } => {
-                let callee_ty = func.ty(self.body, self.tcx);
-                let callee_ty = self.instance.instantiate_mir_and_normalize_erasing_regions(
-                    self.tcx,
+                let callee_ty = func.ty(body, tcx);
+                let callee_ty = instance.instantiate_mir_and_normalize_erasing_regions(
+                    tcx,
                     ty::ParamEnv::reveal_all(),
                     ty::EarlyBinder::bind(callee_ty),
                 );
-                check_call_site_abi(self.tcx, callee_ty, *fn_span, self.body.source.instance);
+                check_call_site_abi(tcx, callee_ty, *fn_span, body.source.instance);
             }
             _ => {}
         }
     }
 }
 
-fn check_callees_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
-    let body = tcx.instance_mir(instance.def);
-    let mut visitor = MirCallesAbiCheck { tcx, body, instance };
-    for (bb, data) in traversal::mono_reachable(body, tcx, instance) {
-        visitor.visit_basic_block_data(bb, data)
-    }
-}
-
-fn check_feature_dependent_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
+pub(crate) fn check_feature_dependent_abi<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    body: &'tcx mir::Body<'tcx>,
+) {
     check_instance_abi(tcx, instance);
-    check_callees_abi(tcx, instance);
-}
-
-pub(super) fn provide(providers: &mut Providers) {
-    *providers = Providers { check_feature_dependent_abi, ..*providers }
+    check_callees_abi(tcx, instance, body);
 }
diff --git a/compiler/rustc_monomorphize/src/mono_checks/mod.rs b/compiler/rustc_monomorphize/src/mono_checks/mod.rs
new file mode 100644
index 00000000000..1ecda824fb8
--- /dev/null
+++ b/compiler/rustc_monomorphize/src/mono_checks/mod.rs
@@ -0,0 +1,23 @@
+//! This implements a single query, `check_mono_fn`, that gets fired for each
+//! monomorphization of all functions. This lets us implement monomorphization-time
+//! checks in a way that is friendly to incremental compilation.
+
+use rustc_middle::query::Providers;
+use rustc_middle::ty::{Instance, TyCtxt};
+
+mod abi_check;
+mod move_check;
+
+fn check_mono_item<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
+    let body = tcx.instance_mir(instance.def);
+    abi_check::check_feature_dependent_abi(tcx, instance, body);
+    move_check::check_moves(tcx, instance, body);
+}
+
+pub(super) fn provide(providers: &mut Providers) {
+    *providers = Providers {
+        check_mono_item,
+        skip_move_check_fns: move_check::skip_move_check_fns,
+        ..*providers
+    }
+}
diff --git a/compiler/rustc_monomorphize/src/collector/move_check.rs b/compiler/rustc_monomorphize/src/mono_checks/move_check.rs
index b86ef1e7373..7f04bdf46f1 100644
--- a/compiler/rustc_monomorphize/src/collector/move_check.rs
+++ b/compiler/rustc_monomorphize/src/mono_checks/move_check.rs
@@ -1,38 +1,74 @@
+use rustc_abi::Size;
+use rustc_data_structures::fx::FxIndexSet;
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir::visit::Visitor as MirVisitor;
+use rustc_middle::mir::{self, Location, traversal};
+use rustc_middle::ty::{self, AssocKind, Instance, Ty, TyCtxt, TypeFoldable};
+use rustc_session::Limit;
 use rustc_session::lint::builtin::LARGE_ASSIGNMENTS;
-use tracing::debug;
+use rustc_span::Span;
+use rustc_span::source_map::Spanned;
+use rustc_span::symbol::{Ident, sym};
+use tracing::{debug, trace};
 
-use super::*;
 use crate::errors::LargeAssignmentsLint;
 
-pub(super) struct MoveCheckState {
+struct MoveCheckVisitor<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    body: &'tcx mir::Body<'tcx>,
     /// Spans for move size lints already emitted. Helps avoid duplicate lints.
     move_size_spans: Vec<Span>,
-    /// Set of functions for which it is OK to move large data into.
-    skip_move_check_fns: Option<Vec<DefId>>,
 }
 
-impl MoveCheckState {
-    pub(super) fn new() -> Self {
-        MoveCheckState { move_size_spans: vec![], skip_move_check_fns: None }
+pub(crate) fn check_moves<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    body: &'tcx mir::Body<'tcx>,
+) {
+    let mut visitor = MoveCheckVisitor { tcx, instance, body, move_size_spans: vec![] };
+    for (bb, data) in traversal::mono_reachable(body, tcx, instance) {
+        visitor.visit_basic_block_data(bb, data)
     }
 }
 
-impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
-    pub(super) fn check_operand_move_size(
-        &mut self,
-        operand: &mir::Operand<'tcx>,
-        location: Location,
-    ) {
-        let limit = self.tcx.move_size_limit();
-        if limit.0 == 0 {
-            return;
+impl<'tcx> MirVisitor<'tcx> for MoveCheckVisitor<'tcx> {
+    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, location: Location) {
+        match terminator.kind {
+            mir::TerminatorKind::Call { ref func, ref args, ref fn_span, .. }
+            | mir::TerminatorKind::TailCall { ref func, ref args, ref fn_span } => {
+                let callee_ty = func.ty(self.body, self.tcx);
+                let callee_ty = self.monomorphize(callee_ty);
+                self.check_fn_args_move_size(callee_ty, args, *fn_span, location);
+            }
+            _ => {}
         }
 
-        // This function is called by visit_operand() which visits _all_
-        // operands, including TerminatorKind::Call operands. But if
-        // check_fn_args_move_size() has been called, the operands have already
-        // been visited. Do not visit them again.
-        if self.visiting_call_terminator {
+        // We deliberately do *not* visit the nested operands here, to avoid
+        // hitting `visit_operand` for function arguments.
+    }
+
+    fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, location: Location) {
+        self.check_operand_move_size(operand, location);
+    }
+}
+
+impl<'tcx> MoveCheckVisitor<'tcx> {
+    fn monomorphize<T>(&self, value: T) -> T
+    where
+        T: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        trace!("monomorphize: self.instance={:?}", self.instance);
+        self.instance.instantiate_mir_and_normalize_erasing_regions(
+            self.tcx,
+            ty::ParamEnv::reveal_all(),
+            ty::EarlyBinder::bind(value),
+        )
+    }
+
+    fn check_operand_move_size(&mut self, operand: &mir::Operand<'tcx>, location: Location) {
+        let limit = self.tcx.move_size_limit();
+        if limit.0 == 0 {
             return;
         }
 
@@ -64,12 +100,7 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
         let ty::FnDef(def_id, _) = *callee_ty.kind() else {
             return;
         };
-        if self
-            .move_check
-            .skip_move_check_fns
-            .get_or_insert_with(|| build_skip_move_check_fns(self.tcx))
-            .contains(&def_id)
-        {
+        if self.tcx.skip_move_check_fns(()).contains(&def_id) {
             return;
         }
 
@@ -116,14 +147,12 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
         span: Span,
     ) {
         let source_info = self.body.source_info(location);
-        debug!(?source_info);
-        for reported_span in &self.move_check.move_size_spans {
+        for reported_span in &self.move_size_spans {
             if reported_span.overlaps(span) {
                 return;
             }
         }
         let lint_root = source_info.scope.lint_root(&self.body.source_scopes);
-        debug!(?lint_root);
         let Some(lint_root) = lint_root else {
             // This happens when the issue is in a function from a foreign crate that
             // we monomorphized in the current crate. We can't get a `HirId` for things
@@ -137,11 +166,25 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
             size: too_large_size.bytes(),
             limit: limit as u64,
         });
-        self.move_check.move_size_spans.push(span);
+        self.move_size_spans.push(span);
+    }
+}
+
+fn assoc_fn_of_type<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, fn_ident: Ident) -> Option<DefId> {
+    for impl_def_id in tcx.inherent_impls(def_id) {
+        if let Some(new) = tcx.associated_items(impl_def_id).find_by_name_and_kind(
+            tcx,
+            fn_ident,
+            AssocKind::Fn,
+            def_id,
+        ) {
+            return Some(new.def_id);
+        }
     }
+    None
 }
 
-fn build_skip_move_check_fns(tcx: TyCtxt<'_>) -> Vec<DefId> {
+pub(crate) fn skip_move_check_fns(tcx: TyCtxt<'_>, _: ()) -> FxIndexSet<DefId> {
     let fns = [
         (tcx.lang_items().owned_box(), "new"),
         (tcx.get_diagnostic_item(sym::Rc), "new"),
@@ -151,5 +194,5 @@ fn build_skip_move_check_fns(tcx: TyCtxt<'_>) -> Vec<DefId> {
         .filter_map(|(def_id, fn_name)| {
             def_id.and_then(|def_id| assoc_fn_of_type(tcx, def_id, Ident::from_str(fn_name)))
         })
-        .collect::<Vec<_>>()
+        .collect()
 }