4 files changed, 37 insertions, 32 deletions

diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index 2ad580d6330..3ff96d710e5 100644
--- a/src/driver/aot.rs
+++ b/src/driver/aot.rs
@@ -107,7 +107,6 @@ fn module_codegen(
 
     let isa = crate::build_isa(tcx.sess, &backend_config);
     let mut module = crate::backend::make_module(tcx.sess, isa, cgu_name.as_str().to_string());
-    assert_eq!(pointer_ty(tcx), module.target_config().pointer_type());
 
     let mut cx = crate::CodegenCx::new(
         tcx,
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index 4b12c1aec99..f7fe02bb016 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -23,6 +23,33 @@ thread_local! {
     static LAZY_JIT_STATE: RefCell<Option<JitState>> = RefCell::new(None);
 }
 
+fn create_jit_module<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    backend_config: &BackendConfig,
+    hotswap: bool,
+) -> (JITModule, CodegenCx<'tcx>) {
+    let imported_symbols = load_imported_symbols_for_jit(tcx);
+
+    let isa = crate::build_isa(tcx.sess, backend_config);
+    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
+    jit_builder.hotswap(hotswap);
+    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
+    jit_builder.symbols(imported_symbols);
+    let mut jit_module = JITModule::new(jit_builder);
+
+    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);
+
+    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,
+    );
+
+    (jit_module, cx)
+}
+
 pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
     if !tcx.sess.opts.output_types.should_codegen() {
         tcx.sess.fatal("JIT mode doesn't work with `cargo check`");
@@ -32,15 +59,11 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         tcx.sess.fatal("can't jit non-executable crate");
     }
 
-    let imported_symbols = load_imported_symbols_for_jit(tcx);
-
-    let isa = crate::build_isa(tcx.sess, &backend_config);
-    let mut jit_builder = JITBuilder::with_isa(isa, cranelift_module::default_libcall_names());
-    jit_builder.hotswap(matches!(backend_config.codegen_mode, CodegenMode::JitLazy));
-    crate::compiler_builtins::register_functions_for_jit(&mut jit_builder);
-    jit_builder.symbols(imported_symbols);
-    let mut jit_module = JITModule::new(jit_builder);
-    assert_eq!(pointer_ty(tcx), jit_module.target_config().pointer_type());
+    let (mut jit_module, mut cx) = create_jit_module(
+        tcx,
+        &backend_config,
+        matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
+    );
 
     let (_, cgus) = tcx.collect_and_partition_mono_items(LOCAL_CRATE);
     let mono_items = cgus
@@ -51,8 +74,6 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         .into_iter()
         .collect::<Vec<(_, (_, _))>>();
 
-    let mut cx = crate::CodegenCx::new(tcx, backend_config.clone(), jit_module.isa(), false);
-
     super::time(tcx, backend_config.display_cg_time, "codegen mono items", || {
         super::predefine_mono_items(tcx, &mut jit_module, &mono_items);
         for (mono_item, _) in mono_items {
@@ -77,20 +98,10 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         }
     });
 
-    jit_module.finalize_definitions();
-
     if !cx.global_asm.is_empty() {
         tcx.sess.fatal("Inline asm is not supported in JIT mode");
     }
 
-    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();
 
     jit_module.finalize_definitions();
diff --git a/src/lib.rs b/src/lib.rs
index 12d8f9c1d73..2c989520f6a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -134,6 +134,8 @@ impl<'tcx> CodegenCx<'tcx> {
         isa: &dyn TargetIsa,
         debug_info: bool,
     ) -> Self {
+        assert_eq!(pointer_ty(tcx), isa.pointer_type());
+
         let unwind_context =
             UnwindContext::new(tcx, isa, matches!(backend_config.codegen_mode, CodegenMode::Aot));
         let debug_context = if debug_info { Some(DebugContext::new(tcx, isa)) } else { None };
diff --git a/src/main_shim.rs b/src/main_shim.rs
index ff386710cd1..4beb6c65a77 100644
--- a/src/main_shim.rs
+++ b/src/main_shim.rs
@@ -10,7 +10,7 @@ pub(crate) fn maybe_create_entry_wrapper(
     tcx: TyCtxt<'_>,
     module: &mut impl Module,
     unwind_context: &mut UnwindContext,
-    ignore_lang_start_wrapper: bool,
+    is_jit: bool,
 ) {
     let (main_def_id, is_main_fn) = match tcx.entry_fn(LOCAL_CRATE) {
         Some((def_id, entry_ty)) => (
@@ -24,18 +24,11 @@ pub(crate) fn maybe_create_entry_wrapper(
     };
 
     let instance = Instance::mono(tcx, main_def_id).polymorphize(tcx);
-    if module.get_name(&*tcx.symbol_name(instance).name).is_none() {
+    if !is_jit && module.get_name(&*tcx.symbol_name(instance).name).is_none() {
         return;
     }
 
-    create_entry_fn(
-        tcx,
-        module,
-        unwind_context,
-        main_def_id,
-        ignore_lang_start_wrapper,
-        is_main_fn,
-    );
+    create_entry_fn(tcx, module, unwind_context, main_def_id, is_jit, is_main_fn);
 
     fn create_entry_fn(
         tcx: TyCtxt<'_>,