Auto merge of #132843 - RalfJung:mono-time-checks, r=lcnr

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

The mono item collector currently also drives two mono-time checks: the lint for "large moves", and the check whether function calls are done with all the required target features.

Instead of doing this "inside" the collector, this PR refactors things so that we have a new `rustc_monomorphize::mono_checks` module providing a per-instance query that does these checks. We already have a per-instance query for the ABI checks, so this should be "free" for incremental builds. Non-incremental builds might do a bit more work now since we now have two separate MIR visits (in the collector and the mono-time checks) -- but one of them is cached in case the MIR doesn't change, which is nice.

This slightly changes behavior of the large-move check since the "move_size_spans" deduplication logic now only works per-instance, not globally across the entire collector.

Cc `@saethlin` since you're also doing some work related to queries and caching and monomorphization, though I don't know if there's any interaction here.

compiler/rustc_middle/src/query/mod.rs

@@ -2326,13 +2326,19 @@ rustc_queries! {
         separate_provide_extern
     }
 
-    /// Check the signature of this function as well as all the call expressions inside of it
+    /// Perform monomorphization-time checking on this item.
-    /// to ensure that any target features required by the ABI are enabled.
+    /// This is used for lints/errors that can only be checked once the instance is fully
-    /// Should be called on a fully monomorphized instance.
+    /// monomorphized.
-    query check_feature_dependent_abi(key: ty::Instance<'tcx>) {
+    query check_mono_item(key: ty::Instance<'tcx>) {
-        desc { "check for feature-dependent ABI" }
+        desc { "monomorphization-time checking" }
         cache_on_disk_if { true }
     }
+
+    /// Builds the set of functions that should be skipped for the move-size check.
+    query skip_move_check_fns(_: ()) -> &'tcx FxIndexSet<DefId> {
+        arena_cache
+        desc { "functions to skip for move-size check" }
+    }
 }
 
 rustc_query_append! { define_callbacks! }

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,
+    self, GenericArgs, GenericParamDefKind, Instance, InstanceKind, Ty, TyCtxt, TypeFoldable,
-    TypeFoldable, TypeVisitableExt, VtblEntry,
+    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::Call { ref func, .. }
-            | mir::TerminatorKind::TailCall { ref func, ref args, ref fn_span } => {
+            | 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);
 }

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);
 }

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, traversal};
-use rustc_middle::mir::{self, Location, traversal};
-use rustc_middle::query::Providers;
 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> {
+fn check_callees_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>, body: &mir::Body<'tcx>) {
-    tcx: TyCtxt<'tcx>,
+    // Check all function call terminators.
-    body: &'a mir::Body<'tcx>,
+    for (bb, _data) in traversal::mono_reachable(body, tcx, instance) {
-    instance: Instance<'tcx>,
+        let terminator = body.basic_blocks[bb].terminator();
-}
-
-impl<'a, 'tcx> MirVisitor<'tcx> for MirCallesAbiCheck<'a, 'tcx> {
-    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, _: Location) {
         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 = func.ty(body, tcx);
-                let callee_ty = self.instance.instantiate_mir_and_normalize_erasing_regions(
+                let callee_ty = instance.instantiate_mir_and_normalize_erasing_regions(
-                    self.tcx,
+                    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>) {
+pub(crate) fn check_feature_dependent_abi<'tcx>(
-    let body = tcx.instance_mir(instance.def);
+    tcx: TyCtxt<'tcx>,
-    let mut visitor = MirCallesAbiCheck { tcx, body, instance };
+    instance: Instance<'tcx>,
-    for (bb, data) in traversal::mono_reachable(body, tcx, instance) {
+    body: &'tcx mir::Body<'tcx>,
-        visitor.visit_basic_block_data(bb, data)
+) {
-    }
-}
-
-fn check_feature_dependent_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
     check_instance_abi(tcx, instance);
-    check_callees_abi(tcx, instance);
+    check_callees_abi(tcx, instance, body);
-}
-
-pub(super) fn provide(providers: &mut Providers) {
-    *providers = Providers { check_feature_dependent_abi, ..*providers }
 }

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
+    }
+}

compiler/rustc_monomorphize/src/mono_checks/move_check.rs

@@ -1,41 +1,77 @@
+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(crate) fn check_moves<'tcx>(
-    pub(super) fn new() -> Self {
+    tcx: TyCtxt<'tcx>,
-        MoveCheckState { move_size_spans: vec![], skip_move_check_fns: None }
+    instance: Instance<'tcx>,
-    }
+    body: &'tcx mir::Body<'tcx>,
-}
-
-impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
-    pub(super) fn check_operand_move_size(
-        &mut self,
-        operand: &mir::Operand<'tcx>,
-        location: Location,
 ) {
+    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<'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);
+            }
+            _ => {}
+        }
+
+        // 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;
         }
 
-        // 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 {
-            return;
-        }
-
         let source_info = self.body.source_info(location);
         debug!(?source_info);
 
@@ -64,12 +100,7 @@ impl<'a, 'tcx> MirUsedCollector<'a, 'tcx> {
         let ty::FnDef(def_id, _) = *callee_ty.kind() else {
             return;
         };
-        if self
+        if self.tcx.skip_move_check_fns(()).contains(&def_id) {
-            .move_check
-            .skip_move_check_fns
-            .get_or_insert_with(|| build_skip_move_check_fns(self.tcx))
-            .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_size_spans {
-        for reported_span in &self.move_check.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 build_skip_move_check_fns(tcx: TyCtxt<'_>) -> Vec<DefId> {
+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
+}
+
+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()
 }