1 files changed, 53 insertions, 47 deletions

diff --git a/compiler/rustc_codegen_cranelift/src/driver/jit.rs b/compiler/rustc_codegen_cranelift/src/driver/jit.rs
index a56a91000..6a430b521 100644
--- a/compiler/rustc_codegen_cranelift/src/driver/jit.rs
+++ b/compiler/rustc_codegen_cranelift/src/driver/jit.rs
@@ -61,19 +61,18 @@ impl UnsafeMessage {
     }
 }
 
-fn create_jit_module<'tcx>(
-    tcx: TyCtxt<'tcx>,
+fn create_jit_module(
+    tcx: TyCtxt<'_>,
     backend_config: &BackendConfig,
     hotswap: bool,
-) -> (JITModule, CodegenCx<'tcx>) {
+) -> (JITModule, CodegenCx) {
     let crate_info = CrateInfo::new(tcx, "dummy_target_cpu".to_string());
-    let imported_symbols = load_imported_symbols_for_jit(tcx.sess, crate_info);
 
     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);
+    jit_builder.symbol_lookup_fn(dep_symbol_lookup_fn(tcx.sess, crate_info));
     jit_builder.symbol("__clif_jit_fn", clif_jit_fn as *const u8);
     let mut jit_module = JITModule::new(jit_builder);
 
@@ -111,6 +110,7 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
         &backend_config,
         matches!(backend_config.codegen_mode, CodegenMode::JitLazy),
     );
+    let mut cached_context = Context::new();
 
     let (_, cgus) = tcx.collect_and_partition_mono_items(());
     let mono_items = cgus
@@ -128,11 +128,19 @@ pub(crate) fn run_jit(tcx: TyCtxt<'_>, backend_config: BackendConfig) -> ! {
                 MonoItem::Fn(inst) => match backend_config.codegen_mode {
                     CodegenMode::Aot => unreachable!(),
                     CodegenMode::Jit => {
-                        cx.tcx.sess.time("codegen fn", || {
-                            crate::base::codegen_fn(&mut cx, &mut jit_module, inst)
+                        tcx.sess.time("codegen fn", || {
+                            crate::base::codegen_and_compile_fn(
+                                tcx,
+                                &mut cx,
+                                &mut cached_context,
+                                &mut jit_module,
+                                inst,
+                            )
                         });
                     }
-                    CodegenMode::JitLazy => codegen_shim(&mut cx, &mut jit_module, inst),
+                    CodegenMode::JitLazy => {
+                        codegen_shim(tcx, &mut cx, &mut cached_context, &mut jit_module, inst)
+                    }
                 },
                 MonoItem::Static(def_id) => {
                     crate::constant::codegen_static(tcx, &mut jit_module, def_id);
@@ -259,7 +267,15 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
                 false,
                 Symbol::intern("dummy_cgu_name"),
             );
-            tcx.sess.time("codegen fn", || crate::base::codegen_fn(&mut cx, jit_module, instance));
+            tcx.sess.time("codegen fn", || {
+                crate::base::codegen_and_compile_fn(
+                    tcx,
+                    &mut cx,
+                    &mut Context::new(),
+                    jit_module,
+                    instance,
+                )
+            });
 
             assert!(cx.global_asm.is_empty());
             jit_module.finalize_definitions();
@@ -269,10 +285,10 @@ fn jit_fn(instance_ptr: *const Instance<'static>, trampoline_ptr: *const u8) ->
     })
 }
 
-fn load_imported_symbols_for_jit(
+fn dep_symbol_lookup_fn(
     sess: &Session,
     crate_info: CrateInfo,
-) -> Vec<(String, *const u8)> {
+) -> Box<dyn Fn(&str) -> Option<*const u8>> {
     use rustc_middle::middle::dependency_format::Linkage;
 
     let mut dylib_paths = Vec::new();
@@ -299,44 +315,32 @@ fn load_imported_symbols_for_jit(
         }
     }
 
-    let mut imported_symbols = Vec::new();
-    for path in dylib_paths {
-        use object::{Object, ObjectSymbol};
-        let lib = libloading::Library::new(&path).unwrap();
-        let obj = std::fs::read(path).unwrap();
-        let obj = object::File::parse(&*obj).unwrap();
-        imported_symbols.extend(obj.dynamic_symbols().filter_map(|symbol| {
-            let name = symbol.name().unwrap().to_string();
-            if name.is_empty() || !symbol.is_global() || symbol.is_undefined() {
-                return None;
-            }
-            if name.starts_with("rust_metadata_") {
-                // The metadata is part of a section that is not loaded by the dynamic linker in
-                // case of cg_llvm.
-                return None;
-            }
-            let dlsym_name = if cfg!(target_os = "macos") {
-                // On macOS `dlsym` expects the name without leading `_`.
-                assert!(name.starts_with('_'), "{:?}", name);
-                &name[1..]
-            } else {
-                &name
-            };
-            let symbol: libloading::Symbol<'_, *const u8> =
-                unsafe { lib.get(dlsym_name.as_bytes()) }.unwrap();
-            Some((name, *symbol))
-        }));
-        std::mem::forget(lib)
-    }
+    let imported_dylibs = Box::leak(
+        dylib_paths
+            .into_iter()
+            .map(|path| unsafe { libloading::Library::new(&path).unwrap() })
+            .collect::<Box<[_]>>(),
+    );
 
     sess.abort_if_errors();
 
-    imported_symbols
+    Box::new(move |sym_name| {
+        for dylib in &*imported_dylibs {
+            if let Ok(sym) = unsafe { dylib.get::<*const u8>(sym_name.as_bytes()) } {
+                return Some(*sym);
+            }
+        }
+        None
+    })
 }
 
-fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: Instance<'tcx>) {
-    let tcx = cx.tcx;
-
+fn codegen_shim<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    cx: &mut CodegenCx,
+    cached_context: &mut Context,
+    module: &mut JITModule,
+    inst: Instance<'tcx>,
+) {
     let pointer_type = module.target_config().pointer_type();
 
     let name = tcx.symbol_name(inst).name;
@@ -357,8 +361,9 @@ fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: In
         )
         .unwrap();
 
-    cx.cached_context.clear();
-    let trampoline = &mut cx.cached_context.func;
+    let context = cached_context;
+    context.clear();
+    let trampoline = &mut context.func;
     trampoline.signature = sig.clone();
 
     let mut builder_ctx = FunctionBuilderContext::new();
@@ -381,5 +386,6 @@ fn codegen_shim<'tcx>(cx: &mut CodegenCx<'tcx>, module: &mut JITModule, inst: In
     let ret_vals = trampoline_builder.func.dfg.inst_results(call_inst).to_vec();
     trampoline_builder.ins().return_(&ret_vals);
 
-    module.define_function(func_id, &mut cx.cached_context).unwrap();
+    module.define_function(func_id, context).unwrap();
+    cx.unwind_context.add_function(func_id, context, module.isa());
 }