From a7de42f61b38dd69a4a58dcf2c7ee69bd4aed50f Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Sun, 28 Aug 2022 10:43:19 +0200
Subject: [PATCH] Use pull instead of push based model for getting dylib
 symbols in the jit

This avoids having to parse the dylibs to get all symbols and matches
the way the dynamic linker resolves symbols. Furthermore it fixes the
jit on Windows.
---
 build_system/tests.rs |  3 +--
 src/driver/jit.rs     | 51 +++++++++++++++----------------------------
 2 files changed, 18 insertions(+), 36 deletions(-)

diff --git a/build_system/tests.rs b/build_system/tests.rs
index e21397cece8..3f225b4efa2 100644
--- a/build_system/tests.rs
+++ b/build_system/tests.rs
@@ -465,8 +465,7 @@ pub fn new(host_triple: String, target_triple: String) -> Self {
         out_dir.push("out");
 
         let is_native = host_triple == target_triple;
-        let jit_supported =
-            target_triple.contains("x86_64") && is_native && !host_triple.contains("windows");
+        let jit_supported = target_triple.contains("x86_64") && is_native;
 
         let mut rust_flags = env::var("RUSTFLAGS").ok().unwrap_or("".to_string());
         let mut run_wrapper = Vec::new();
diff --git a/src/driver/jit.rs b/src/driver/jit.rs
index 0e77e4004c0..4ad859b6d83 100644
--- a/src/driver/jit.rs
+++ b/src/driver/jit.rs
@@ -67,13 +67,12 @@ fn create_jit_module(
     hotswap: bool,
 ) -> (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);
 
@@ -286,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();
@@ -316,39 +315,23 @@ 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| 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>(