Auto merge of rust-lang#122568 - RalfJung:mentioned-items, r=<try>

recursively evaluate the constants in everything that is 'mentioned'

This is another attempt at fixing rust-lang#107503. The previous attempt at rust-lang#112879 seems stuck in figuring out where the [perf regression](https://perf.rust-lang.org/compare.html?start=c55d1ee8d4e3162187214692229a63c2cc5e0f31&end=ec8de1ebe0d698b109beeaaac83e60f4ef8bb7d1&stat=instructions:u) comes from. In  rust-lang#122258 I learned some things, which informed the approach this PR is taking.

Quoting from the new collector docs, which explain the high-level idea:
```rust
//! One important role of collection is to evaluate all constants that are used by all the items
//! which are being collected. Codegen can then rely on only encountering constants that evaluate
//! successfully, and if a constant fails to evaluate, the collector has much better context to be
//! able to show where this constant comes up.
//!
//! However, the exact set of "used" items (collected as described above), and therefore the exact
//! set of used constants, can depend on optimizations. Optimizing away dead code may optimize away
//! a function call that uses a failing constant, so an unoptimized build may fail where an
//! optimized build succeeds. This is undesirable.
//!
//! To fix this, the collector has the concept of "mentioned" items. Some time during the MIR
//! pipeline, before any optimization-level-dependent optimizations, we compute a list of all items
//! that syntactically appear in the code. These are considered "mentioned", and even if they are in
//! dead code and get optimized away (which makes them no longer "used"), they are still
//! "mentioned". For every used item, the collector ensures that all mentioned items, recursively,
//! do not use a failing constant. This is reflected via the [`CollectionMode`], which determines
//! whether we are visiting a used item or merely a mentioned item.

enum CollectionMode {
    /// Collect items that are used, i.e., actually needed for codegen.
    ///
    /// Which items are used can depend on optimization levels, as MIR optimizations can remove
    /// uses.
    UsedItems,
    /// Collect items that are mentioned. The goal of this mode is that it is independent of
    /// optimizations: the set of "mentioned" items is computed before optimizations are run.
    ///
    /// The exact contents of this set are *not* a stable guarantee. (For instance, it is currently
    /// computed after drop-elaboration. If we ever do some optimizations even in debug builds, we
    /// might decide to run them before computing mentioned items.) The key property of this set is
    /// that it is optimization-independent.
    MentionedItems,
}
```

Fixes rust-lang#107503

Author: bors

compiler/rustc_data_structures/src/sync.rs

@@ -20,6 +20,7 @@
 //! | ----------------------- | ------------------- | ------------------------------- |
 //! | `Lrc<T>`                | `rc::Rc<T>`         | `sync::Arc<T>`                  |
 //! |` Weak<T>`               | `rc::Weak<T>`       | `sync::Weak<T>`                 |
+//! | `LRef<'a, T>` [^2]      | `&'a mut T`         | `&'a T`                         |
 //! |                         |                     |                                 |
 //! | `AtomicBool`            | `Cell<bool>`        | `atomic::AtomicBool`            |
 //! | `AtomicU32`             | `Cell<u32>`         | `atomic::AtomicU32`             |
@@ -38,7 +39,7 @@
 //! of a `RefCell`. This is appropriate when interior mutability is not
 //! required.
 //!
-//! [^2] `MTLockRef` is a typedef.
+//! [^2] `MTRef`, `MTLockRef` are type aliases.
 
 pub use crate::marker::*;
 use std::collections::HashMap;
