patch unresolved callsites after jit compiled the callee

src/compiler/backend/arch/x86_64/asm_backend.rs

@@ -3238,6 +3238,7 @@ impl CodeGenerator for ASMCodeGen {
         &mut self,
         callsite: String,
         func: MuName,
+        func_id: MuID,
         addr: Address,
         pe: Option<MuName>,
         is_native: bool

src/compiler/backend/arch/x86_64/binary_backend/binary_codegen.rs

@@ -365,6 +365,7 @@ impl BinaryCodeGen {
         inst: XedInstruction,
         callsite_name: MuName,
         dest: Address,
+        func_id: MuID,
         potentially_exception: Option<MuName>
     ) {
         self.add_xed_inst_internal(
@@ -379,8 +380,9 @@ impl BinaryCodeGen {
             } else {
                 BranchTarget::None
             },
-            Some(MyCallsite {
+            Some(JITCallsite {
                 name: callsite_name,
+                callee_id: func_id,
                 target: CallTarget::Direct(dest)
             })
         )
@@ -395,7 +397,7 @@ impl BinaryCodeGen {
         is_move: bool,
         spill_info: Option<SpillMemInfo>,
         branch: BranchTarget,
-        call: Option<MyCallsite>
+        call: Option<JITCallsite>
     ) {
         let ctx = self.cur_mut();
         trace!("xed: {:?}", inst);
@@ -1030,6 +1032,7 @@ impl CodeGenerator for BinaryCodeGen {
         &mut self,
         callsite: String,
         func: MuName,
+        func_id: MuID,
         addr: Address,
         pe: Option<MuName>,
         is_native: bool
@@ -1046,7 +1049,7 @@ impl CodeGenerator for BinaryCodeGen {
             target
         );
 
-        self.add_xed_call_addr(inst, callsite, addr, pe);
+        self.add_xed_call_addr(inst, callsite, addr, func_id, pe);
 
         unsafe { ValueLocation::null() }
     }
@@ -1287,7 +1290,7 @@ pub fn emit_code(fv: &mut MuFunctionVersion, vm: &VM) {
 
         // generate code (back patching included)
         let code_start = mm::allocate_code(JIT_CODE_SIZE);
-        let code_size = xed_ctx.generate_code(code_start, JIT_CODE_SIZE);
+        let code_size = xed_ctx.generate_code(code_start, JIT_CODE_SIZE, vm);
 
         (code_start, code_start + code_size)
     };
@@ -1320,3 +1323,36 @@ pub fn emit_code(fv: &mut MuFunctionVersion, vm: &VM) {
 
     info!("code emission done");
 }
+
+pub fn jit_patch_call_addr(dest: Address, inst_start: Address, length: ByteSize) {
+    use std::i32;
+
+    let call_offset = dest.diff_offset(inst_start + length);
+
+    assert!(
+        call_offset >= i32::MIN as isize && call_offset <= i32::MAX as isize,
+        "call offset is not suitable for a CALL NEAR"
+    );
+
+    let inst = xed_inst1(
+        xed_state_64(),
+        xed_iclass_enum_t::XED_ICLASS_CALL_NEAR,
+        64,
+        xed_relbr(call_offset as i32, 32)
+    );
+
+    let new_len = encode_inst(
+        &inst,
+        inst_start.to_ptr_mut(),
+        XED_MAX_INSTRUCTION_BYTES as usize
+    );
+
+    if new_len < length {
+        warn!("need to pad with a few nops");
+        unimplemented!();
+    } else if new_len > length {
+        panic!("patching used more space than we have");
+    } else {
+        // do nothing
+    }
+}
\ No newline at end of file

src/compiler/backend/arch/x86_64/binary_backend/code_context.rs

@@ -13,6 +13,8 @@ use compiler::machine_code::MachineCode;
 use compiler::machine_code::CompiledCallsite;
 use runtime::ValueLocation;
 use runtime::mm;
+use runtime::patchpoint::RuntimePatchpoint;
+use vm::VM;
 use utils::mem::memsec;
 use utils::vec_utils;
 use utils::Address;
@@ -243,38 +245,57 @@ impl XedContext {
         ret
     }
 
-    pub fn generate_code(&mut self, code_ptr: Address, len: ByteSize) -> ByteSize {
+    pub fn generate_code(&mut self, code_ptr: Address, len: ByteSize, vm: &VM) -> ByteSize {
         info!("code mission");
 
         // generate code and recording branches and block start address for back patching
-        let mut cur_ptr = code_ptr;
-        let mut used = 0;
+        let code_length = {
+            let mut cur_ptr = code_ptr;
+            let mut used = 0;
 
-        for i in 0..self.insts.len() {
-            trace!("encode inst: {}", self.pretty_print_inst(i));
+            for i in 0..self.insts.len() {
+                trace!("encode inst: {}", self.pretty_print_inst(i));
 
-            let my_len = encode_inst(
-                &self.insts[i],
-                cur_ptr.to_ptr_mut(),
-                XED_MAX_INSTRUCTION_BYTES as usize
-            );
+                let my_len = encode_inst(
+                    &self.insts[i],
+                    cur_ptr.to_ptr_mut(),
+                    XED_MAX_INSTRUCTION_BYTES as usize
+                );
 
-            trace!("emitted");
-            trace!(
-                "0x{:x}: {}",
-                cur_ptr,
-                XedContext::inspect_code(cur_ptr, my_len)
-            );
-            trace!("");
+                trace!("emitted");
+                trace!(
+                    "0x{:x}: {}",
+                    cur_ptr,
+                    XedContext::inspect_code(cur_ptr, my_len)
+                );
+                trace!("");
 
-            self.insts_addr.push(cur_ptr);
-            self.insts_len.push(my_len as ByteSize);
-            cur_ptr = cur_ptr + my_len;
-            used += my_len;
-        }
+                self.insts_addr.push(cur_ptr);
+                self.insts_len.push(my_len as ByteSize);
+                cur_ptr = cur_ptr + my_len;
+                used += my_len;
+            }
 
-        trace!("generated code (before backpatching): ");
-        trace!("{}", XedContext::inspect_code(code_ptr, used));
+            trace!("generated code (before backpatching): ");
+            trace!("{}", XedContext::inspect_code(code_ptr, used));
+
+            used
+        };
+
+        // another pass through the code, record patchpoints
+        info!("record patch point...");
+        for i in 0..self.insts.len() {
+            if let Some(ref callsite) = self.insts_info[i].call {
+                let inst_addr = self.insts_addr[i];
+                let patchpoint = RuntimePatchpoint::inst(inst_addr, self.insts_len[i]);
+                trace!(
+                    "added callsite for func {} at {}",
+                    callsite.callee_id,
+                    inst_addr
+                );
+                vm.add_callsite_patchpoint(callsite.callee_id, patchpoint);
+            }
+        }
 
         info!("back patching...");
         for i in 0..self.insts.len() {
@@ -291,7 +312,7 @@ impl XedContext {
             // we may need to patch offsets
             let new_inst = match info.call {
                 // if we are calling an address, we need to patch offset
-                Some(MyCallsite {
+                Some(JITCallsite {
                     target: CallTarget::Direct(callee_addr),
                     ..
                 }) => {
@@ -353,10 +374,10 @@ impl XedContext {
         }
 
         trace!("generated code (after backpatching): ");
-        trace!("{}", XedContext::inspect_code(code_ptr, used));
+        trace!("{}", XedContext::inspect_code(code_ptr, code_length));
         trace!("back patching done");
 
-        used
+        code_length
     }
 
     /// builds a callsite table for runtime query
@@ -592,7 +613,7 @@ pub struct InstInfo {
     /// branch target of this instruction (intra-procedural)
     pub branch: BranchTarget,
     /// call target of this instruction (inter-procedural)
-    pub call: Option<MyCallsite>
+    pub call: Option<JITCallsite>
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -980,8 +1001,9 @@ pub enum BranchTarget {
 }
 
 #[derive(Clone, Debug)]
-pub struct MyCallsite {
+pub struct JITCallsite {
     pub name: MuName,
+    pub callee_id: MuID,
     pub target: CallTarget
 }
 

src/compiler/backend/arch/x86_64/binary_backend/mod.rs

@@ -19,6 +19,7 @@ mod binary_codegen;
 pub use self::binary_codegen::BinaryCodeGen;
 pub use self::binary_codegen::init_jit;
 pub use self::binary_codegen::emit_code;
+pub use self::binary_codegen::jit_patch_call_addr;
 
 mod code_context;
 

src/compiler/backend/arch/x86_64/codegen.rs

@@ -226,6 +226,7 @@ pub trait CodeGenerator {
         &mut self,
         callsite: String,
         func: MuName,
+        func_id: MuID,
         addr: Address,
         pe: Option<MuName>,
         is_native: bool

src/compiler/backend/arch/x86_64/inst_sel.rs

@@ -3730,6 +3730,7 @@ impl<'a> InstructionSelection {
             self.backend.emit_call_near_rel32(
                 callsite.clone(),
                 func_name,
+                0,
                 unsafe { Address::zero() },
                 None,
                 true
@@ -3906,6 +3907,7 @@ impl<'a> InstructionSelection {
                     self.backend.emit_call_near_rel32(
                         callsite,
                         target.name(),
+                        target_id,
                         unsafe { Address::zero() },
                         potentially_excepting,
                         false
@@ -3917,6 +3919,7 @@ impl<'a> InstructionSelection {
                             self.backend.emit_call_near_rel32(
                                 callsite,
                                 target.name(),
+                                target_id,
                                 addr,
                                 potentially_excepting,
                                 false
@@ -3927,6 +3930,7 @@ impl<'a> InstructionSelection {
                             self.backend.emit_call_near_rel32(
                                 callsite,
                                 target.name(),
+                                target_id,
                                 entrypoints::LAZY_RESOLVE_FUNC.jit,
                                 potentially_excepting,
                                 false

src/compiler/backend/mod.rs

@@ -130,6 +130,9 @@ pub use compiler::backend::x86_64::asm_backend::spill_rewrite as asm_spill_rewri
 #[cfg(feature = "binary-codegen")]
 pub use compiler::backend::x86_64::binary_backend::spill_rewrite as binary_spill_rewrite;
 #[cfg(target_arch = "x86_64")]
+#[cfg(feature = "binary-codegen")]
+pub use compiler::backend::x86_64::binary_backend::jit_patch_call_addr;
+#[cfg(target_arch = "x86_64")]
 pub use compiler::backend::x86_64::ARGUMENT_GPRS;
 #[cfg(target_arch = "x86_64")]
 pub use compiler::backend::x86_64::ARGUMENT_FPRS;

src/compiler/mod.rs

@@ -183,7 +183,20 @@ pub extern "C" fn lazy_resolve_function_internal(callsite: Address) -> Address {
     };
 
     trace!("lazy compiling function version {}", func_id);
-    vm.jit_function(func_id)
+    let func_start = vm.jit_function(func_id);
+
+    // patch all callsites of this function
+    {
+        use runtime::patchpoint::*;
+        let patch_table = vm.callsite_patch_table().read().unwrap();
+        for patchpoint in patch_table.get(&func_id).unwrap().iter() {
+            assert!(patchpoint.v == RuntimePatchpoint_::Instruction);
+
+            backend::jit_patch_call_addr(func_start, patchpoint.start, patchpoint.length);
+        }
+    }
+
+    func_start
 }
 
 // rewrite parts of the hprof crates to print via log (instead of print!())

src/runtime/mod.rs

@@ -39,6 +39,8 @@ pub mod math;
 pub mod entrypoints;
 /// exception handling
 pub mod exception;
+/// code patching point
+pub mod patchpoint;
 
 
 /// returns name for a function address

src/runtime/patchpoint.rs

+use utils::Address;
+use utils::ByteSize;
+
+#[derive(Clone, Copy, Debug)]
+pub struct RuntimePatchpoint {
+    pub start: Address,
+    pub length: ByteSize,
+    pub v: RuntimePatchpoint_
+}
+
+impl RuntimePatchpoint {
+    pub fn inst(start: Address, length: ByteSize) -> RuntimePatchpoint {
+        RuntimePatchpoint {
+            start: start,
+            length: length,
+            v: RuntimePatchpoint_::Instruction
+        }
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq)]
+pub enum RuntimePatchpoint_ {
+    Instruction,
+    Offset(ByteSize)
+}
\ No newline at end of file

src/vm/vm.rs

@@ -25,12 +25,13 @@ use compiler::backend;
 use compiler::backend::BackendType;
 use compiler::machine_code::{CompiledCallsite, CompiledFunction};
 
-use runtime::thread::*;
-use runtime::*;
 use utils::ByteSize;
 use utils::BitSize;
 use utils::Address;
+use runtime::thread::*;
+use runtime::*;
 use runtime::mm as gc;
+use runtime::patchpoint::RuntimePatchpoint;
 use vm::handle::*;
 use vm::vm_options::VMOptions;
 use vm::vm_options::MuLogLevel;
@@ -119,6 +120,13 @@ pub struct VM {
     /// a map from callsite address to CompiledCallsite
     compiled_callsite_table: RwLock<HashMap<Address, CompiledCallsite>>, // 896
 
+    /// patchpoints for callsites. This is used for two scenarios:
+    /// * for unresolved functions, we insert a call to trampoline, but later after we JIT compiled
+    ///   the callee, we need to patch the callsite
+    /// * for function redefinition, we need to patch the callsite as well
+    /// We keep this as a separate table because we want to index it with function ID
+    callsite_patchpoints: RwLock<HashMap<MuID, Vec<RuntimePatchpoint>>>,
+
     /// Nnmber of callsites in the callsite tables
     callsite_count: AtomicUsize
 }
@@ -162,6 +170,12 @@ unsafe impl rodal::Dump for VM {
         dumper.dump_object_here(&RwLock::new(
             rodal::EmptyHashMap::<Address, CompiledCallsite>::new()
         ));
+
+        dumper.dump_padding(&self.callsite_patchpoints);
+        dumper.dump_object_here(&RwLock::new(
+            rodal::EmptyHashMap::<MuID, Vec<RuntimePatchpoint>>::new()
+        ));
+
         dumper.dump_object(&self.callsite_count);
     }
 }
@@ -225,6 +239,7 @@ impl<'a> VM {
             primordial: RwLock::new(None),
             aot_pending_funcref_store: RwLock::new(HashMap::new()),
             compiled_callsite_table: RwLock::new(HashMap::new()),
+            callsite_patchpoints: RwLock::new(HashMap::new()),
             callsite_count: ATOMIC_USIZE_INIT
         };
 
@@ -341,6 +356,17 @@ impl<'a> VM {
         self.callsite_count.fetch_add(1, Ordering::Relaxed);
     }
 
+    /// adds a callsite patchpoint
+    pub fn add_callsite_patchpoint(&self, function_id: MuID, patchpoint: RuntimePatchpoint) {
+        let mut table = self.callsite_patchpoints.write().unwrap();
+
+        if table.contains_key(&function_id) {
+            table.get_mut(&function_id).unwrap().push(patchpoint);
+        } else {
+            table.insert(function_id, vec![patchpoint]);
+        }
+    }
+
     /// resumes persisted VM. Ideally the VM should be back to the status when we start
     /// persisting it except a few fields that we do not want to persist.
     pub fn resume_vm(dumped_vm: *mut Arc<VM>) -> Arc<VM> {
@@ -953,6 +979,11 @@ impl<'a> VM {
         &self.callsite_table
     }
 
+    /// returns the lock for callsite patch table
+    pub fn callsite_patch_table(&self) -> &RwLock<HashMap<MuID, Vec<RuntimePatchpoint>>> {
+        &self.callsite_patchpoints
+    }
+
     /// set info (entry function, arguments) for primordial thread for boot image
     pub fn set_primordial_thread(&self, func_id: MuID, has_const_args: bool, args: Vec<Constant>) {
         let mut guard = self.primordial.write().unwrap();

tests/test_jit_rust/test_jit_call.rs

@@ -70,6 +70,42 @@ fn test_call_add_unresolved() {
     }
 }
 
+#[test]
+fn test_call_add_callsite_patching() {
+    VM::start_logging_trace();
+
+    let mut vm = add();
+    call_add(&vm);
+    // do not inline functions otherwise add() will get inlined
+    vm.vm_options.flag_disable_inline = true;
+
+    let vm = Arc::new(vm);
+
+    // compile call_add() - at this time, add() is NOT compiled (we will lazily compile add())
+    let call_add_id = vm.id_of("call_add");
+    let code_start = vm.jit_function(call_add_id);
+    info!("call_add() is at 0x{:x}", code_start);
+    info!("add() fvid = {}", vm.id_of("add"));
+
+    unsafe {
+        let need_dealloc = MuThread::current_thread_as_mu_thread(Address::zero(), vm.clone());
+
+        let call_add: unsafe extern "C" fn(u64, u64) -> u64 = mem::transmute(code_start);
+
+        let res = call_add(1, 1);
+        println!("add(1, 1) = {}", res);
+        assert!(res == 2);
+
+        let res = call_add(1, 2);
+        println!("add(1, 2) = {}", res);
+        assert!(res == 3);
+
+        if need_dealloc {
+            MuThread::cleanup_current_mu_thread();
+        }
+    }
+}
+
 fn call_add(vm: &VM) {
     let int64 = vm.get_type(vm.id_of("int64"));