Directly invoke the main function in JIT mode

Fixes #1151

Author: bjorn3

src/driver/jit.rs

@@ -7,6 +7,7 @@ use std::os::raw::{c_char, c_int};
 
 use rustc_codegen_ssa::CrateInfo;
 use rustc_middle::mir::mono::MonoItem;
+use rustc_session::config::EntryFnType;
 
 use cranelift_jit::{JITBuilder, JITModule};
 
@@ -32,16 +33,6 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
     let mut jit_module = JITModule::new(jit_builder);
     assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
 
-    let sig = Signature {
-        params: vec![
-            AbiParam::new(jit_module.target_config().pointer_type()),
-            AbiParam::new(jit_module.target_config().pointer_type()),
-        ],
-        returns: vec![AbiParam::new(jit_module.target_config().pointer_type() /*isize*/)],
-        call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
-    };
-    let main_func_id = jit_module.declare_function("main", Linkage::Import, &sig).unwrap();
-
     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
     let mono_items = cgus
         .iter()
@@ -86,24 +77,17 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         tcx.sess.fatal("Inline asm is not supported in JIT mode");
     }
 
-    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut unwind_context);
     crate::allocator::codegen(tcx, &mut jit_module, &mut unwind_context);
 
     tcx.sess.abort_if_errors();
 
     jit_module.finalize_definitions();
-
     let _unwind_register_guard = unsafe { unwind_context.register_jit(&jit_module) };
 
-    let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
-
     println!(
         "Rustc codegen cranelift will JIT run the executable, because -Cllvm-args=mode=jit was passed"
     );
 
-    let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
-        unsafe { ::std::mem::transmute(finalized_main) };
-
     let args = ::std::env::var("CG_CLIF_JIT_ARGS").unwrap_or_else(|_| String::new());
     let args = std::iter::once(&*tcx.crate_name(LOCAL_CRATE).as_str().to_string())
         .chain(args.split(' '))
@@ -118,12 +102,58 @@ pub(super) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
     BACKEND_CONFIG.with(|tls_backend_config| {
         assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
     });
-    CURRENT_MODULE
-        .with(|current_module| assert!(current_module.borrow_mut().replace(jit_module).is_none()));
 
-    let ret = f(args.len() as c_int, argv.as_ptr());
+    let (main_def_id, entry_ty) = tcx.entry_fn(LOCAL_CRATE).unwrap();
+    let instance = Instance::mono(tcx, main_def_id.to_def_id()).polymorphize(tcx);
+
+    match entry_ty {
+        EntryFnType::Main => {
+            // FIXME set program arguments somehow
 
-    std::process::exit(ret);
+            let main_sig = Signature {
+                params: vec![],
+                returns: vec![],
+                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
+            };
+            let main_func_id = jit_module
+                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &main_sig)
+                .unwrap();
+            let finalized_main: *const u8 = jit_module.get_finalized_function(main_func_id);
+
+            CURRENT_MODULE.with(|current_module| {
+                assert!(current_module.borrow_mut().replace(jit_module).is_none())
+            });
+
+            let f: extern "C" fn() = unsafe { ::std::mem::transmute(finalized_main) };
+            f();
+            std::process::exit(0);
+        }
+        EntryFnType::Start => {
+            let start_sig = Signature {
+                params: vec![
+                    AbiParam::new(jit_module.target_config().pointer_type()),
+                    AbiParam::new(jit_module.target_config().pointer_type()),
+                ],
+                returns: vec![AbiParam::new(
+                    jit_module.target_config().pointer_type(), /*isize*/
+                )],
+                call_conv: CallConv::triple_default(&crate::target_triple(tcx.sess)),
+            };
+            let start_func_id = jit_module
+                .declare_function(tcx.symbol_name(instance).name, Linkage::Import, &start_sig)
+                .unwrap();
+            let finalized_start: *const u8 = jit_module.get_finalized_function(start_func_id);
+
+            CURRENT_MODULE.with(|current_module| {
+                assert!(current_module.borrow_mut().replace(jit_module).is_none())
+            });
+
+            let f: extern "C" fn(c_int, *const *const c_char) -> c_int =
+                unsafe { ::std::mem::transmute(finalized_start) };
+            let ret = f(args.len() as c_int, argv.as_ptr());
+            std::process::exit(ret);
+        }
+    }
 }
 
 #[no_mangle]
@@ -220,7 +250,7 @@ fn load_imported_symbols_for_jit(tcx: TyCtxt<'_>) -> Vec<(String, *const u8)> {
     imported_symbols
 }
 
-pub(super) fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
+fn codegen_shim<'tcx>(cx: &mut CodegenCx<'_, 'tcx>, inst: Instance<'tcx>) {
     let tcx = cx.tcx;
 
     let pointer_type = cx.module.target_config().pointer_type();