@@ -208,7 +209,7 @@ cfg_match! {
 
         use std::cell::RefCell as InnerRwLock;
 
-        pub type MTLockRef<'a, T> = &'a mut MTLock<T>;
+        pub type LRef<'a, T> = &'a mut T;
 
         #[derive(Debug, Default)]
         pub struct MTLock<T>(T);
@@ -274,7 +275,7 @@ cfg_match! {
         pub use std::sync::Arc as Lrc;
         pub use std::sync::Weak as Weak;
 
-        pub type MTLockRef<'a, T> = &'a MTLock<T>;
+        pub type LRef<'a, T> = &'a T;
 
         #[derive(Debug, Default)]
         pub struct MTLock<T>(Lock<T>);
@@ -314,6 +315,8 @@ cfg_match! {
     }
 }
 
+pub type MTLockRef<'a, T> = LRef<'a, MTLock<T>>;
+
 #[derive(Default)]
 #[cfg_attr(parallel_compiler, repr(align(64)))]
 pub struct CacheAligned<T>(pub T);

compiler/rustc_middle/src/mir/mod.rs

@@ -20,6 +20,7 @@ use rustc_hir::def_id::{DefId, CRATE_DEF_ID};
 use rustc_hir::{self, CoroutineDesugaring, CoroutineKind, ImplicitSelfKind};
 use rustc_hir::{self as hir, HirId};
 use rustc_session::Session;
+use rustc_span::source_map::Spanned;
 use rustc_target::abi::{FieldIdx, VariantIdx};
 
 use polonius_engine::Atom;
@@ -312,6 +313,21 @@ impl<'tcx> CoroutineInfo<'tcx> {
     }
 }
 
