Auto merge of #123937 - RalfJung:miri-link-section, r=oli-obk

Miri on Windows: run .CRT$XLB linker section on thread-end

Hopefully fixes #123583

First commit is originally by `@bjorn3`

r? `@oli-obk`
Cc `@ChrisDenton`

library/std/src/sys/pal/windows/thread_local_key.rs

@@ -276,6 +276,7 @@ unsafe fn register_dtor(key: &'static StaticKey) {
 // the address of the symbol to ensure it sticks around.
 
 #[link_section = ".CRT$XLB"]
+#[cfg_attr(miri, used)] // Miri only considers explicitly `#[used]` statics for `lookup_link_section`
 pub static p_thread_callback: unsafe extern "system" fn(c::LPVOID, c::DWORD, c::LPVOID) =
     on_tls_callback;
 

src/tools/miri/src/bin/miri.rs

@@ -32,8 +32,9 @@ use rustc_driver::Compilation;
 use rustc_hir::{self as hir, Node};
 use rustc_interface::interface::Config;
 use rustc_middle::{
-    middle::exported_symbols::{
-        ExportedSymbol, SymbolExportInfo, SymbolExportKind, SymbolExportLevel,
+    middle::{
+        codegen_fn_attrs::CodegenFnAttrFlags,
+        exported_symbols::{ExportedSymbol, SymbolExportInfo, SymbolExportKind, SymbolExportLevel},
     },
     query::LocalCrate,
     ty::TyCtxt,
@@ -136,6 +137,7 @@ impl rustc_driver::Callbacks for MiriBeRustCompilerCalls {
             config.override_queries = Some(|_, local_providers| {
                 // `exported_symbols` and `reachable_non_generics` provided by rustc always returns
                 // an empty result if `tcx.sess.opts.output_types.should_codegen()` is false.
+                // In addition we need to add #[used] symbols to exported_symbols for `lookup_link_section`.
                 local_providers.exported_symbols = |tcx, LocalCrate| {
                     let reachable_set = tcx.with_stable_hashing_context(|hcx| {
                         tcx.reachable_set(()).to_sorted(&hcx, true)
@@ -160,19 +162,28 @@ impl rustc_driver::Callbacks for MiriBeRustCompilerCalls {
                                 })
                                 if !tcx.generics_of(local_def_id).requires_monomorphization(tcx)
                             );
-                            (is_reachable_non_generic
-                                && tcx.codegen_fn_attrs(local_def_id).contains_extern_indicator())
-                            .then_some((
-                                ExportedSymbol::NonGeneric(local_def_id.to_def_id()),
-                                // Some dummy `SymbolExportInfo` here. We only use
-                                // `exported_symbols` in shims/foreign_items.rs and the export info
-                                // is ignored.
-                                SymbolExportInfo {
-                                    level: SymbolExportLevel::C,
-                                    kind: SymbolExportKind::Text,
-                                    used: false,
-                                },
-                            ))
+                            if !is_reachable_non_generic {
+                                return None;
+                            }
+                            let codegen_fn_attrs = tcx.codegen_fn_attrs(local_def_id);
+                            if codegen_fn_attrs.contains_extern_indicator()
+                                || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::USED)
+                                || codegen_fn_attrs.flags.contains(CodegenFnAttrFlags::USED_LINKER)
+                            {
+                                Some((
+                                    ExportedSymbol::NonGeneric(local_def_id.to_def_id()),
+                                    // Some dummy `SymbolExportInfo` here. We only use
+                                    // `exported_symbols` in shims/foreign_items.rs and the export info
+                                    // is ignored.
+                                    SymbolExportInfo {
+                                        level: SymbolExportLevel::C,
+                                        kind: SymbolExportKind::Text,
+                                        used: false,
+                                    },
+                                ))
+                            } else {
+                                None
+                            }
                         }),
                     )
                 }

src/tools/miri/src/concurrency/thread.rs

@@ -158,7 +158,7 @@ pub struct Thread<'mir, 'tcx> {
 }
 
 pub type StackEmptyCallback<'mir, 'tcx> =
-    Box<dyn FnMut(&mut MiriInterpCx<'mir, 'tcx>) -> InterpResult<'tcx, Poll<()>>>;
+    Box<dyn FnMut(&mut MiriInterpCx<'mir, 'tcx>) -> InterpResult<'tcx, Poll<()>> + 'tcx>;
 
 impl<'mir, 'tcx> Thread<'mir, 'tcx> {
     /// Get the name of the current thread if it was set.

src/tools/miri/src/eval.rs

@@ -192,18 +192,18 @@ impl Default for MiriConfig {
 
 /// The state of the main thread. Implementation detail of `on_main_stack_empty`.
 #[derive(Default, Debug)]
-enum MainThreadState {
+enum MainThreadState<'tcx> {
     #[default]
     Running,
-    TlsDtors(tls::TlsDtorsState),
+    TlsDtors(tls::TlsDtorsState<'tcx>),
     Yield {
         remaining: u32,
     },
     Done,
 }
 
-impl MainThreadState {
-    fn on_main_stack_empty<'tcx>(
+impl<'tcx> MainThreadState<'tcx> {
+    fn on_main_stack_empty(
         &mut self,
         this: &mut MiriInterpCx<'_, 'tcx>,
     ) -> InterpResult<'tcx, Poll<()>> {

src/tools/miri/src/helpers.rs

@@ -9,16 +9,21 @@ use rand::RngCore;
 
 use rustc_apfloat::ieee::{Double, Single};
 use rustc_apfloat::Float;
-use rustc_hir::def::{DefKind, Namespace};
-use rustc_hir::def_id::{DefId, CRATE_DEF_INDEX};
+use rustc_hir::{
+    def::{DefKind, Namespace},
+    def_id::{CrateNum, DefId, CRATE_DEF_INDEX, LOCAL_CRATE},
+};
 use rustc_index::IndexVec;
+use rustc_middle::middle::dependency_format::Linkage;
+use rustc_middle::middle::exported_symbols::ExportedSymbol;
 use rustc_middle::mir;
 use rustc_middle::ty::{
     self,
     layout::{LayoutOf, TyAndLayout},
     FloatTy, IntTy, Ty, TyCtxt, UintTy,
 };
-use rustc_span::{def_id::CrateNum, sym, Span, Symbol};
+use rustc_session::config::CrateType;
+use rustc_span::{sym, Span, Symbol};
 use rustc_target::abi::{Align, FieldIdx, FieldsShape, Size, Variants};
 use rustc_target::spec::abi::Abi;
 
@@ -142,6 +147,38 @@ fn try_resolve_did(tcx: TyCtxt<'_>, path: &[&str], namespace: Option<Namespace>)
     None
 }
 
+/// Call `f` for each exported symbol.
+pub fn iter_exported_symbols<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    mut f: impl FnMut(CrateNum, DefId) -> InterpResult<'tcx>,
+) -> InterpResult<'tcx> {
+    // `dependency_formats` includes all the transitive informations needed to link a crate,
+    // which is what we need here since we need to dig out `exported_symbols` from all transitive
+    // dependencies.
+    let dependency_formats = tcx.dependency_formats(());
+    let dependency_format = dependency_formats
+        .iter()
+        .find(|(crate_type, _)| *crate_type == CrateType::Executable)
+        .expect("interpreting a non-executable crate");
+    for cnum in iter::once(LOCAL_CRATE).chain(dependency_format.1.iter().enumerate().filter_map(
+        |(num, &linkage)| {
+            // We add 1 to the number because that's what rustc also does everywhere it
+            // calls `CrateNum::new`...
+            #[allow(clippy::arithmetic_side_effects)]
+            (linkage != Linkage::NotLinked).then_some(CrateNum::new(num + 1))
+        },
+    )) {
+        // We can ignore `_export_info` here: we are a Rust crate, and everything is exported
+        // from a Rust crate.
+        for &(symbol, _export_info) in tcx.exported_symbols(cnum) {
+            if let ExportedSymbol::NonGeneric(def_id) = symbol {
+                f(cnum, def_id)?;
+            }
+        }
+    }
+    Ok(())
+}
+
 /// Convert a softfloat type to its corresponding hostfloat type.
 pub trait ToHost {
     type HostFloat;
@@ -1180,6 +1217,37 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
         }
         Ok(())
     }
+
+    /// Lookup an array of immediates stored as a linker section of name `name`.
+    fn lookup_link_section(
+        &mut self,
+        name: &str,
+    ) -> InterpResult<'tcx, Vec<ImmTy<'tcx, Provenance>>> {
+        let this = self.eval_context_mut();
+        let tcx = this.tcx.tcx;
+
+        let mut array = vec![];
+
+        iter_exported_symbols(tcx, |_cnum, def_id| {
+            let attrs = tcx.codegen_fn_attrs(def_id);
+            let Some(link_section) = attrs.link_section else {
+                return Ok(());
+            };
+            if link_section.as_str() == name {
+                let instance = ty::Instance::mono(tcx, def_id);
+                let const_val = this.eval_global(instance).unwrap_or_else(|err| {
+                    panic!(
+                        "failed to evaluate static in required link_section: {def_id:?}\n{err:?}"
+                    )
+                });
+                let val = this.read_immediate(&const_val)?;
+                array.push(val);
+            }
+            Ok(())
+        })?;
+
+        Ok(array)
+    }
 }
 
 impl<'mir, 'tcx> MiriMachine<'mir, 'tcx> {

src/tools/miri/src/shims/foreign_items.rs

@@ -2,17 +2,10 @@ use std::{collections::hash_map::Entry, io::Write, iter, path::Path};
 
 use rustc_apfloat::Float;
 use rustc_ast::expand::allocator::AllocatorKind;
-use rustc_hir::{
-    def::DefKind,
-    def_id::{CrateNum, LOCAL_CRATE},
-};
-use rustc_middle::middle::{
-    codegen_fn_attrs::CodegenFnAttrFlags, dependency_format::Linkage,
-    exported_symbols::ExportedSymbol,
-};
+use rustc_hir::{def::DefKind, def_id::CrateNum};
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::mir;
 use rustc_middle::ty;
-use rustc_session::config::CrateType;
 use rustc_span::Symbol;
 use rustc_target::{
     abi::{Align, Size},
@@ -174,74 +167,48 @@ pub trait EvalContextExt<'mir, 'tcx: 'mir>: crate::MiriInterpCxExt<'mir, 'tcx> {
             Entry::Vacant(e) => {
                 // Find it if it was not cached.
                 let mut instance_and_crate: Option<(ty::Instance<'_>, CrateNum)> = None;
-                // `dependency_formats` includes all the transitive informations needed to link a crate,
-                // which is what we need here since we need to dig out `exported_symbols` from all transitive
-                // dependencies.
-                let dependency_formats = tcx.dependency_formats(());
-                let dependency_format = dependency_formats
-                    .iter()
-                    .find(|(crate_type, _)| *crate_type == CrateType::Executable)
-                    .expect("interpreting a non-executable crate");
-                for cnum in iter::once(LOCAL_CRATE).chain(
-                    dependency_format.1.iter().enumerate().filter_map(|(num, &linkage)| {
-                        // We add 1 to the number because that's what rustc also does everywhere it
-                        // calls `CrateNum::new`...
-                        #[allow(clippy::arithmetic_side_effects)]
-                        (linkage != Linkage::NotLinked).then_some(CrateNum::new(num + 1))
-                    }),
-                ) {
-                    // We can ignore `_export_info` here: we are a Rust crate, and everything is exported
-                    // from a Rust crate.
-                    for &(symbol, _export_info) in tcx.exported_symbols(cnum) {
-                        if let ExportedSymbol::NonGeneric(def_id) = symbol {
-                            let attrs = tcx.codegen_fn_attrs(def_id);
-                            let symbol_name = if let Some(export_name) = attrs.export_name {
-                                export_name
-                            } else if attrs.flags.contains(CodegenFnAttrFlags::NO_MANGLE) {
-                                tcx.item_name(def_id)
+                helpers::iter_exported_symbols(tcx, |cnum, def_id| {
+                    let attrs = tcx.codegen_fn_attrs(def_id);
+                    let symbol_name = if let Some(export_name) = attrs.export_name {
+                        export_name
+                    } else if attrs.flags.contains(CodegenFnAttrFlags::NO_MANGLE) {
+                        tcx.item_name(def_id)
+                    } else {
+                        // Skip over items without an explicitly defined symbol name.
+                        return Ok(());
+                    };
+                    if symbol_name == link_name {
+                        if let Some((original_instance, original_cnum)) = instance_and_crate {
+                            // Make sure we are consistent wrt what is 'first' and 'second'.
+                            let original_span = tcx.def_span(original_instance.def_id()).data();
+                            let span = tcx.def_span(def_id).data();
+                            if original_span < span {
+                                throw_machine_stop!(TerminationInfo::MultipleSymbolDefinitions {
+                                    link_name,
+                                    first: original_span,
+                                    first_crate: tcx.crate_name(original_cnum),
+                                    second: span,
+                                    second_crate: tcx.crate_name(cnum),
+                                });
                             } else {
-                                // Skip over items without an explicitly defined symbol name.
-                                continue;
-                            };
-                            if symbol_name == link_name {
-                                if let Some((original_instance, original_cnum)) = instance_and_crate
-                                {
-                                    // Make sure we are consistent wrt what is 'first' and 'second'.
-                                    let original_span =
-                                        tcx.def_span(original_instance.def_id()).data();
-                                    let span = tcx.def_span(def_id).data();
-                                    if original_span < span {
-                                        throw_machine_stop!(
-                                            TerminationInfo::MultipleSymbolDefinitions {
-                                                link_name,
-                                                first: original_span,
-                                                first_crate: tcx.crate_name(original_cnum),
-                                                second: span,
-                                                second_crate: tcx.crate_name(cnum),
-                                            }
-                                        );
-                                    } else {
-                                        throw_machine_stop!(
-                                            TerminationInfo::MultipleSymbolDefinitions {
-                                                link_name,
-                                                first: span,
-                                                first_crate: tcx.crate_name(cnum),
-                                                second: original_span,
-                                                second_crate: tcx.crate_name(original_cnum),
-                                            }
-                                        );
-                                    }
-                                }
-                                if !matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
-                                    throw_ub_format!(
-                                        "attempt to call an exported symbol that is not defined as a function"
-                                    );
-                                }
-                                instance_and_crate = Some((ty::Instance::mono(tcx, def_id), cnum));
+                                throw_machine_stop!(TerminationInfo::MultipleSymbolDefinitions {
+                                    link_name,
+                                    first: span,
+                                    first_crate: tcx.crate_name(cnum),
+                                    second: original_span,
+                                    second_crate: tcx.crate_name(original_cnum),
+                                });
                             }
                         }
+                        if !matches!(tcx.def_kind(def_id), DefKind::Fn | DefKind::AssocFn) {
+                            throw_ub_format!(
+                                "attempt to call an exported symbol that is not defined as a function"
+                            );
+                        }
+                        instance_and_crate = Some((ty::Instance::mono(tcx, def_id), cnum));
                     }
-                }
+                    Ok(())
+                })?;
 
                 e.insert(instance_and_crate.map(|ic| ic.0))
             }