Use maybe_create_entry_wrapper again in jit mode

bjorn3 · bjorn3 · commit 86530f889ead · 2021-04-14T14:59:36.000+02:00
This simplifies the jit driver a lot
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
@@ -135,7 +135,7 @@ fn module_codegen(
             }
         }
     }
-    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut cx.unwind_context);
+    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut cx.unwind_context, false);
 
     let codegen_result = emit_module(
         tcx,
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
@@ -8,7 +8,6 @@ use std::os::raw::{c_char, c_int};
 use cranelift_codegen::binemit::{NullStackMapSink, NullTrapSink};
 use rustc_codegen_ssa::CrateInfo;
 use rustc_middle::mir::mono::MonoItem;
-use rustc_session::config::EntryFnType;
 
 use cranelift_jit::{JITBuilder, JITModule};
 
@@ -81,6 +80,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
     }
 
     crate::allocator::codegen(tcx, &mut jit_module, &mut cx.unwind_context);
+    crate::main_shim::maybe_create_entry_wrapper(tcx, &mut jit_module, &mut cx.unwind_context, true);
 
     tcx.sess.abort_if_errors();
 
@@ -105,57 +105,24 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         assert!(tls_backend_config.borrow_mut().replace(backend_config).is_none())
     });
 
-    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
-
-            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);
-        }
-    }
+    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("main", 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]
diff --git a/src/main_shim.rs b/src/main_shim.rs
@@ -10,8 +10,9 @@ pub(crate) fn maybe_create_entry_wrapper(
     tcx: TyCtxt<'_>,
     module: &mut impl Module,
     unwind_context: &mut UnwindContext,
+    ignore_lang_start_wrapper: bool,
 ) {
-    let (main_def_id, use_start_lang_item) = match tcx.entry_fn(LOCAL_CRATE) {
+    let (main_def_id, is_main_fn) = match tcx.entry_fn(LOCAL_CRATE) {
         Some((def_id, entry_ty)) => (
             def_id.to_def_id(),
             match entry_ty {
@@ -27,14 +28,22 @@ pub(crate) fn maybe_create_entry_wrapper(
         return;
     }
 
-    create_entry_fn(tcx, module, unwind_context, main_def_id, use_start_lang_item);
+    create_entry_fn(
+        tcx,
+        module,
+        unwind_context,
+        main_def_id,
+        ignore_lang_start_wrapper,
+        is_main_fn,
+    );
 
     fn create_entry_fn(
         tcx: TyCtxt<'_>,
         m: &mut impl Module,
         unwind_context: &mut UnwindContext,
         rust_main_def_id: DefId,
-        use_start_lang_item: bool,
+        ignore_lang_start_wrapper: bool,
+        is_main_fn: bool,
     ) {
         let main_ret_ty = tcx.fn_sig(rust_main_def_id).output();
         // Given that `main()` has no arguments,
@@ -74,7 +83,12 @@ pub(crate) fn maybe_create_entry_wrapper(
 
             let main_func_ref = m.declare_func_in_func(main_func_id, &mut bcx.func);
 
-            let call_inst = if use_start_lang_item {
+            let result = if is_main_fn && ignore_lang_start_wrapper {
+                // regular main fn, but ignoring #[lang = "start"] as we are running in the jit
+                // FIXME set program arguments somehow
+                bcx.ins().call(main_func_ref, &[]);
+                bcx.ins().iconst(m.target_config().pointer_type(), 0)
+            } else if is_main_fn {
                 let start_def_id = tcx.require_lang_item(LangItem::Start, None);
                 let start_instance = Instance::resolve(
                     tcx,
@@ -90,13 +104,14 @@ pub(crate) fn maybe_create_entry_wrapper(
                 let main_val = bcx.ins().func_addr(m.target_config().pointer_type(), main_func_ref);
 
                 let func_ref = m.declare_func_in_func(start_func_id, &mut bcx.func);
-                bcx.ins().call(func_ref, &[main_val, arg_argc, arg_argv])
+                let call_inst = bcx.ins().call(func_ref, &[main_val, arg_argc, arg_argv]);
+                bcx.inst_results(call_inst)[0]
             } else {
                 // using user-defined start fn
-                bcx.ins().call(main_func_ref, &[arg_argc, arg_argv])
+                let call_inst = bcx.ins().call(main_func_ref, &[arg_argc, arg_argv]);
+                bcx.inst_results(call_inst)[0]
             };
 
-            let result = bcx.inst_results(call_inst)[0];
             bcx.ins().return_(&[result]);
             bcx.seal_all_blocks();
             bcx.finalize();

Original file line number	Diff line number	Diff line change
`@@ -135,7 +135,7 @@ fn module_codegen(`
`135`	`135`	`}`
`136`	`136`	`}`
`137`	`137`	`}`
`138`		`- crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut cx.unwind_context);`
	`138`	`+ crate::main_shim::maybe_create_entry_wrapper(tcx, &mut module, &mut cx.unwind_context, false);`
`139`	`139`
`140`	`140`	`let codegen_result = emit_module(`
`141`	`141`	`tcx,`