mono_checks: merge the two traversals into one

RalfJung · RalfJung · commit f25b73b67078 · 2024-11-10T16:17:00.000+01:00
diff --git a/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs b/compiler/rustc_monomorphize/src/mono_checks/abi_check.rs
@@ -1,7 +1,6 @@
 //! 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::{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;
@@ -56,7 +55,7 @@ fn do_check_abi<'tcx>(
 
 /// Checks that the ABI of a given instance of a function does not contain vector-passed arguments
 /// or return values for which the corresponding target feature is not enabled.
-fn check_instance_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
+pub(super) fn check_instance_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
     let param_env = ParamEnv::reveal_all();
     let Ok(abi) = tcx.fn_abi_of_instance(param_env.and((instance, ty::List::empty()))) else {
         // An error will be reported during codegen if we cannot determine the ABI of this
@@ -76,7 +75,7 @@ fn check_instance_abi<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
 
 /// Checks that a call expression does not try to pass a vector-passed argument which requires a
 /// target feature that the caller does not have, as doing so causes UB because of ABI mismatch.
-fn check_call_site_abi<'tcx>(
+pub(super) fn check_call_site_abi<'tcx>(
     tcx: TyCtxt<'tcx>,
     callee: Ty<'tcx>,
     span: Span,
@@ -117,32 +116,3 @@ fn check_call_site_abi<'tcx>(
         );
     });
 }
-
-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(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(tcx, callee_ty, *fn_span, body.source.instance);
-            }
-            _ => {}
-        }
-    }
-}
-
-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, body);
-}
diff --git a/compiler/rustc_monomorphize/src/mono_checks/mod.rs b/compiler/rustc_monomorphize/src/mono_checks/mod.rs
@@ -2,16 +2,81 @@
 //! monomorphization of all functions. This lets us implement monomorphization-time
 //! checks in a way that is friendly to incremental compilation.
 
+use rustc_middle::mir::visit::Visitor as MirVisitor;
+use rustc_middle::mir::{self, Location, traversal};
 use rustc_middle::query::Providers;
-use rustc_middle::ty::{Instance, TyCtxt};
+use rustc_middle::ty::{self, Instance, TyCtxt, TypeFoldable};
 
 mod abi_check;
 mod move_check;
 
+struct MonoCheckVisitor<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    instance: Instance<'tcx>,
+    body: &'tcx mir::Body<'tcx>,
+    move_check: move_check::State,
+    visiting_call_terminator: bool,
+}
+
+impl<'tcx> MirVisitor<'tcx> for MonoCheckVisitor<'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);
+
+                abi_check::check_call_site_abi(
+                    self.tcx,
+                    callee_ty,
+                    *fn_span,
+                    self.body.source.instance,
+                );
+                self.check_fn_args_move_size(callee_ty, args, *fn_span, location);
+
+                self.visiting_call_terminator = true;
+            }
+            _ => {
+                debug_assert!(!self.visiting_call_terminator);
+            }
+        }
+
+        self.super_terminator(terminator, location);
+        self.visiting_call_terminator = false;
+    }
+
+    fn visit_operand(&mut self, operand: &mir::Operand<'tcx>, location: Location) {
+        self.check_operand_move_size(operand, location);
+    }
+}
+
+impl<'tcx> MonoCheckVisitor<'tcx> {
+    fn monomorphize<T>(&self, value: T) -> T
+    where
+        T: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        self.instance.instantiate_mir_and_normalize_erasing_regions(
+            self.tcx,
+            ty::ParamEnv::reveal_all(),
+            ty::EarlyBinder::bind(value),
+        )
+    }
+}
+
 fn check_mono_item<'tcx>(tcx: TyCtxt<'tcx>, instance: Instance<'tcx>) {
+    abi_check::check_instance_abi(tcx, instance);
+
     let body = tcx.instance_mir(instance.def);
-    abi_check::check_feature_dependent_abi(tcx, instance, body);
-    move_check::check_moves(tcx, instance, body);
+    let mut visitor = MonoCheckVisitor {
+        tcx,
+        instance,
+        body,
+        move_check: Default::default(),
+        visiting_call_terminator: false,
+    };
+    for (bb, data) in traversal::mono_reachable(body, tcx, instance) {
+        visitor.visit_basic_block_data(bb, data)
+    }
 }
 
 pub(super) fn provide(providers: &mut Providers) {
diff --git a/compiler/rustc_monomorphize/src/mono_checks/move_check.rs b/compiler/rustc_monomorphize/src/mono_checks/move_check.rs
@@ -1,77 +1,43 @@
 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_middle::mir::{self, Location};
+use rustc_middle::ty::{self, AssocKind, Ty, TyCtxt};
 use rustc_session::Limit;
 use rustc_session::lint::builtin::LARGE_ASSIGNMENTS;
 use rustc_span::Span;
 use rustc_span::source_map::Spanned;
 use rustc_span::symbol::{Ident, sym};
-use tracing::{debug, trace};
+use tracing::debug;
 
+use super::MonoCheckVisitor;
 use crate::errors::LargeAssignmentsLint;
 
-struct MoveCheckVisitor<'tcx> {
-    tcx: TyCtxt<'tcx>,
-    instance: Instance<'tcx>,
-    body: &'tcx mir::Body<'tcx>,
+#[derive(Default)]
+pub(super) struct State {
     /// Spans for move size lints already emitted. Helps avoid duplicate lints.
     move_size_spans: Vec<Span>,
 }
 
-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<'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) {
+impl<'tcx> MonoCheckVisitor<'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;
         }
 
+        // 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);
 
@@ -147,7 +113,7 @@ impl<'tcx> MoveCheckVisitor<'tcx> {
         span: Span,
     ) {
         let source_info = self.body.source_info(location);
-        for reported_span in &self.move_size_spans {
+        for reported_span in &self.move_check.move_size_spans {
             if reported_span.overlaps(span) {
                 return;
             }
@@ -166,7 +132,7 @@ impl<'tcx> MoveCheckVisitor<'tcx> {
             size: too_large_size.bytes(),
             limit: limit as u64,
         });
-        self.move_size_spans.push(span);
+        self.move_check.move_size_spans.push(span);
     }
 }
 
