From 2ca01ae9478f199337cf743a6ab543e8c3f3b238 Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Sat, 19 Feb 2022 17:33:39 +0100
Subject: [PATCH] Add a way to define a symbol lookup fn for the JIT (#2745)

* Couple of cranelift-jit cleanups

* Add a way to define a symbol lookup fn for the JIT
---
 cranelift/jit/src/backend.rs | 53 ++++++++++++++++++++++++++----------
 1 file changed, 38 insertions(+), 15 deletions(-)

diff --git a/cranelift/jit/src/backend.rs b/cranelift/jit/src/backend.rs
index 82a38ec0d1..3e0cc7368b 100644
--- a/cranelift/jit/src/backend.rs
+++ b/cranelift/jit/src/backend.rs
@@ -14,6 +14,7 @@ use cranelift_module::{
     ModuleDeclarations, ModuleError, ModuleResult,
 };
 use log::info;
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::convert::{TryFrom, TryInto};
 use std::ffi::CString;
@@ -31,6 +32,7 @@ const READONLY_DATA_ALIGNMENT: u64 = 0x1;
 pub struct JITBuilder {
     isa: Box<dyn TargetIsa>,
     symbols: HashMap<String, *const u8>,
+    lookup_symbols: Vec<Box<dyn Fn(&str) -> Option<*const u8>>>,
     libcall_names: Box<dyn Fn(ir::LibCall) -> String + Send + Sync>,
     hotswap_enabled: bool,
 }
@@ -73,9 +75,11 @@ impl JITBuilder {
         libcall_names: Box<dyn Fn(ir::LibCall) -> String + Send + Sync>,
     ) -> Self {
         let symbols = HashMap::new();
+        let lookup_symbols = vec![Box::new(lookup_with_dlsym) as Box<_>];
         Self {
             isa,
             symbols,
+            lookup_symbols,
             libcall_names,
             hotswap_enabled: false,
         }
@@ -95,7 +99,7 @@ impl JITBuilder {
     /// back to a platform-specific search (this typically involves searching
     /// the current process for public symbols, followed by searching the
     /// platform's C runtime).
-    pub fn symbol<K>(&mut self, name: K, ptr: *const u8) -> &Self
+    pub fn symbol<K>(&mut self, name: K, ptr: *const u8) -> &mut Self
     where
         K: Into<String>,
     {
@@ -106,7 +110,7 @@ impl JITBuilder {
     /// Define multiple symbols in the internal symbol table.
     ///
     /// Using this is equivalent to calling `symbol` on each element.
-    pub fn symbols<It, K>(&mut self, symbols: It) -> &Self
+    pub fn symbols<It, K>(&mut self, symbols: It) -> &mut Self
     where
         It: IntoIterator<Item = (K, *const u8)>,
         K: Into<String>,
@@ -117,6 +121,18 @@ impl JITBuilder {
         self
     }
 
+    /// Add a symbol lookup fn.
+    ///
+    /// Symbol lookup fn's are used to lookup symbols when they couldn't be found in the internal
+    /// symbol table. Symbol lookup fn's are called in reverse of the order in which they were added.
+    pub fn symbol_lookup_fn(
+        &mut self,
+        symbol_lookup_fn: Box<dyn Fn(&str) -> Option<*const u8>>,
+    ) -> &mut Self {
+        self.lookup_symbols.push(symbol_lookup_fn);
+        self
+    }
+
     /// Enable or disable hotswap support. See [`JITModule::prepare_for_function_redefine`]
     /// for more information.
     ///
@@ -143,7 +159,8 @@ struct GotUpdate {
 pub struct JITModule {
     isa: Box<dyn TargetIsa>,
     hotswap_enabled: bool,
-    symbols: HashMap<String, *const u8>,
+    symbols: RefCell<HashMap<String, *const u8>>,
+    lookup_symbols: Vec<Box<dyn Fn(&str) -> Option<*const u8>>>,
     libcall_names: Box<dyn Fn(ir::LibCall) -> String>,
     memory: MemoryHandle,
     declarations: ModuleDeclarations,
@@ -184,10 +201,20 @@ impl JITModule {
     }
 
     fn lookup_symbol(&self, name: &str) -> Option<*const u8> {
-        self.symbols
-            .get(name)
-            .copied()
-            .or_else(|| lookup_with_dlsym(name))
+        match self.symbols.borrow_mut().entry(name.to_owned()) {
+            std::collections::hash_map::Entry::Occupied(occ) => Some(*occ.get()),
+            std::collections::hash_map::Entry::Vacant(vac) => {
+                let ptr = self
+                    .lookup_symbols
+                    .iter()
+                    .rev() // Try last lookup function first
+                    .find_map(|lookup| lookup(name));
+                if let Some(ptr) = ptr {
+                    vac.insert(ptr);
+                }
+                ptr
+            }
+        }
     }
 
     fn new_got_entry(&mut self, val: *const u8) -> NonNull<AtomicPtr<u8>> {
@@ -277,7 +304,7 @@ impl JITModule {
                         }
                     }
                 };
-                if let Some(ptr) = self.lookup_symbol(&name) {
+                if let Some(ptr) = self.lookup_symbol(name) {
                     ptr
                 } else if linkage == Linkage::Preemptible {
                     0 as *const u8
@@ -455,7 +482,8 @@ impl JITModule {
         let mut module = Self {
             isa: builder.isa,
             hotswap_enabled: builder.hotswap_enabled,
-            symbols: builder.symbols,
+            symbols: RefCell::new(builder.symbols),
+            lookup_symbols: builder.lookup_symbols,
             libcall_names: builder.libcall_names,
             memory: MemoryHandle {
                 code: Memory::new(),
@@ -483,12 +511,7 @@ impl JITModule {
         };
         for &libcall in all_libcalls {
             let sym = (module.libcall_names)(libcall);
-            let addr = if let Some(addr) = module
-                .symbols
-                .get(&sym)
-                .copied()
-                .or_else(|| lookup_with_dlsym(&sym))
-            {
+            let addr = if let Some(addr) = module.lookup_symbol(&sym) {
                 addr
             } else {
                 continue;