add support for global constructors (i.e. life before main)

5 files changed, 165 insertions, 32 deletions

diff --git a/src/tools/miri/src/eval.rs b/src/tools/miri/src/eval.rs
index 425a136dfa5..d986f2830aa 100644
--- a/src/tools/miri/src/eval.rs
+++ b/src/tools/miri/src/eval.rs
@@ -11,14 +11,14 @@ use rustc_abi::ExternAbi;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_hir::def::Namespace;
 use rustc_hir::def_id::DefId;
-use rustc_middle::ty::layout::LayoutCx;
+use rustc_middle::ty::layout::{HasTyCtxt, HasTypingEnv, LayoutCx};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_session::config::EntryFnType;
 
 use crate::concurrency::GenmcCtx;
 use crate::concurrency::thread::TlsAllocAction;
 use crate::diagnostics::report_leaks;
-use crate::shims::tls;
+use crate::shims::{ctor, tls};
 use crate::*;
 
 #[derive(Copy, Clone, Debug)]
@@ -216,9 +216,15 @@ impl Default for MiriConfig {
 }
 
 /// The state of the main thread. Implementation detail of `on_main_stack_empty`.
-#[derive(Default, Debug)]
+#[derive(Debug)]
 enum MainThreadState<'tcx> {
-    #[default]
+    GlobalCtors {
+        ctor_state: ctor::GlobalCtorState<'tcx>,
+        entry_id: DefId,
+        entry_type: MiriEntryFnType,
+        argc: ImmTy<'tcx>,
+        argv: ImmTy<'tcx>,
+    },
     Running,
     TlsDtors(tls::TlsDtorsState<'tcx>),
     Yield {
@@ -234,6 +240,15 @@ impl<'tcx> MainThreadState<'tcx> {
     ) -> InterpResult<'tcx, Poll<()>> {
         use MainThreadState::*;
         match self {
+            GlobalCtors { ctor_state, entry_id, entry_type, argc, argv } => {
+                match ctor_state.on_stack_empty(this)? {
+                    Poll::Pending => {} // just keep going
+                    Poll::Ready(()) => {
+                        call_main(this, *entry_id, *entry_type, argc.clone(), argv.clone())?;
+                        *self = Running;
+                    }
+                }
+            }
             Running => {
                 *self = TlsDtors(Default::default());
             }
@@ -309,26 +324,6 @@ pub fn create_ecx<'tcx>(
         MiriMachine::new(config, layout_cx, genmc_ctx),
     );
 
-    // Some parts of initialization require a full `InterpCx`.
-    MiriMachine::late_init(&mut ecx, config, {
-        let mut state = MainThreadState::default();
-        // Cannot capture anything GC-relevant here.
-        Box::new(move |m| state.on_main_stack_empty(m))
-    })?;
-
-    // Make sure we have MIR. We check MIR for some stable monomorphic function in libcore.
-    let sentinel =
-        helpers::try_resolve_path(tcx, &["core", "ascii", "escape_default"], Namespace::ValueNS);
-    if !matches!(sentinel, Some(s) if tcx.is_mir_available(s.def.def_id())) {
-        tcx.dcx().fatal(
-            "the current sysroot was built without `-Zalways-encode-mir`, or libcore seems missing. \
-            Use `cargo miri setup` to prepare a sysroot that is suitable for Miri."
-        );
-    }
-
-    // Setup first stack frame.
-    let entry_instance = ty::Instance::mono(tcx, entry_id);
-
     // First argument is constructed later, because it's skipped for `miri_start.`
 
     // Second argument (argc): length of `config.args`.
@@ -395,11 +390,51 @@ pub fn create_ecx<'tcx>(
         ImmTy::from_immediate(imm, layout)
     };
 
+    // Some parts of initialization require a full `InterpCx`.
+    MiriMachine::late_init(&mut ecx, config, {
+        let mut state = MainThreadState::GlobalCtors {
+            entry_id,
+            entry_type,
+            argc,
+            argv,
+            ctor_state: ctor::GlobalCtorState::default(),
+        };
+
+        // Cannot capture anything GC-relevant here.
+        Box::new(move |m| state.on_main_stack_empty(m))
+    })?;
+
+    // Make sure we have MIR. We check MIR for some stable monomorphic function in libcore.
+    let sentinel =
+        helpers::try_resolve_path(tcx, &["core", "ascii", "escape_default"], Namespace::ValueNS);
+    if !matches!(sentinel, Some(s) if tcx.is_mir_available(s.def.def_id())) {
+        tcx.dcx().fatal(
+            "the current sysroot was built without `-Zalways-encode-mir`, or libcore seems missing. \
+            Use `cargo miri setup` to prepare a sysroot that is suitable for Miri."
+        );
+    }
+
+    interp_ok(ecx)
+}
+
+// Call the entry function.
+fn call_main<'tcx>(
+    ecx: &mut MiriInterpCx<'tcx>,
+    entry_id: DefId,
+    entry_type: MiriEntryFnType,
+    argc: ImmTy<'tcx>,
+    argv: ImmTy<'tcx>,
+) -> InterpResult<'tcx, ()> {
+    let tcx = ecx.tcx();
+
+    // Setup first stack frame.
+    let entry_instance = ty::Instance::mono(tcx, entry_id);
+
     // Return place (in static memory so that it does not count as leak).
     let ret_place = ecx.allocate(ecx.machine.layouts.isize, MiriMemoryKind::Machine.into())?;
     ecx.machine.main_fn_ret_place = Some(ret_place.clone());
