Implement JIT-to-JIT calls (https://github.com/Shopify/zjit/pull/109)

* Implement JIT-to-JIT calls

* Use a closer dummy address for Arm64

* Revert an obsoleted change

* Revert a few more obsoleted changes

* Fix outdated comments

* Explain PosMarkers for CCall

* s/JIT code/machine code/

* Get rid of ParallelMov

test/ruby/test_zjit.rb

@@ -389,6 +389,21 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_method_call
+    assert_compiles '12', %q{
+      def callee(a, b)
+        a - b
+      end
+
+      def test
+        callee(4, 2) + 10
+      end
+
+      test # profile test
+      test
+    }, call_threshold: 2
+  end
+
   def test_recursive_fact
     assert_compiles '[1, 6, 720]', %q{
       def fact(n)
@@ -401,6 +416,19 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_profiled_fact
+    assert_compiles '[1, 6, 720]', %q{
+      def fact(n)
+        if n == 0
+          return 1
+        end
+        return n * fact(n-1)
+      end
+      fact(1) # profile fact
+      [fact(0), fact(3), fact(6)]
+    }, call_threshold: 3, num_profiles: 2
+  end
+
   def test_recursive_fib
     assert_compiles '[0, 2, 3]', %q{
       def fib(n)
@@ -413,11 +441,24 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_profiled_fib
+    assert_compiles '[0, 2, 3]', %q{
+      def fib(n)
+        if n < 2
+          return n
+        end
+        return fib(n-1) + fib(n-2)
+      end
+      fib(3) # profile fib
+      [fib(0), fib(3), fib(4)]
+    }, call_threshold: 5, num_profiles: 3
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT
   # at a given call_threshold
-  def assert_compiles(expected, test_script, call_threshold: 1)
+  def assert_compiles(expected, test_script, **opts)
     pipe_fd = 3
 
     script = <<~RUBY
@@ -429,7 +470,7 @@ class TestZJIT < Test::Unit::TestCase
       IO.open(#{pipe_fd}).write(result.inspect)
     RUBY
 
-    status, out, err, actual = eval_with_jit(script, call_threshold:, pipe_fd:)
+    status, out, err, actual = eval_with_jit(script, pipe_fd:, **opts)
 
     message = "exited with status #{status.to_i}"
     message << "\nstdout:\n```\n#{out}```\n" unless out.empty?
@@ -440,10 +481,11 @@ class TestZJIT < Test::Unit::TestCase
   end
 
   # Run a Ruby process with ZJIT options and a pipe for writing test results
-  def eval_with_jit(script, call_threshold: 1, timeout: 1000, pipe_fd:, debug: true)
+  def eval_with_jit(script, call_threshold: 1, num_profiles: 1, timeout: 1000, pipe_fd:, debug: true)
     args = [
       "--disable-gems",
       "--zjit-call-threshold=#{call_threshold}",
+      "--zjit-num-profiles=#{num_profiles}",
     ]
     args << "--zjit-debug" if debug
     args << "-e" << script_shell_encode(script)

zjit/src/asm/mod.rs

@@ -111,6 +111,28 @@ impl CodeBlock {
         self.get_ptr(self.write_pos)
     }
 
+    /// Set the current write position from a pointer
+    fn set_write_ptr(&mut self, code_ptr: CodePtr) {
+        let pos = code_ptr.as_offset() - self.mem_block.borrow().start_ptr().as_offset();
+        self.write_pos = pos.try_into().unwrap();
+    }
+
+    /// Invoke a callback with write_ptr temporarily adjusted to a given address
+    pub fn with_write_ptr(&mut self, code_ptr: CodePtr, callback: impl Fn(&mut CodeBlock)) {
+        // Temporarily update the write_pos. Ignore the dropped_bytes flag at the old address.
+        let old_write_pos = self.write_pos;
+        let old_dropped_bytes = self.dropped_bytes;
+        self.set_write_ptr(code_ptr);
+        self.dropped_bytes = false;
+
+        // Invoke the callback
+        callback(self);
+
+        // Restore the original write_pos and dropped_bytes flag.
+        self.dropped_bytes = old_dropped_bytes;
+        self.write_pos = old_write_pos;
+    }
+
     /// Get a (possibly dangling) direct pointer into the executable memory block
     pub fn get_ptr(&self, offset: usize) -> CodePtr {
         self.mem_block.borrow().start_ptr().add_bytes(offset)

zjit/src/backend/arm64/mod.rs

@@ -491,19 +491,21 @@ impl Assembler
                     // register.
                     // Note: the iteration order is reversed to avoid corrupting x0,
                     // which is both the return value and first argument register
-                    let mut args: Vec<(Reg, Opnd)> = vec![];
-                    for (idx, opnd) in opnds.into_iter().enumerate().rev() {
-                        // If the value that we're sending is 0, then we can use
-                        // the zero register, so in this case we'll just send
-                        // a UImm of 0 along as the argument to the move.
-                        let value = match opnd {
-                            Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
-                            Opnd::Mem(_) => split_memory_address(asm, *opnd),
-                            _ => *opnd
-                        };
-                        args.push((C_ARG_OPNDS[idx].unwrap_reg(), value));
+                    if !opnds.is_empty() {
+                        let mut args: Vec<(Reg, Opnd)> = vec![];
+                        for (idx, opnd) in opnds.into_iter().enumerate().rev() {
+                            // If the value that we're sending is 0, then we can use
+                            // the zero register, so in this case we'll just send
+                            // a UImm of 0 along as the argument to the move.
+                            let value = match opnd {
+                                Opnd::UImm(0) | Opnd::Imm(0) => Opnd::UImm(0),
+                                Opnd::Mem(_) => split_memory_address(asm, *opnd),
+                                _ => *opnd
+                            };
+                            args.push((C_ARG_OPNDS[idx].unwrap_reg(), value));
+                        }
+                        asm.parallel_mov(args);
                     }
-                    asm.parallel_mov(args);
 
                     // Now we push the CCall without any arguments so that it
                     // just performs the call.

zjit/src/backend/lir.rs

@@ -345,7 +345,19 @@ pub enum Insn {
     CPushAll,
 
     // C function call with N arguments (variadic)
-    CCall { opnds: Vec<Opnd>, fptr: *const u8, out: Opnd },
+    CCall {
+        opnds: Vec<Opnd>,
+        fptr: *const u8,
+        /// Optional PosMarker to remember the start address of the C call.
+        /// It's embedded here to insert the PosMarker after push instructions
+        /// that are split from this CCall on alloc_regs().
+        start_marker: Option<PosMarkerFn>,
+        /// Optional PosMarker to remember the end address of the C call.
+        /// It's embedded here to insert the PosMarker before pop instructions
+        /// that are split from this CCall on alloc_regs().
+        end_marker: Option<PosMarkerFn>,
+        out: Opnd,
+    },
 
     // C function return
     CRet(Opnd),
@@ -1455,6 +1467,23 @@ impl Assembler
         let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), live_ranges.len());
 
         while let Some((index, mut insn)) = iterator.next() {
+            let before_ccall = match (&insn, iterator.peek().map(|(_, insn)| insn)) {
+                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
+                (Insn::CCall { .. }, _) if !pool.is_empty() => {
+                    // If C_RET_REG is in use, move it to another register.
+                    // This must happen before last-use registers are deallocated.
+                    if let Some(vreg_idx) = pool.vreg_for(&C_RET_REG) {
+                        let new_reg = pool.alloc_reg(vreg_idx).unwrap(); // TODO: support spill
+                        asm.mov(Opnd::Reg(new_reg), C_RET_OPND);
+                        pool.dealloc_reg(&C_RET_REG);
+                        reg_mapping[vreg_idx] = Some(new_reg);
+                    }
+
+                    true
+                },
+                _ => false,
+            };
+
             // Check if this is the last instruction that uses an operand that
             // spans more than one instruction. In that case, return the
             // allocated register to the pool.
@@ -1477,32 +1506,20 @@ impl Assembler
                 }
             }
 
-            // If we're about to make a C call, save caller-saved registers
-            match (&insn, iterator.peek().map(|(_, insn)| insn)) {
-                (Insn::ParallelMov { .. }, Some(Insn::CCall { .. })) |
-                (Insn::CCall { .. }, _) if !pool.is_empty() => {
-                    // If C_RET_REG is in use, move it to another register
-                    if let Some(vreg_idx) = pool.vreg_for(&C_RET_REG) {
-                        let new_reg = pool.alloc_reg(vreg_idx).unwrap(); // TODO: support spill
-                        asm.mov(Opnd::Reg(new_reg), C_RET_OPND);
-                        pool.dealloc_reg(&C_RET_REG);
-                        reg_mapping[vreg_idx] = Some(new_reg);
-                    }
+            // Save caller-saved registers on a C call.
+            if before_ccall {
+                // Find all live registers
+                saved_regs = pool.live_regs();
 
-                    // Find all live registers
-                    saved_regs = pool.live_regs();
-
-                    // Save live registers
-                    for &(reg, _) in saved_regs.iter() {
-                        asm.cpush(Opnd::Reg(reg));
-                        pool.dealloc_reg(&reg);
-                    }
-                    // On x86_64, maintain 16-byte stack alignment
-                    if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
-                        asm.cpush(Opnd::Reg(saved_regs.last().unwrap().0));
-                    }
+                // Save live registers
+                for &(reg, _) in saved_regs.iter() {
+                    asm.cpush(Opnd::Reg(reg));
+                    pool.dealloc_reg(&reg);
+                }
+                // On x86_64, maintain 16-byte stack alignment
+                if cfg!(target_arch = "x86_64") && saved_regs.len() % 2 == 1 {
+                    asm.cpush(Opnd::Reg(saved_regs.last().unwrap().0));
                 }
-                _ => {},
             }
 
             // If the output VReg of this instruction is used by another instruction,
@@ -1590,13 +1607,24 @@ impl Assembler
 
             // Push instruction(s)
             let is_ccall = matches!(insn, Insn::CCall { .. });
-            if let Insn::ParallelMov { moves } = insn {
-                // Now that register allocation is done, it's ready to resolve parallel moves.
-                for (reg, opnd) in Self::resolve_parallel_moves(&moves) {
-                    asm.load_into(Opnd::Reg(reg), opnd);
+            match insn {
+                Insn::ParallelMov { moves } => {
+                    // Now that register allocation is done, it's ready to resolve parallel moves.
+                    for (reg, opnd) in Self::resolve_parallel_moves(&moves) {
+                        asm.load_into(Opnd::Reg(reg), opnd);
+                    }
                 }
-            } else {
-                asm.push_insn(insn);
+                Insn::CCall { opnds, fptr, start_marker, end_marker, out } => {
+                    // Split start_marker and end_marker here to avoid inserting push/pop between them.
+                    if let Some(start_marker) = start_marker {
+                        asm.push_insn(Insn::PosMarker(start_marker));
+                    }
+                    asm.push_insn(Insn::CCall { opnds, fptr, start_marker: None, end_marker: None, out });
+                    if let Some(end_marker) = end_marker {
+                        asm.push_insn(Insn::PosMarker(end_marker));
+                    }
+                }
+                _ => asm.push_insn(insn),
             }
 
             // After a C call, restore caller-saved registers
@@ -1720,38 +1748,30 @@ impl Assembler {
         self.push_insn(Insn::Breakpoint);
     }
 
+    /// Call a C function without PosMarkers
     pub fn ccall(&mut self, fptr: *const u8, opnds: Vec<Opnd>) -> Opnd {
-        /*
-        // Let vm_check_canary() assert this ccall's leafness if leaf_ccall is set
-        let canary_opnd = self.set_stack_canary(&opnds);
-
-        let old_temps = self.ctx.get_reg_mapping(); // with registers
-        // Spill stack temp registers since they are caller-saved registers.
-        // Note that this doesn't spill stack temps that are already popped
-        // but may still be used in the C arguments.
-        self.spill_regs();
-        let new_temps = self.ctx.get_reg_mapping(); // all spilled
-
-        // Temporarily manipulate RegMappings so that we can use registers
-        // to pass stack operands that are already spilled above.
-        self.ctx.set_reg_mapping(old_temps);
-        */
-
-        // Call a C function
         let out = self.new_vreg(Opnd::match_num_bits(&opnds));
-        self.push_insn(Insn::CCall { fptr, opnds, out });
-
-        /*
-        // Registers in old_temps may be clobbered by the above C call,
-        // so rollback the manipulated RegMappings to a spilled version.
-        self.ctx.set_reg_mapping(new_temps);
-
-        // Clear the canary after use
-        if let Some(canary_opnd) = canary_opnd {
-            self.mov(canary_opnd, 0.into());
-        }
-        */
+        self.push_insn(Insn::CCall { fptr, opnds, start_marker: None, end_marker: None, out });
+        out
+    }
 
+    /// Call a C function with PosMarkers. This is used for recording the start and end
+    /// addresses of the C call and rewriting it with a different function address later.
+    pub fn ccall_with_pos_markers(
+        &mut self,
+        fptr: *const u8,
+        opnds: Vec<Opnd>,
+        start_marker: impl Fn(CodePtr, &CodeBlock) + 'static,
+        end_marker: impl Fn(CodePtr, &CodeBlock) + 'static,
+    ) -> Opnd {
+        let out = self.new_vreg(Opnd::match_num_bits(&opnds));
+        self.push_insn(Insn::CCall {
+            fptr,
+            opnds,
+            start_marker: Some(Box::new(start_marker)),
+            end_marker: Some(Box::new(end_marker)),
+            out,
+        });
         out
     }
 

zjit/src/backend/x86_64/mod.rs

@@ -82,8 +82,8 @@ impl From<&Opnd> for X86Opnd {
 /// This has the same number of registers for x86_64 and arm64.
 /// SCRATCH_REG is excluded.
 pub const ALLOC_REGS: &'static [Reg] = &[
-    RSI_REG,
     RDI_REG,
+    RSI_REG,
     RDX_REG,
     RCX_REG,
     R8_REG,
@@ -338,11 +338,13 @@ impl Assembler
                     assert!(opnds.len() <= C_ARG_OPNDS.len());
 
                     // Load each operand into the corresponding argument register.
-                    let mut args: Vec<(Reg, Opnd)> = vec![];
-                    for (idx, opnd) in opnds.into_iter().enumerate() {
-                        args.push((C_ARG_OPNDS[idx].unwrap_reg(), *opnd));
+                    if !opnds.is_empty() {
+                        let mut args: Vec<(Reg, Opnd)> = vec![];
+                        for (idx, opnd) in opnds.into_iter().enumerate() {
+                            args.push((C_ARG_OPNDS[idx].unwrap_reg(), *opnd));
+                        }
+                        asm.parallel_mov(args);
                     }
-                    asm.parallel_mov(args);
 
                     // Now we push the CCall without any arguments so that it
                     // just performs the call.

zjit/src/codegen.rs

@@ -1,4 +1,8 @@
+use std::cell::Cell;
+use std::rc::Rc;
+
 use crate::backend::current::{Reg, ALLOC_REGS};
+use crate::profile::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
@@ -17,6 +21,9 @@ struct JITState {
 
     /// Labels for each basic block indexed by the BlockId
     labels: Vec<Option<Target>>,
+
+    /// Branches to an ISEQ that need to be compiled later
+    branch_iseqs: Vec<(Rc<Branch>, IseqPtr)>,
 }
 
 impl JITState {
@@ -26,6 +33,7 @@ impl JITState {
             iseq,
             opnds: vec![None; num_insns],
             labels: vec![None; num_blocks],
+            branch_iseqs: Vec::default(),
         }
     }
 
@@ -83,33 +91,47 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
     code_ptr
 }
 
-
 /// Compile an entry point for a given ISEQ
 fn gen_iseq_entry_point(iseq: IseqPtr) -> *const u8 {
     // Compile ISEQ into High-level IR
-    let mut function = match iseq_to_hir(iseq) {
-        Ok(function) => function,
-        Err(err) => {
-            debug!("ZJIT: iseq_to_hir: {err:?}");
-            return std::ptr::null();
-        }
+    let function = match compile_iseq(iseq) {
+        Some(function) => function,
+        None => return std::ptr::null(),
     };
-    function.optimize();
 
     // Compile the High-level IR
     let cb = ZJITState::get_code_block();
-    let function_ptr = gen_function(cb, iseq, &function);
-    // TODO: Reuse function_ptr for JIT-to-JIT calls
+    let (start_ptr, mut branch_iseqs) = match gen_function(cb, iseq, &function) {
+        Some((start_ptr, branch_iseqs)) => {
+            // Remember the block address to reuse it later
+            let payload = get_or_create_iseq_payload(iseq);
+            payload.start_ptr = Some(start_ptr);
 
-    // Compile an entry point to the JIT code
-    let start_ptr = match function_ptr {
-        Some(function_ptr) => gen_entry(cb, iseq, &function, function_ptr),
-        None => None,
+            // Compile an entry point to the JIT code
+            (gen_entry(cb, iseq, &function, start_ptr), branch_iseqs)
+        },
+        None => (None, vec![]),
     };
 
+    // Recursively compile callee ISEQs
+    while let Some((branch, iseq)) = branch_iseqs.pop() {
+        // Disable profiling. This will be the last use of the profiling information for the ISEQ.
+        unsafe { rb_zjit_profile_disable(iseq); }
+
+        // Compile the ISEQ
+        if let Some((callee_ptr, callee_branch_iseqs)) = gen_iseq(cb, iseq) {
+            let callee_addr = callee_ptr.raw_ptr(cb);
+            branch.regenerate(cb, |asm| {
+                asm.ccall(callee_addr, vec![]);
+            });
+            branch_iseqs.extend(callee_branch_iseqs);
+        }
+    }
+
     // Always mark the code region executable if asm.compile() has been used
     cb.mark_all_executable();
 
+    // Return a JIT code address or a null pointer
     start_ptr.map(|start_ptr| start_ptr.raw_ptr(cb)).unwrap_or(std::ptr::null())
 }
 
@@ -117,18 +139,52 @@ fn gen_iseq_entry_point(iseq: IseqPtr) -> *const u8 {
 fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_ptr: CodePtr) -> Option<CodePtr> {
     // Set up registers for CFP, EC, SP, and basic block arguments
     let mut asm = Assembler::new();
-    gen_entry_prologue(iseq, &mut asm);
+    gen_entry_prologue(&mut asm, iseq);
     gen_method_params(&mut asm, iseq, function.block(BlockId(0)));
 
-    // Jump to the function. We can't remove this jump by calling gen_entry() first and
-    // then calling gen_function() because gen_function() writes side exit code first.
-    asm.jmp(function_ptr.into());
+    // Jump to the first block using a call instruction
+    asm.ccall(function_ptr.raw_ptr(cb) as *const u8, vec![]);
+
+    // Restore registers for CFP, EC, and SP after use
+    asm_comment!(asm, "exit to the interpreter");
+    // On x86_64, maintain 16-byte stack alignment
+    if cfg!(target_arch = "x86_64") {
+        asm.cpop_into(SP);
+    }
+    asm.cpop_into(SP);
+    asm.cpop_into(EC);
+    asm.cpop_into(CFP);
+    asm.frame_teardown();
+    asm.cret(C_RET_OPND);
 
     asm.compile(cb).map(|(start_ptr, _)| start_ptr)
 }
 
+/// Compile an ISEQ into machine code
+fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<(Rc<Branch>, IseqPtr)>)> {
+    // Return an existing pointer if it's already compiled
+    let payload = get_or_create_iseq_payload(iseq);
+    if let Some(start_ptr) = payload.start_ptr {
+        return Some((start_ptr, vec![]));
+    }
+
+    // Convert ISEQ into High-level IR
+    let mut function = match compile_iseq(iseq) {
+        Some(function) => function,
+        None => return None,
+    };
+    function.optimize();
+
+    // Compile the High-level IR
+    let result = gen_function(cb, iseq, &function);
+    if let Some((start_ptr, _)) = result {
+        payload.start_ptr = Some(start_ptr);
+    }
+    result
+}
+
 /// Compile a function
-fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Option<CodePtr> {
+fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Option<(CodePtr, Vec<(Rc<Branch>, IseqPtr)>)> {
     let mut jit = JITState::new(iseq, function.num_insns(), function.num_blocks());
     let mut asm = Assembler::new();
 
@@ -142,6 +198,11 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
         let label = jit.get_label(&mut asm, block_id);
         asm.write_label(label);
 
+        // Set up the frame at the first block
+        if block_id == BlockId(0) {
+            asm.frame_setup();
+        }
+
         // Compile all parameters
         for &insn_id in block.params() {
             match function.find(insn_id) {
@@ -155,7 +216,7 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
         // Compile all instructions
         for &insn_id in block.insns() {
             let insn = function.find(insn_id);
-            if gen_insn(&mut jit, &mut asm, function, insn_id, &insn).is_none() {
+            if gen_insn(cb, &mut jit, &mut asm, function, insn_id, &insn).is_none() {
                 debug!("Failed to compile insn: {insn_id} {insn:?}");
                 return None;
             }
@@ -163,11 +224,11 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
     }
 
     // Generate code if everything can be compiled
-    asm.compile(cb).map(|(start_ptr, _)| start_ptr)
+    asm.compile(cb).map(|(start_ptr, _)| (start_ptr, jit.branch_iseqs))
 }
 
 /// Compile an instruction
-fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_id: InsnId, insn: &Insn) -> Option<()> {
+fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_id: InsnId, insn: &Insn) -> Option<()> {
     // Convert InsnId to lir::Opnd
     macro_rules! opnd {
         ($insn_id:ident) => {
@@ -188,8 +249,8 @@ fn gen_insn(jit: &mut JITState, asm: &mut Assembler, function: &Function, insn_i
         Insn::Jump(branch) => return gen_jump(jit, asm, branch),
         Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
         Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
-        Insn::SendWithoutBlock { call_info, cd, state, .. } | Insn::SendWithoutBlockDirect { call_info, cd, state, .. }
-            => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state))?,
+        Insn::SendWithoutBlock { call_info, cd, state, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state))?,
+        Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
@@ -229,7 +290,7 @@ fn gen_ccall(jit: &mut JITState, asm: &mut Assembler, cfun: *const u8, args: &[I
 }
 
 /// Compile an interpreter entry block to be inserted into an ISEQ
-fn gen_entry_prologue(iseq: IseqPtr, asm: &mut Assembler) {
+fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
     asm_comment!(asm, "ZJIT entry point: {}", iseq_get_location(iseq, 0));
     asm.frame_setup();
 
@@ -237,6 +298,10 @@ fn gen_entry_prologue(iseq: IseqPtr, asm: &mut Assembler) {
     asm.cpush(CFP);
     asm.cpush(EC);
     asm.cpush(SP);
+    // On x86_64, maintain 16-byte stack alignment
+    if cfg!(target_arch = "x86_64") {
+        asm.cpush(SP);
+    }
 
     // EC and CFP are pased as arguments
     asm.mov(EC, C_ARG_OPNDS[0]);
@@ -397,6 +462,36 @@ fn gen_send_without_block(
     Some(ret)
 }
 
+/// Compile a direct jump to an ISEQ call without block
+fn gen_send_without_block_direct(
+    cb: &mut CodeBlock,
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    iseq: IseqPtr,
+    recv: Opnd,
+    args: &Vec<InsnId>,
+) -> Option<lir::Opnd> {
+    // Set up the new frame
+    gen_push_frame(asm, recv);
+
+    asm_comment!(asm, "switch to new CFP");
+    let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+
+    // Set up arguments
+    let mut c_args: Vec<Opnd> = vec![];
+    for &arg in args.iter() {
+        c_args.push(jit.get_opnd(arg)?);
+    }
+
+    // Make a method call. The target address will be rewritten once compiled.
+    let branch = Branch::new();
+    let dummy_ptr = cb.get_write_ptr().raw_ptr(cb);
+    jit.branch_iseqs.push((branch.clone(), iseq));
+    Some(asm.ccall_with_branch(dummy_ptr, c_args, &branch))
+}
+
 /// Compile an array duplication instruction
 fn gen_array_dup(
     asm: &mut Assembler,
@@ -423,17 +518,10 @@ fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
     asm.mov(CFP, incr_cfp);
     asm.mov(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
 
-    // Set a return value to the register. We do this before popping SP, EC,
-    // and CFP registers because ret_val may depend on them.
-    asm.mov(C_RET_OPND, val);
-
-    asm_comment!(asm, "exit from leave");
-    asm.cpop_into(SP);
-    asm.cpop_into(EC);
-    asm.cpop_into(CFP);
     asm.frame_teardown();
-    asm.cret(C_RET_OPND);
 
+    // Return from the function
+    asm.cret(val);
     Some(())
 }
 
@@ -560,6 +648,18 @@ fn gen_save_sp(asm: &mut Assembler, state: &FrameState) {
     asm.mov(cfp_sp, sp_addr);
 }
 
+/// Compile an interpreter frame
+fn gen_push_frame(asm: &mut Assembler, recv: Opnd) {
+    // Write to a callee CFP
+    fn cfp_opnd(offset: i32) -> Opnd {
+        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
+    }
+
+    asm_comment!(asm, "push callee control frame");
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), recv);
+    // TODO: Write more fields as needed
+}
+
 /// Return a register we use for the basic block argument at a given index
 fn param_reg(idx: usize) -> Reg {
     // To simplify the implementation, allocate a fixed register for each basic block argument for now.
@@ -574,3 +674,66 @@ fn local_idx_to_ep_offset(iseq: IseqPtr, local_idx: usize) -> i32 {
         .unwrap();
     local_table_size - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
 }
+
+/// Convert ISEQ into High-level IR
+fn compile_iseq(iseq: IseqPtr) -> Option<Function> {
+    let mut function = match iseq_to_hir(iseq) {
+        Ok(function) => function,
+        Err(err) => {
+            debug!("ZJIT: iseq_to_hir: {err:?}");
+            return None;
+        }
+    };
+    function.optimize();
+    Some(function)
+}
+
+impl Assembler {
+    /// Make a C call while marking the start and end positions of it
+    fn ccall_with_branch(&mut self, fptr: *const u8, opnds: Vec<Opnd>, branch: &Rc<Branch>) -> Opnd {
+        // We need to create our own branch rc objects so that we can move the closure below
+        let start_branch = branch.clone();
+        let end_branch = branch.clone();
+
+        self.ccall_with_pos_markers(
+            fptr,
+            opnds,
+            move |code_ptr, _| {
+                start_branch.start_addr.set(Some(code_ptr));
+            },
+            move |code_ptr, _| {
+                end_branch.end_addr.set(Some(code_ptr));
+            },
+        )
+    }
+}
+
+/// Store info about an outgoing branch in a code segment
+#[derive(Debug)]
+struct Branch {
+    /// Position where the generated code starts
+    start_addr: Cell<Option<CodePtr>>,
+
+    /// Position where the generated code ends (exclusive)
+    end_addr: Cell<Option<CodePtr>>,
+}
+
+impl Branch {
+    /// Allocate a new branch
+    fn new() -> Rc<Self> {
+        Rc::new(Branch {
+            start_addr: Cell::new(None),
+            end_addr: Cell::new(None),
+        })
+    }
+
+    /// Regenerate a branch with a given callback
+    fn regenerate(&self, cb: &mut CodeBlock, callback: impl Fn(&mut Assembler)) {
+        cb.with_write_ptr(self.start_addr.get().unwrap(), |cb| {
+            let mut asm = Assembler::new();
+            callback(&mut asm);
+            asm.compile(cb).unwrap();
+            assert_eq!(self.end_addr.get().unwrap(), cb.get_write_ptr());
+        });
+    }
+}

zjit/src/profile.rs

@@ -4,7 +4,7 @@
 use core::ffi::c_void;
 use std::collections::HashMap;
 
-use crate::{cruby::*, hir_type::{types::{Empty, Fixnum}, Type}};
+use crate::{cruby::*, hir_type::{types::{Empty, Fixnum}, Type}, virtualmem::CodePtr};
 
 /// Ephemeral state for profiling runtime information
 struct Profiler {
@@ -95,6 +95,9 @@ fn profile_operands(profiler: &mut Profiler, n: usize) {
 pub struct IseqPayload {
     /// Type information of YARV instruction operands, indexed by the instruction index
     opnd_types: HashMap<usize, Vec<Type>>,
+
+    /// JIT code address of the first block
+    pub start_ptr: Option<CodePtr>,
 }
 
 impl IseqPayload {