From 66a885b25dbaf01232834a4c81344422750a5c6b Mon Sep 17 00:00:00 2001
From: Ralf Jung <post@ralfj.de>
Date: Sat, 29 Jun 2024 12:17:10 +0200
Subject: [PATCH] iter_exported_symbols: also walk used statics in local crate

---
 src/tools/miri/src/helpers.rs                 | 34 ++++++++++++++-----
 .../miri/tests/pass/tls/win_tls_callback.rs   | 16 +++++++++
 .../tests/pass/tls/win_tls_callback.stderr    |  1 +
 3 files changed, 43 insertions(+), 8 deletions(-)
 create mode 100644 src/tools/miri/tests/pass/tls/win_tls_callback.rs
 create mode 100644 src/tools/miri/tests/pass/tls/win_tls_callback.stderr

diff --git a/src/tools/miri/src/helpers.rs b/src/tools/miri/src/helpers.rs
index 3d2b102b279..a7a6f8cfd87 100644
--- a/src/tools/miri/src/helpers.rs
+++ b/src/tools/miri/src/helpers.rs
@@ -14,6 +14,7 @@
     def_id::{CrateNum, DefId, CRATE_DEF_INDEX, LOCAL_CRATE},
 };
 use rustc_index::IndexVec;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 use rustc_middle::middle::dependency_format::Linkage;
 use rustc_middle::middle::exported_symbols::ExportedSymbol;
 use rustc_middle::mir;
@@ -163,22 +164,39 @@ pub fn iter_exported_symbols<'tcx>(
     tcx: TyCtxt<'tcx>,
     mut f: impl FnMut(CrateNum, DefId) -> InterpResult<'tcx>,
 ) -> InterpResult<'tcx> {
+    // First, the symbols in the local crate. We can't use `exported_symbols` here as that
+    // skips `#[used]` statics (since `reachable_set` skips them in binary crates).
+    // So we walk all HIR items ourselves instead.
+    let crate_items = tcx.hir_crate_items(());
+    for def_id in crate_items.definitions() {
+        let exported = tcx.def_kind(def_id).has_codegen_attrs() && {
+            let codegen_attrs = tcx.codegen_fn_attrs(def_id);
+            codegen_attrs.contains_extern_indicator()
+                || codegen_attrs.flags.contains(CodegenFnAttrFlags::RUSTC_STD_INTERNAL_SYMBOL)
+                || codegen_attrs.flags.contains(CodegenFnAttrFlags::USED)
+                || codegen_attrs.flags.contains(CodegenFnAttrFlags::USED_LINKER)
+        };
+        if exported {
+            f(LOCAL_CRATE, def_id.into())?;
+        }
+    }
+
+    // Next, all our dependencies.
     // `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(());
+    // Find the dependencies of the executable we are running.
     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))
-        },
-    )) {
+    for cnum in 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) {
diff --git a/src/tools/miri/tests/pass/tls/win_tls_callback.rs b/src/tools/miri/tests/pass/tls/win_tls_callback.rs
new file mode 100644
index 00000000000..99a8de29e91
--- /dev/null
+++ b/src/tools/miri/tests/pass/tls/win_tls_callback.rs
@@ -0,0 +1,16 @@
+//! Ensure that we call Windows TLS callbacks in the local crate.
+//@only-target-windows
+// Calling eprintln in the callback seems to (re-)initialize some thread-local storage
+// and then leak the memory allocated for that. Let's just ignore these leaks,
+// that's not what this test is about.
+//@compile-flags: -Zmiri-ignore-leaks
+
+#[link_section = ".CRT$XLB"]
+#[used] // Miri only considers explicitly `#[used]` statics for `lookup_link_section`
+pub static CALLBACK: unsafe extern "system" fn(*const (), u32, *const ()) = tls_callback;
+
+unsafe extern "system" fn tls_callback(_h: *const (), _dw_reason: u32, _pv: *const ()) {
+    eprintln!("in tls_callback");
+}
+
+fn main() {}
diff --git a/src/tools/miri/tests/pass/tls/win_tls_callback.stderr b/src/tools/miri/tests/pass/tls/win_tls_callback.stderr
new file mode 100644
index 00000000000..84795589544
--- /dev/null
+++ b/src/tools/miri/tests/pass/tls/win_tls_callback.stderr
@@ -0,0 +1 @@
+in tls_callback