-    // Call start function.
 
+    // Call start function.
     match entry_type {
         MiriEntryFnType::Rustc(EntryFnType::Main { .. }) => {
             let start_id = tcx.lang_items().start_fn().unwrap_or_else(|| {
@@ -409,7 +444,7 @@ pub fn create_ecx<'tcx>(
             let main_ret_ty = main_ret_ty.no_bound_vars().unwrap();
             let start_instance = ty::Instance::try_resolve(
                 tcx,
-                typing_env,
+                ecx.typing_env(),
                 start_id,
                 tcx.mk_args(&[ty::GenericArg::from(main_ret_ty)]),
             )
@@ -427,7 +462,7 @@ pub fn create_ecx<'tcx>(
                 ExternAbi::Rust,
                 &[
                     ImmTy::from_scalar(
-                        Scalar::from_pointer(main_ptr, &ecx),
+                        Scalar::from_pointer(main_ptr, ecx),
                         // FIXME use a proper fn ptr type
                         ecx.machine.layouts.const_raw_ptr,
                     ),
@@ -450,7 +485,7 @@ pub fn create_ecx<'tcx>(
         }
     }
 
-    interp_ok(ecx)
+    interp_ok(())
 }
 
 /// Evaluates the entry function specified by `entry_id`.
diff --git a/src/tools/miri/src/helpers.rs b/src/tools/miri/src/helpers.rs
index 216b7b1e4bb..35d3e7e647c 100644
--- a/src/tools/miri/src/helpers.rs
+++ b/src/tools/miri/src/helpers.rs
@@ -1235,8 +1235,11 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         interp_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>>> {
+    /// Lookup an array of immediates from any linker sections matching the provided predicate.
+    fn lookup_link_section(
+        &mut self,
+        include_name: impl Fn(&str) -> bool,
+    ) -> InterpResult<'tcx, Vec<ImmTy<'tcx>>> {
         let this = self.eval_context_mut();
         let tcx = this.tcx.tcx;
 
@@ -1247,7 +1250,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
             let Some(link_section) = attrs.link_section else {
                 return interp_ok(());
             };
-            if link_section.as_str() == name {
+            if include_name(link_section.as_str()) {
                 let instance = ty::Instance::mono(tcx, def_id);
                 let const_val = this.eval_global(instance).unwrap_or_else(|err| {
                     panic!(
diff --git a/src/tools/miri/src/shims/ctor.rs b/src/tools/miri/src/shims/ctor.rs
new file mode 100644
index 00000000000..75f564989bc
--- /dev/null
+++ b/src/tools/miri/src/shims/ctor.rs
@@ -0,0 +1,94 @@
+//! Implement global constructors.
+
+use std::task::Poll;
+
+use rustc_abi::ExternAbi;
+use rustc_target::spec::BinaryFormat;
+
+use crate::*;
+
+#[derive(Debug, Default)]
+pub struct GlobalCtorState<'tcx>(GlobalCtorStatePriv<'tcx>);
+
+#[derive(Debug, Default)]
+enum GlobalCtorStatePriv<'tcx> {
+    #[default]
+    Init,
+    /// The list of constructor functions that we still have to call.
+    Ctors(Vec<ImmTy<'tcx>>),
+    Done,
+}
+
+impl<'tcx> GlobalCtorState<'tcx> {
+    pub fn on_stack_empty(
+        &mut self,
+        this: &mut MiriInterpCx<'tcx>,
+    ) -> InterpResult<'tcx, Poll<()>> {
+        use GlobalCtorStatePriv::*;
+        let new_state = 'new_state: {
+            match &mut self.0 {
+                Init => {
+                    let this = this.eval_context_mut();
+
+                    // Lookup constructors from the relevant magic link section.
+                    let ctors = match this.tcx.sess.target.binary_format {
+                        // Read the CRT library section on Windows.
+                        BinaryFormat::Coff =>
+                            this.lookup_link_section(|section| section == ".CRT$XCU")?,
+
+                        // Read the `__mod_init_func` section on macOS.
+                        BinaryFormat::MachO =>
+                            this.lookup_link_section(|section| {
+                                let mut parts = section.splitn(3, ',');
+                                let (segment_name, section_name, section_type) =
+                                    (parts.next(), parts.next(), parts.next());
+
+                                segment_name == Some("__DATA")
+                                    && section_name == Some("__mod_init_func")
+                                    // The `mod_init_funcs` directive ensures that the `S_MOD_INIT_FUNC_POINTERS` flag
+                                    // is set on the section, but it is not strictly required.
+                                    && matches!(section_type, None | Some("mod_init_funcs"))
+                            })?,
+
+                        // Read the standard `.init_array` section on platforms that use ELF, or WASM,
+                        // which supports the same linker directive.
+                        // FIXME: Add support for `.init_array.N`.
+                        BinaryFormat::Elf | BinaryFormat::Wasm =>
+                            this.lookup_link_section(|section| section == ".init_array")?,
+
+                        // Other platforms have no global ctor support.
+                        _ => break 'new_state Done,
+                    };
+
+                    break 'new_state Ctors(ctors);
+                }
+                Ctors(ctors) => {
+                    if let Some(ctor) = ctors.pop() {
+                        let this = this.eval_context_mut();
+
+                        let ctor = ctor.to_scalar().to_pointer(this)?;
+                        let thread_callback = this.get_ptr_fn(ctor)?.as_instance()?;
+
+                        // The signature of this function is `unsafe extern "C" fn()`.
+                        this.call_function(
+                            thread_callback,
+                            ExternAbi::C { unwind: false },
+                            &[],
+                            None,
+                            ReturnContinuation::Stop { cleanup: true },
+                        )?;
+
+                        return interp_ok(Poll::Pending); // we stay in this state (but `ctors` got shorter)
+                    }
+
+                    // No more constructors to run.
+                    break 'new_state Done;
+                }
+                Done => return interp_ok(Poll::Ready(())),
+            }
+        };
+
+        self.0 = new_state;
+        interp_ok(Poll::Pending)
+    }
+}
diff --git a/src/tools/miri/src/shims/mod.rs b/src/tools/miri/src/shims/mod.rs
index b08ab101e94..0faf298bfaa 100644
--- a/src/tools/miri/src/shims/mod.rs
+++ b/src/tools/miri/src/shims/mod.rs
@@ -11,6 +11,7 @@ mod wasi;
 mod windows;
 mod x86;
 
+pub mod ctor;
 pub mod env;
 pub mod extern_static;
 pub mod foreign_items;
diff --git a/src/tools/miri/src/shims/tls.rs b/src/tools/miri/src/shims/tls.rs
index 7182637437a..1200029692d 100644
--- a/src/tools/miri/src/shims/tls.rs
+++ b/src/tools/miri/src/shims/tls.rs
@@ -302,7 +302,7 @@ trait EvalContextPrivExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
 
         // Windows has a special magic linker section that is run on certain events.
         // We don't support most of that, but just enough to make thread-local dtors in `std` work.
-        interp_ok(this.lookup_link_section(".CRT$XLB")?)
+        interp_ok(this.lookup_link_section(|section| section == ".CRT$XLB")?)
     }
 
     fn schedule_windows_tls_dtor(&mut self, dtor: ImmTy<'tcx>) -> InterpResult<'tcx> {