+/// Some item that needs to monomorphize successfully for a MIR body to be considered well-formed.
+#[derive(Copy, Clone, PartialEq, Eq, Debug, Hash, HashStable, TyEncodable, TyDecodable)]
+#[derive(TypeFoldable, TypeVisitable)]
+pub enum MentionedItem<'tcx> {
+    Fn(DefId, GenericArgsRef<'tcx>),
+    Drop(Ty<'tcx>),
+    /// Unsizing casts might require vtables, so we have to record them.
+    UnsizeCast {
+        source_ty: Ty<'tcx>,
+        target_ty: Ty<'tcx>,
+    },
+    /// A closure that is coerced to a function pointer.
+    Closure(DefId, GenericArgsRef<'tcx>),
+}
+
 /// The lowered representation of a single function.
 #[derive(Clone, TyEncodable, TyDecodable, Debug, HashStable, TypeFoldable, TypeVisitable)]
 pub struct Body<'tcx> {
@@ -375,8 +391,21 @@ pub struct Body<'tcx> {
 
     /// Constants that are required to evaluate successfully for this MIR to be well-formed.
     /// We hold in this field all the constants we are not able to evaluate yet.
+    ///
+    /// This is soundness-critical, we make a guarantee that all consts syntactically mentioned in a
+    /// function have successfully evaluated if the function ever gets executed at runtime.
     pub required_consts: Vec<ConstOperand<'tcx>>,
 
+    /// Further items that were mentioned in this function and hence *may* become monomorphized,
+    /// depending on optimizations. We use this to avoid optimization-dependent compile errors: the
+    /// collector recursively traverses all "mentioned" items and evaluates all their
+    /// `requiered_consts`.
+    ///
+    /// This is *not* soundness-critical and the contents of this list are *not* a stable guarantee.
+    /// See the documentation of `CollectionMode` in `compiler/rustc_monomorphize/src/collector.rs`
+    /// for more context.
+    pub mentioned_items: Vec<Spanned<MentionedItem<'tcx>>>,
+
     /// Does this body use generic parameters. This is used for the `ConstEvaluatable` check.
     ///
     /// Note that this does not actually mean that this body is not computable right now.
@@ -453,6 +482,7 @@ impl<'tcx> Body<'tcx> {
             var_debug_info,
             span,
             required_consts: Vec::new(),
+            mentioned_items: Vec::new(),
             is_polymorphic: false,
             injection_phase: None,
             tainted_by_errors,
@@ -482,6 +512,7 @@ impl<'tcx> Body<'tcx> {
             spread_arg: None,
             span: DUMMY_SP,
             required_consts: Vec::new(),
+            mentioned_items: Vec::new(),
             var_debug_info: Vec::new(),
             is_polymorphic: false,
             injection_phase: None,

compiler/rustc_mir_build/src/build/custom/mod.rs

@@ -56,6 +56,7 @@ pub(super) fn build_custom_mir<'tcx>(
         var_debug_info: Vec::new(),
         span,
         required_consts: Vec::new(),
+        mentioned_items: Vec::new(),
         is_polymorphic: false,
         tainted_by_errors: None,
         injection_phase: None,

compiler/rustc_mir_transform/src/inline.rs

@@ -565,7 +565,8 @@ impl<'tcx> Inliner<'tcx> {
         mut callee_body: Body<'tcx>,
     ) {
         let terminator = caller_body[callsite.block].terminator.take().unwrap();
-        let TerminatorKind::Call { args, destination, unwind, target, .. } = terminator.kind else {
+        let TerminatorKind::Call { func, args, destination, unwind, target, .. } = terminator.kind
+        else {
             bug!("unexpected terminator kind {:?}", terminator.kind);
         };
 
@@ -717,6 +718,29 @@ impl<'tcx> Inliner<'tcx> {
                 Const::Val(..) | Const::Unevaluated(..) => true,
             },
         ));
+        // Now that we incorporated the callee's `required_consts`, we can remove the callee from
+        // `mentioned_items` -- but we have to take their `mentioned_items` in return.
+        // This does some extra work here to save the monomorphization collector work later.
+        // FIXME: benchmark which option is better.
+        let callee_item = {
+            // We need to reconstruct the `required_item` for the callee so that we can find and
+            // remove it.
+            let func_ty = func.ty(caller_body, self.tcx);
+            match func_ty.kind() {
+                ty::FnDef(def_id, args) => MentionedItem::Fn(*def_id, args),
+                _ => bug!(),
+            }
+        };
+        if let Some(idx) =
+            caller_body.mentioned_items.iter().position(|item| item.node == callee_item)
+        {
+            // We found the callee, so remove it and add its items instead.
+            caller_body.mentioned_items.remove(idx);
+            caller_body.mentioned_items.extend(callee_body.mentioned_items);
+        } else {
+            // If we can't find the callee, there's no point in adding its items.
+            // Probably it already got removed by being inlined elsewhere in the same function.
+        }
     }
 
     fn make_call_args(

compiler/rustc_mir_transform/src/lib.rs

@@ -89,6 +89,7 @@ mod lint;
 mod lower_intrinsics;
 mod lower_slice_len;
 mod match_branches;
+mod mentioned_items;
 mod multiple_return_terminators;
 mod normalize_array_len;
 mod nrvo;
@@ -566,6 +567,10 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         tcx,
         body,
         &[
+            // Before doing anything, remember which items are being mentioned so that the set of items
+            // visited does not depend on the optimization level.
+            &mentioned_items::MentionedItems,
+            // Add some UB checks before any UB gets optimized away.
             &check_alignment::CheckAlignment,
             // Before inlining: trim down MIR with passes to reduce inlining work.
 

compiler/rustc_mir_transform/src/mentioned_items.rs

@@ -0,0 +1,126 @@
+use rustc_hir::def_id::DefId;
+use rustc_middle::mir::visit::Visitor;
+use rustc_middle::mir::{self, ConstOperand, Location, MentionedItem, MirPass};
+use rustc_middle::ty::{self, adjustment::PointerCoercion, TyCtxt};
+use rustc_session::Session;
+use rustc_span::source_map::Spanned;
+
+pub struct MentionedItems;
+
+struct MentionedItemsVisitor<'a, 'tcx> {
+    tcx: TyCtxt<'tcx>,
+    body: &'a mir::Body<'tcx>,
+    mentioned_items: &'a mut Vec<Spanned<MentionedItem<'tcx>>>,
+}
+
+impl<'tcx> MirPass<'tcx> for MentionedItems {
+    fn is_enabled(&self, _sess: &Session) -> bool {
+        // If this pass is skipped the collector assume that nothing got mentioned! We could
+        // potentially skip it in opt-level 0 if we are sure that opt-level will never *remove* uses
+        // of anything, but that still seems fragile. Furthermore, even debug builds use level 1, so
+        // special-casing level 0 is just not worth it.
+        true
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
+        debug_assert!(body.mentioned_items.is_empty());
+        let mut mentioned_items = Vec::new();
+        MentionedItemsVisitor { tcx, body, mentioned_items: &mut mentioned_items }.visit_body(body);
+        body.mentioned_items = mentioned_items;
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for MentionedItemsVisitor<'_, 'tcx> {
+    fn visit_constant(&mut self, constant: &ConstOperand<'tcx>, _: Location) {
+        let const_ = constant.const_;
+        // This is how function items get referenced: via constants of `FnDef` type. This handles
+        // both functions that are called and those that are just turned to function pointers.
+        if let ty::FnDef(def_id, args) = *const_.ty().kind()
+            && may_codegen_locally(self.tcx, def_id)
+        {
+            debug!("adding to required_items: {def_id:?}");
+            self.mentioned_items
+                .push(Spanned { node: MentionedItem::Fn(def_id, args), span: constant.span });
+        }
+    }
+
+    fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, location: Location) {
+        self.super_terminator(terminator, location);
+        match terminator.kind {
+            // We don't need to handle `Call` as we already handled all function type operands in
+            // `visit_constant`. But we do need to handle `Drop`.
+            mir::TerminatorKind::Drop { place, .. } => {
+                let ty = place.ty(self.body, self.tcx).ty;
+                let span = self.body.source_info(location).span;
+                self.mentioned_items.push(Spanned { node: MentionedItem::Drop(ty), span });
+            }
+            _ => {}
+        }
+    }
+
+    fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: Location) {
+        self.super_rvalue(rvalue, location);
+        match *rvalue {
+            // We need to detect unsizing casts that required vtables.
+            mir::Rvalue::Cast(
+                mir::CastKind::PointerCoercion(PointerCoercion::Unsize),
+                ref operand,
+                target_ty,
+            )
+            | mir::Rvalue::Cast(mir::CastKind::DynStar, ref operand, target_ty) => {
+                let span = self.body.source_info(location).span;
+                self.mentioned_items.push(Spanned {
+                    node: MentionedItem::UnsizeCast {
+                        source_ty: operand.ty(self.body, self.tcx),
+                        target_ty,
+                    },
+                    span,
+                });
+            }
+            // Similarly, record closures that are turned into function pointers.
+            mir::Rvalue::Cast(
+                mir::CastKind::PointerCoercion(PointerCoercion::ClosureFnPointer(_)),
+                ref operand,
+                _,
+            ) => {
+                let span = self.body.source_info(location).span;
+                let source_ty = operand.ty(self.body, self.tcx);
+                match *source_ty.kind() {
+                    ty::Closure(def_id, args) => {
+                        self.mentioned_items
+                            .push(Spanned { node: MentionedItem::Closure(def_id, args), span });
+                    }
+                    _ => bug!(),
+                }
+            }
+            // Function pointer casts are already handled by `visit_constant` above.
+            _ => {}
+        }
+    }
+}
+
+/// Returns `true` if we should codegen an item in the local crate, or returns `false` if we
+/// can just link to the upstream crate and therefore don't need a mono item.
+///
+/// This is an approximation of the collector's `should_codegen_locally`, in the sense that this
+/// here may return `true` even if `should_codegen_locally` says `false`. The point is to let us
+/// filter out items that definitely will not be considered by the collector.
+fn may_codegen_locally<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> bool {
+    if tcx.is_foreign_item(def_id) {
+        // Foreign items are always linked against, there's no way of instantiating them.
+        return false;
+    }
+
+    if def_id.is_local() {
+        // Local items cannot be referred to locally without monomorphizing them locally.
+        return true;
+    }
+
+    if tcx.is_reachable_non_generic(def_id) {
+        // We can link to the item in question, no instance needed in this crate.
+        return false;
+    }
+
+    // Conservative fall-back.
+    true
+}

compiler/rustc_mir_transform/src/shim.rs

@@ -17,7 +17,7 @@ use std::iter;
 
 use crate::{
     abort_unwinding_calls, add_call_guards, add_moves_for_packed_drops, deref_separator,
-    pass_manager as pm, remove_noop_landing_pads, simplify,
+    mentioned_items, pass_manager as pm, remove_noop_landing_pads, simplify,
 };
 use rustc_middle::mir::patch::MirPatch;
 use rustc_mir_dataflow::elaborate_drops::{self, DropElaborator, DropFlagMode, DropStyle};
@@ -147,6 +147,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
                     tcx,
                     &mut body,
                     &[
+                        &mentioned_items::MentionedItems,
                         &abort_unwinding_calls::AbortUnwindingCalls,
                         &add_call_guards::CriticalCallEdges,
                     ],
@@ -178,6 +179,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceDef<'tcx>) -> Body<'
         tcx,
         &mut result,
         &[
+            &mentioned_items::MentionedItems,
             &add_moves_for_packed_drops::AddMovesForPackedDrops,
             &deref_separator::Derefer,
             &remove_noop_landing_pads::RemoveNoopLandingPads,