Add `miri_start` support

7 files changed, 125 insertions, 11 deletions

diff --git a/src/tools/miri/src/bin/miri.rs b/src/tools/miri/src/bin/miri.rs
index 25b154a8206..2cea8e75117 100644
--- a/src/tools/miri/src/bin/miri.rs
+++ b/src/tools/miri/src/bin/miri.rs
@@ -14,11 +14,14 @@ extern crate tracing;
 extern crate rustc_data_structures;
 extern crate rustc_driver;
 extern crate rustc_hir;
+extern crate rustc_hir_analysis;
 extern crate rustc_interface;
 extern crate rustc_log;
 extern crate rustc_metadata;
 extern crate rustc_middle;
 extern crate rustc_session;
+extern crate rustc_span;
+extern crate rustc_target;
 
 use std::env::{self, VarError};
 use std::num::NonZero;
@@ -27,9 +30,12 @@ use std::str::FromStr;
 
 use tracing::debug;
 
+use miri::{BacktraceStyle, BorrowTrackerMethod, ProvenanceMode, RetagFields};
 use rustc_data_structures::sync::Lrc;
 use rustc_driver::Compilation;
+use rustc_hir::def_id::LOCAL_CRATE;
 use rustc_hir::{self as hir, Node};
+use rustc_hir_analysis::check::check_function_signature;
 use rustc_interface::interface::Config;
 use rustc_middle::{
     middle::{
@@ -37,14 +43,15 @@ use rustc_middle::{
         exported_symbols::{ExportedSymbol, SymbolExportInfo, SymbolExportKind, SymbolExportLevel},
     },
     query::LocalCrate,
-    ty::TyCtxt,
+    traits::{ObligationCause, ObligationCauseCode},
+    ty::{self, Ty, TyCtxt},
     util::Providers,
 };
-use rustc_session::config::{CrateType, ErrorOutputType, OptLevel};
+use rustc_session::config::{CrateType, EntryFnType, ErrorOutputType, OptLevel};
 use rustc_session::search_paths::PathKind;
 use rustc_session::{CtfeBacktrace, EarlyDiagCtxt};
-
-use miri::{BacktraceStyle, BorrowTrackerMethod, ProvenanceMode, RetagFields};
+use rustc_span::def_id::DefId;
+use rustc_target::spec::abi::Abi;
 
 struct MiriCompilerCalls {
     miri_config: miri::MiriConfig,
@@ -82,11 +89,7 @@ impl rustc_driver::Callbacks for MiriCompilerCalls {
                 tcx.dcx().fatal("miri only makes sense on bin crates");
             }
 
-            let (entry_def_id, entry_type) = if let Some(entry_def) = tcx.entry_fn(()) {
-                entry_def
-            } else {
-                tcx.dcx().fatal("miri can only run programs that have a main function");
-            };
+            let (entry_def_id, entry_type) = entry_fn(tcx);
             let mut config = self.miri_config.clone();
 
             // Add filename to `miri` arguments.
@@ -351,6 +354,54 @@ fn jemalloc_magic() {
     }
 }
 
+fn entry_fn(tcx: TyCtxt<'_>) -> (DefId, EntryFnType) {
+    if let Some(entry_def) = tcx.entry_fn(()) {
+        return entry_def;
+    }
+    // Look for a symbol in the local crate named `miri_start`, and threat that as the entry point.
+    let sym = tcx.exported_symbols(LOCAL_CRATE).iter().find_map(|(sym, _)| {
+        if sym.symbol_name_for_local_instance(tcx).name == "miri_start" { Some(sym) } else { None }
+    });
+    if let Some(ExportedSymbol::NonGeneric(id)) = sym {
+        let start_def_id = id.expect_local();
+        let start_span = tcx.def_span(start_def_id);
+
+        let expected_sig = ty::Binder::dummy(tcx.mk_fn_sig(
+            [tcx.types.isize, Ty::new_imm_ptr(tcx, Ty::new_imm_ptr(tcx, tcx.types.u8))],
+            tcx.types.isize,
+            false,
+            hir::Safety::Safe,
+            Abi::Rust,
+        ));
+
+        let correct_func_sig = check_function_signature(
+            tcx,
+            ObligationCause::new(start_span, start_def_id, ObligationCauseCode::Misc),
+            *id,
+            expected_sig,
+        )
+        .is_ok();
+
+        if correct_func_sig {
+            (*id, EntryFnType::Start)
+        } else {
+            tcx.dcx().fatal(
+                "`miri_start` must have the following signature:\n\
+                        fn miri_start(argc: isize, argv: *const *const u8) -> isize",
+            );
+        }
+    } else {
+        tcx.dcx().fatal("Miri can only run programs that have a main function.\n\
+                    Alternatively, you can export a `miri_start` function:\n\
+                    \n\
+                    #[cfg(miri)]\n\
+                    #[no_mangle]\n\
+                    fn miri_start(argc: isize, argv: *const *const u8) -> isize {\
+                    \n    // Call the actual start function that your project implements, based on your target's conventions.\n\
+                    }");
+    }
+}
+
 fn main() {
     #[cfg(any(target_os = "linux", target_os = "macos"))]
     jemalloc_magic();
diff --git a/src/tools/miri/tests/fail/miri_start_wrong_sig.rs b/src/tools/miri/tests/fail/miri_start_wrong_sig.rs
new file mode 100644
index 00000000000..dac83d817ba
--- /dev/null
+++ b/src/tools/miri/tests/fail/miri_start_wrong_sig.rs
@@ -0,0 +1,21 @@
+//@compile-flags: -Cpanic=abort
+//@error-in-other-file: `miri_start` must have the following signature:
+#![no_main]
+#![no_std]
+
+use core::fmt::Write;
+
+#[path = "../utils/mod.no_std.rs"]
+mod utils;
+
+#[no_mangle]
+fn miri_start() -> isize {
+    //~^ ERROR: mismatched types
+    writeln!(utils::MiriStdout, "Hello from miri_start!").unwrap();
+    0
+}
+
+#[panic_handler]
+fn panic_handler(_: &core::panic::PanicInfo) -> ! {
+    loop {}
+}
diff --git a/src/tools/miri/tests/fail/miri_start_wrong_sig.stderr b/src/tools/miri/tests/fail/miri_start_wrong_sig.stderr
new file mode 100644
index 00000000000..62171917116
--- /dev/null
+++ b/src/tools/miri/tests/fail/miri_start_wrong_sig.stderr
@@ -0,0 +1,15 @@
+error[E0308]: mismatched types
+  --> $DIR/miri_start_wrong_sig.rs:LL:CC
+   |
+LL | fn miri_start() -> isize {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^ incorrect number of function parameters
+   |
+   = note: expected signature `fn(isize, *const *const u8) -> _`
+              found signature `fn() -> _`
+
+error: `miri_start` must have the following signature:
+       fn miri_start(argc: isize, argv: *const *const u8) -> isize
+
+error: aborting due to 2 previous errors
+
+For more information about this error, try `rustc --explain E0308`.
diff --git a/src/tools/miri/tests/fail/no_main.rs b/src/tools/miri/tests/fail/no_main.rs
index 01b8c7bd66b..4bbd820a636 100644
--- a/src/tools/miri/tests/fail/no_main.rs
+++ b/src/tools/miri/tests/fail/no_main.rs
@@ -1,2 +1,2 @@
-//@error-in-other-file: miri can only run programs that have a main function
+//@error-in-other-file: Miri can only run programs that have a main function.
 #![no_main]
diff --git a/src/tools/miri/tests/fail/no_main.stderr b/src/tools/miri/tests/fail/no_main.stderr
index 1c4fc88989b..e9b9e5d65b1 100644
--- a/src/tools/miri/tests/fail/no_main.stderr
+++ b/src/tools/miri/tests/fail/no_main.stderr
@@ -1,4 +1,11 @@
-error: miri can only run programs that have a main function
+error: Miri can only run programs that have a main function.
+       Alternatively, you can export a `miri_start` function:
+       
+       #[cfg(miri)]
+       #[no_mangle]
+       fn miri_start(argc: isize, argv: *const *const u8) -> isize {
+           // Call the actual start function that your project implements, based on your target's conventions.
+       }
 
 error: aborting due to 1 previous error
 
diff --git a/src/tools/miri/tests/pass/miri_start.rs b/src/tools/miri/tests/pass/miri_start.rs
new file mode 100644
index 00000000000..756a1f60be1
--- /dev/null
+++ b/src/tools/miri/tests/pass/miri_start.rs
@@ -0,0 +1,19 @@
+//@compile-flags: -Cpanic=abort
+#![no_main]
+#![no_std]
+
+use core::fmt::Write;
+
+#[path = "../utils/mod.no_std.rs"]
+mod utils;
+
+#[no_mangle]
+fn miri_start(_argc: isize, _argv: *const *const u8) -> isize {
+    writeln!(utils::MiriStdout, "Hello from miri_start!").unwrap();
+    0
+}
+
+#[panic_handler]
+fn panic_handler(_: &core::panic::PanicInfo) -> ! {
+    loop {}
+}
diff --git a/src/tools/miri/tests/pass/miri_start.stdout b/src/tools/miri/tests/pass/miri_start.stdout
new file mode 100644
index 00000000000..1c9e8489b57
--- /dev/null
+++ b/src/tools/miri/tests/pass/miri_start.stdout
@@ -0,0 +1 @@
+Hello from miri_start!