14 files changed, 363 insertions, 132 deletions
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index 4aff3193f0..cf328fc68c 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -108,7 +108,6 @@ fn main() {
 
         // From shape.h
         .allowlist_function("rb_obj_shape_id")
-        .allowlist_function("rb_shape_lookup")
         .allowlist_function("rb_shape_id_offset")
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
diff --git a/zjit/src/asm/mod.rs b/zjit/src/asm/mod.rs
index a7f2705af1..0b571f9aff 100644
--- a/zjit/src/asm/mod.rs
+++ b/zjit/src/asm/mod.rs
@@ -1,5 +1,5 @@
 use std::collections::BTreeMap;
-//use std::fmt;
+use std::fmt;
 use std::rc::Rc;
 use std::cell::RefCell;
 use std::mem;
@@ -260,6 +260,18 @@ impl CodeBlock {
     }
 }
 
+/// Produce hex string output from the bytes in a code block
+impl fmt::LowerHex for CodeBlock {
+    fn fmt(&self, fmtr: &mut fmt::Formatter) -> fmt::Result {
+        for pos in 0..self.write_pos {
+            let mem_block = &*self.mem_block.borrow();
+            let byte = unsafe { mem_block.start_ptr().raw_ptr(mem_block).add(pos).read() };
+            fmtr.write_fmt(format_args!("{:02x}", byte))?;
+        }
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 impl CodeBlock {
     /// Stubbed CodeBlock for testing. Can't execute generated code.
diff --git a/zjit/src/asm/x86_64/mod.rs b/zjit/src/asm/x86_64/mod.rs
index efc58dfdb8..fea66c8a3b 100644
--- a/zjit/src/asm/x86_64/mod.rs
+++ b/zjit/src/asm/x86_64/mod.rs
@@ -1024,7 +1024,10 @@ pub fn mov(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
             }
 
             let output_num_bits:u32 = if mem.num_bits > 32 { 32 } else { mem.num_bits.into() };
-            assert!(imm_num_bits(imm.value) <= (output_num_bits as u8));
+            assert!(
+                mem.num_bits < 64 || imm_num_bits(imm.value) <= (output_num_bits as u8),
+                "immediate value should be small enough to survive sign extension"
+            );
             cb.write_int(imm.value as u64, output_num_bits);
         },
         // M + UImm
@@ -1039,7 +1042,10 @@ pub fn mov(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
             }
 
             let output_num_bits = if mem.num_bits > 32 { 32 } else { mem.num_bits.into() };
-            assert!(imm_num_bits(uimm.value as i64) <= (output_num_bits as u8));
+            assert!(
+                mem.num_bits < 64 || imm_num_bits(uimm.value as i64) <= (output_num_bits as u8),
+                "immediate value should be small enough to survive sign extension"
+            );
             cb.write_int(uimm.value, output_num_bits);
         },
         // * + Imm/UImm
diff --git a/zjit/src/asm/x86_64/tests.rs b/zjit/src/asm/x86_64/tests.rs
index f2b949b7f7..ec490fd330 100644
--- a/zjit/src/asm/x86_64/tests.rs
+++ b/zjit/src/asm/x86_64/tests.rs
@@ -1,11 +1,10 @@
 #![cfg(test)]
 
-//use crate::asm::x86_64::*;
+use crate::asm::x86_64::*;
 
-/*
 /// Check that the bytes for an instruction sequence match a hex string
 fn check_bytes<R>(bytes: &str, run: R) where R: FnOnce(&mut super::CodeBlock) {
-    let mut cb = super::CodeBlock::new_dummy(4096);
+    let mut cb = super::CodeBlock::new_dummy();
     run(&mut cb);
     assert_eq!(format!("{:x}", cb), bytes);
 }
@@ -194,6 +193,7 @@ fn test_mov() {
     check_bytes("48c7470801000000", |cb| mov(cb, mem_opnd(64, RDI, 8), imm_opnd(1)));
     //check_bytes("67c7400411000000", |cb| mov(cb, mem_opnd(32, EAX, 4), imm_opnd(0x34))); // We don't distinguish between EAX and RAX here - that's probably fine?
     check_bytes("c7400411000000", |cb| mov(cb, mem_opnd(32, RAX, 4), imm_opnd(17)));
+    check_bytes("c7400401000080", |cb| mov(cb, mem_opnd(32, RAX, 4), uimm_opnd(0x80000001)));
     check_bytes("41895814", |cb| mov(cb, mem_opnd(32, R8, 20), EBX));
     check_bytes("4d8913", |cb| mov(cb, mem_opnd(64, R11, 0), R10));
     check_bytes("48c742f8f4ffffff", |cb| mov(cb, mem_opnd(64, RDX, -8), imm_opnd(-12)));
@@ -439,9 +439,10 @@ fn basic_capstone_usage() -> std::result::Result<(), capstone::Error> {
 }
 
 #[test]
+#[ignore]
 #[cfg(feature = "disasm")]
 fn block_comments() {
-    let mut cb = super::CodeBlock::new_dummy(4096);
+    let mut cb = super::CodeBlock::new_dummy();
 
     let first_write_ptr = cb.get_write_ptr().raw_addr(&cb);
     cb.add_comment("Beginning");
@@ -458,4 +459,3 @@ fn block_comments() {
     assert_eq!(&vec!( "Two bytes in".to_string(), "Still two bytes in".to_string() ), cb.comments_at(second_write_ptr).unwrap());
     assert_eq!(&vec!( "Ten bytes in".to_string() ), cb.comments_at(third_write_ptr).unwrap());
 }
-*/
diff --git a/zjit/src/assertions.rs b/zjit/src/assertions.rs
new file mode 100644
index 0000000000..0dacc938fc
--- /dev/null
+++ b/zjit/src/assertions.rs
@@ -0,0 +1,21 @@
+/// Assert that CodeBlock has the code specified with hex. In addition, if tested with
+/// `cargo test --all-features`, it also checks it generates the specified disasm.
+#[cfg(test)]
+macro_rules! assert_disasm {
+    ($cb:expr, $hex:expr, $disasm:expr) => {
+        #[cfg(feature = "disasm")]
+        {
+            use $crate::disasm::disasm_addr_range;
+            use $crate::cruby::unindent;
+            let disasm = disasm_addr_range(
+                &$cb,
+                $cb.get_ptr(0).raw_addr(&$cb),
+                $cb.get_write_ptr().raw_addr(&$cb),
+            );
+            assert_eq!(unindent(&disasm, false), unindent(&$disasm, true));
+        }
+        assert_eq!(format!("{:x}", $cb), $hex);
+    };
+}
+#[cfg(test)]
+pub(crate) use assert_disasm;
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
index f7e871523e..dd1eb52d34 100644
--- a/zjit/src/backend/arm64/mod.rs
+++ b/zjit/src/backend/arm64/mod.rs
@@ -211,11 +211,6 @@ impl Assembler
         vec![X1_REG, X9_REG, X10_REG, X11_REG, X12_REG, X13_REG, X14_REG, X15_REG]
     }
 
-    /// Get the address that the current frame returns to
-    pub fn return_addr_opnd() -> Opnd {
-        Opnd::Reg(X30_REG)
-    }
-
     /// Split platform-specific instructions
     /// The transformations done here are meant to make our lives simpler in later
     /// stages of the compilation pipeline.
@@ -1345,14 +1340,30 @@ impl Assembler
     }
 }
 
-/*
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::disasm::*;
+    use crate::assertions::assert_disasm;
+
+    static TEMP_REGS: [Reg; 5] = [X1_REG, X9_REG, X10_REG, X14_REG, X15_REG];
 
     fn setup_asm() -> (Assembler, CodeBlock) {
-        (Assembler::new(0), CodeBlock::new_dummy(1024))
+        (Assembler::new(), CodeBlock::new_dummy())
+    }
+
+    #[test]
+    fn test_mul_with_immediate() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let out = asm.mul(Opnd::Reg(TEMP_REGS[1]), 3.into());
+        asm.mov(Opnd::Reg(TEMP_REGS[0]), out);
+        asm.compile_with_num_regs(&mut cb, 2);
+
+        assert_disasm!(cb, "600080d2207d009be10300aa", {"
+            0x0: mov x0, #3
+            0x4: mul x0, x9, x0
+            0x8: mov x1, x0
+        "});
     }
 
     #[test]
@@ -1361,7 +1372,7 @@ mod tests {
 
         let opnd = asm.add(Opnd::Reg(X0_REG), Opnd::Reg(X1_REG));
         asm.store(Opnd::mem(64, Opnd::Reg(X2_REG), 0), opnd);
-        asm.compile_with_regs(&mut cb, None, vec![X3_REG]);
+        asm.compile_with_regs(&mut cb, vec![X3_REG]);
 
         // Assert that only 2 instructions were written.
         assert_eq!(8, cb.get_write_pos());
@@ -1425,6 +1436,7 @@ mod tests {
         asm.compile_with_num_regs(&mut cb, 0);
     }
 
+    /*
     #[test]
     fn test_emit_lea_label() {
         let (mut asm, mut cb) = setup_asm();
@@ -1438,6 +1450,7 @@ mod tests {
 
         asm.compile_with_num_regs(&mut cb, 1);
     }
+    */
 
     #[test]
     fn test_emit_load_mem_disp_fits_into_load() {
@@ -1648,6 +1661,7 @@ mod tests {
         asm.compile_with_num_regs(&mut cb, 2);
     }
 
+    /*
     #[test]
     fn test_bcond_straddling_code_pages() {
         const LANDING_PAGE: usize = 65;
@@ -1784,20 +1798,5 @@ mod tests {
             0x8: mov x1, x11
         "});
     }
-
-    #[test]
-    fn test_mul_with_immediate() {
-        let (mut asm, mut cb) = setup_asm();
-
-        let out = asm.mul(Opnd::Reg(TEMP_REGS[1]), 3.into());
-        asm.mov(Opnd::Reg(TEMP_REGS[0]), out);
-        asm.compile_with_num_regs(&mut cb, 2);
-
-        assert_disasm!(cb, "6b0080d22b7d0b9be1030baa", {"
-            0x0: mov x11, #3
-            0x4: mul x11, x9, x11
-            0x8: mov x1, x11
-        "});
-    }
+    */
 }
-*/
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index e9ae8730f6..f46b35ded5 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 use std::fmt;
 use std::mem::take;
-use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, VM_ENV_DATA_SIZE};
-use crate::state::ZJITState;
+use crate::codegen::local_size_and_idx_to_ep_offset;
+use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32};
 use crate::{cruby::VALUE};
 use crate::backend::current::*;
 use crate::virtualmem::CodePtr;
@@ -1751,6 +1751,15 @@ impl Assembler
         ret
     }
 
+    /// Compile with a limited number of registers. Used only for unit tests.
+    #[cfg(test)]
+    pub fn compile_with_num_regs(self, cb: &mut CodeBlock, num_regs: usize) -> (CodePtr, Vec<u32>)
+    {
+        let mut alloc_regs = Self::get_alloc_regs();
+        let alloc_regs = alloc_regs.drain(0..num_regs).collect();
+        self.compile_with_regs(cb, alloc_regs).unwrap()
+    }
+
     /// Compile Target::SideExit and convert it into Target::CodePtr for all instructions
     #[must_use]
     pub fn compile_side_exits(&mut self) -> Option<()> {
@@ -1788,7 +1797,7 @@ impl Assembler
                 asm_comment!(self, "write locals: {locals:?}");
                 for (idx, &opnd) in locals.iter().enumerate() {
                     let opnd = split_store_source(self, opnd);
-                    self.store(Opnd::mem(64, SP, (-(VM_ENV_DATA_SIZE as i32) - locals.len() as i32 + idx as i32) * SIZEOF_VALUE_I32), opnd);
+                    self.store(Opnd::mem(64, SP, (-local_size_and_idx_to_ep_offset(locals.len(), idx) - 1) * SIZEOF_VALUE_I32), opnd);
                 }
 
                 asm_comment!(self, "save cfp->pc");
@@ -1800,10 +1809,6 @@ impl Assembler
                 let cfp_sp = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP);
                 self.store(cfp_sp, Opnd::Reg(Assembler::SCRATCH_REG));
 
-                asm_comment!(self, "rewind caller frames");
-                self.mov(C_ARG_OPNDS[0], Assembler::return_addr_opnd());
-                self.ccall(Self::rewind_caller_frames as *const u8, vec![]);
-
                 asm_comment!(self, "exit to the interpreter");
                 self.frame_teardown();
                 self.mov(C_RET_OPND, Opnd::UImm(Qundef.as_u64()));
@@ -1814,13 +1819,6 @@ impl Assembler
         }
         Some(())
     }
-
-    #[unsafe(no_mangle)]
-    extern "C" fn rewind_caller_frames(addr: *const u8) {
-        if ZJITState::is_iseq_return_addr(addr) {
-            unimplemented!("Can't side-exit from JIT-JIT call: rewind_caller_frames is not implemented yet");
-        }
-    }
 }
 
 impl fmt::Debug for Assembler {
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
index cf62cdd7f5..d83fc184f9 100644
--- a/zjit/src/backend/x86_64/mod.rs
+++ b/zjit/src/backend/x86_64/mod.rs
@@ -109,11 +109,6 @@ impl Assembler
         vec![RAX_REG, RCX_REG, RDX_REG, RSI_REG, RDI_REG, R8_REG, R9_REG, R10_REG, R11_REG]
     }
 
-    /// Get the address that the current frame returns to
-    pub fn return_addr_opnd() -> Opnd {
-        Opnd::mem(64, Opnd::Reg(RSP_REG), 0)
-    }
-
     // These are the callee-saved registers in the x86-64 SysV ABI
     // RBX, RSP, RBP, and R12–R15
 
@@ -298,19 +293,24 @@ impl Assembler
                             let opnd1 = asm.load(*src);
                             asm.mov(*dest, opnd1);
                         },
-                        (Opnd::Mem(_), Opnd::UImm(value)) => {
-                            // 32-bit values will be sign-extended
-                            if imm_num_bits(*value as i64) > 32 {
+                        (Opnd::Mem(Mem { num_bits, .. }), Opnd::UImm(value)) => {
+                            // For 64 bit destinations, 32-bit values will be sign-extended
+                            if *num_bits == 64 && imm_num_bits(*value as i64) > 32 {
                                 let opnd1 = asm.load(*src);
                                 asm.mov(*dest, opnd1);
                             } else {
                                 asm.mov(*dest, *src);
                             }
                         },
-                        (Opnd::Mem(_), Opnd::Imm(value)) => {
-                            if imm_num_bits(*value) > 32 {
+                        (Opnd::Mem(Mem { num_bits, .. }), Opnd::Imm(value)) => {
+                            // For 64 bit destinations, 32-bit values will be sign-extended
+                            if *num_bits == 64 && imm_num_bits(*value) > 32 {
                                 let opnd1 = asm.load(*src);
                                 asm.mov(*dest, opnd1);
+                            } else if uimm_num_bits(*value as u64) <= *num_bits {
+                                // If the bit string is short enough for the destination, use the unsigned representation.
+                                // Note that 64-bit and negative values are ruled out.
+                                asm.mov(*dest, Opnd::UImm(*value as u64));
                             } else {
                                 asm.mov(*dest, *src);
                             }
@@ -859,20 +859,17 @@ impl Assembler
     }
 }
 
-/*
 #[cfg(test)]
 mod tests {
-    use crate::disasm::assert_disasm;
-    #[cfg(feature = "disasm")]
-    use crate::disasm::{unindent, disasm_addr_range};
-
+    use crate::assertions::assert_disasm;
     use super::*;
 
     fn setup_asm() -> (Assembler, CodeBlock) {
-        (Assembler::new(0), CodeBlock::new_dummy(1024))
+        (Assembler::new(), CodeBlock::new_dummy())
     }
 
     #[test]
+    #[ignore]
     fn test_emit_add_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -883,6 +880,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_add_gt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -893,6 +891,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_and_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -903,6 +902,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_and_gt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -957,6 +957,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_or_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -967,6 +968,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_or_gt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -977,6 +979,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_sub_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -987,6 +990,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_sub_gt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1017,6 +1021,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_xor_lt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1027,6 +1032,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_emit_xor_gt_32_bits() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1050,6 +1056,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_merge_lea_mem() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1064,6 +1071,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_replace_cmp_0() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1216,6 +1224,7 @@ mod tests {
     }
 
     #[test]
+    #[ignore]
     fn test_reorder_c_args_with_insn_out() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1259,15 +1268,16 @@ mod tests {
 
         asm.compile_with_num_regs(&mut cb, 1);
 
-        assert_disasm!(cb, "48837b1001b804000000480f4f03488903", {"
+        assert_disasm!(cb, "48837b1001bf04000000480f4f3b48893b", {"
             0x0: cmp qword ptr [rbx + 0x10], 1
-            0x5: mov eax, 4
-            0xa: cmovg rax, qword ptr [rbx]
-            0xe: mov qword ptr [rbx], rax
+            0x5: mov edi, 4
+            0xa: cmovg rdi, qword ptr [rbx]
+            0xe: mov qword ptr [rbx], rdi
         "});
     }
 
     #[test]
+    #[ignore]
     fn test_csel_split() {
         let (mut asm, mut cb) = setup_asm();
 
@@ -1284,6 +1294,19 @@ mod tests {
             0x13: mov qword ptr [rbx], rax
         "});
     }
-}
 
-*/
+    #[test]
+    fn test_mov_m32_imm32() {
+        let (mut asm, mut cb) = setup_asm();
+
+        let shape_opnd = Opnd::mem(32, C_RET_OPND, 0);
+        asm.mov(shape_opnd, Opnd::UImm(0x8000_0001));
+        asm.mov(shape_opnd, Opnd::Imm(0x8000_0001));
+
+        asm.compile_with_num_regs(&mut cb, 0);
+        assert_disasm!(cb, "c70001000080c70001000080", {"
+            0x0: mov dword ptr [rax], 0x80000001
+            0x6: mov dword ptr [rax], 0x80000001
+        "});
+    }
+}
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 0dbe815c71..f274a64ca6 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -258,7 +258,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         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, self_val, args, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state), self_val, args)?,
-        Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
+        Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), args, &function.frame_state(*state))?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
@@ -275,13 +275,18 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::PatchPoint(_) => return Some(()), // For now, rb_zjit_bop_redefined() panics. TODO: leave a patch point and fix rb_zjit_bop_redefined()
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(jit, asm, *cfun, args)?,
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
-        Insn::SetIvar { self_val, id, val, state: _ } => gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
+        Insn::SetGlobal { id, val, state: _ } => gen_setglobal(asm, *id, opnd!(val)),
+        Insn::GetGlobal { id, state: _ } => gen_getglobal(asm, *id),
+        Insn::SetIvar { self_val, id, val, state: _ } => return gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
+        Insn::SideExit { state } => return gen_side_exit(jit, asm, &function.frame_state(*state)),
         _ => {
             debug!("ZJIT: gen_function: unexpected insn {:?}", insn);
             return None;
         }
     };
 
+    assert!(insn.has_output(), "Cannot write LIR output of HIR instruction with no output");
+
     // If the instruction has an output, remember it in jit.opnds
     jit.opnds[insn_id.0] = Some(out_opnd);
 
@@ -309,14 +314,39 @@ fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID) -> Opnd {
 }
 
 /// Emit an uncached instance variable store
-fn gen_setivar(asm: &mut Assembler, recv: Opnd, id: ID, val: Opnd) -> Opnd {
+fn gen_setivar(asm: &mut Assembler, recv: Opnd, id: ID, val: Opnd) -> Option<()> {
     asm_comment!(asm, "call rb_ivar_set");
     asm.ccall(
         rb_ivar_set as *const u8,
         vec![recv, Opnd::UImm(id.0), val],
+    );
+    Some(())
+}
+
+/// Look up global variables
+fn gen_getglobal(asm: &mut Assembler, id: ID) -> Opnd {
+    asm_comment!(asm, "call rb_gvar_get");
+    asm.ccall(
+        rb_gvar_get as *const u8,
+        vec![Opnd::UImm(id.0)],
+    )
+}
+
+/// Set global variables
+fn gen_setglobal(asm: &mut Assembler, id: ID, val: Opnd) -> Opnd {
+    asm_comment!(asm, "call rb_gvar_set");
+    asm.ccall(
+        rb_gvar_set as *const u8,
+        vec![Opnd::UImm(id.0), val],
     )
 }
 
+/// Side-exit into the interpreter
+fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, state: &FrameState) -> Option<()> {
+    asm.jmp(side_exit(jit, state)?);
+    Some(())
+}
+
 /// Compile an interpreter entry block to be inserted into an ISEQ
 fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
     asm_comment!(asm, "ZJIT entry point: {}", iseq_get_location(iseq, 0));
@@ -464,8 +494,16 @@ fn gen_send_without_block(
     self_val: &InsnId,
     args: &Vec<InsnId>,
 ) -> Option<lir::Opnd> {
-    // Spill the receiver and the arguments onto the stack. They need to be marked by GC and may be caller-saved registers.
+    // Spill locals onto the stack.
+    // TODO: Don't spill locals eagerly; lazily reify frames
+    asm_comment!(asm, "spill locals");
+    for (idx, &insn_id) in state.locals().enumerate() {
+        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
+    }
+    // Spill the receiver and the arguments onto the stack.
+    // They need to be on the interpreter stack to let the interpreter access them.
     // TODO: Avoid spilling operands that have been spilled before.
+    asm_comment!(asm, "spill receiver and arguments");
     for (idx, &insn_id) in [*self_val].iter().chain(args.iter()).enumerate() {
         // Currently, we don't move the SP register. So it's equal to the base pointer.
         let stack_opnd = Opnd::mem(64, SP, idx as i32  * SIZEOF_VALUE_I32);
@@ -495,10 +533,40 @@ fn gen_send_without_block_direct(
     cb: &mut CodeBlock,
     jit: &mut JITState,
     asm: &mut Assembler,
+    cme: *const rb_callable_method_entry_t,
     iseq: IseqPtr,
     recv: Opnd,
     args: &Vec<InsnId>,
+    state: &FrameState,
 ) -> Option<lir::Opnd> {
+    // Save cfp->pc and cfp->sp for the caller frame
+    gen_save_pc(asm, state);
+    gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
+
+    // Spill the virtual stack and the locals of the caller onto the stack
+    // TODO: Lazily materialize caller frames on side exits or when needed
+    asm_comment!(asm, "spill locals and stack");
+    for (idx, &insn_id) in state.locals().enumerate() {
+        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
+    }
+    for (idx, &insn_id) in state.stack().enumerate() {
+        asm.mov(Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
+    }
+
+    // Set up the new frame
+    // TODO: Lazily materialize caller frames on side exits or when needed
+    gen_push_frame(asm, args.len(), state, ControlFrame {
+        recv,
+        iseq,
+        cme,
+        frame_type: VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL,
+    });
+
+    asm_comment!(asm, "switch to new SP register");
+    let local_size = unsafe { get_iseq_body_local_table_size(iseq) } as usize;
+    let new_sp = asm.add(SP, ((state.stack().len() + local_size - args.len() + VM_ENV_DATA_SIZE as usize) * SIZEOF_VALUE).into());
+    asm.mov(SP, new_sp);
+
     asm_comment!(asm, "switch to new CFP");
     let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
     asm.mov(CFP, new_cfp);
@@ -517,7 +585,15 @@ fn gen_send_without_block_direct(
     jit.branch_iseqs.push((branch.clone(), iseq));
     // TODO(max): Add a PatchPoint here that can side-exit the function if the callee messed with
     // the frame's locals
-    Some(asm.ccall_with_branch(dummy_ptr, c_args, &branch))
+    let ret = asm.ccall_with_branch(dummy_ptr, c_args, &branch);
+
+    // If a callee side-exits, i.e. returns Qundef, propagate the return value to the caller.
+    // The caller will side-exit the callee into the interpreter.
+    // TODO: Let side exit code pop all JIT frames to optimize away this cmp + je.
+    asm.cmp(ret, Qundef.into());
+    asm.je(ZJITState::get_exit_trampoline().into());
+
+    Some(ret)
 }
 
 /// Compile an array duplication instruction
@@ -729,6 +805,45 @@ fn gen_save_sp(asm: &mut Assembler, stack_size: usize) {
     asm.mov(cfp_sp, sp_addr);
 }
 
+/// Frame metadata written by gen_push_frame()
+struct ControlFrame {
+    recv: Opnd,
+    iseq: IseqPtr,
+    cme: *const rb_callable_method_entry_t,
+    frame_type: u32,
+}
+
+/// Compile an interpreter frame
+fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: ControlFrame) {
+    // Locals are written by the callee frame on side-exits or non-leaf calls
+
+    // See vm_push_frame() for details
+    asm_comment!(asm, "push cme, specval, frame type");
+    // ep[-2]: cref of cme
+    let local_size = unsafe { get_iseq_body_local_table_size(frame.iseq) } as i32;
+    let ep_offset = state.stack().len() as i32 + local_size - argc as i32 + VM_ENV_DATA_SIZE as i32 - 1;
+    asm.store(Opnd::mem(64, SP, (ep_offset - 2) * SIZEOF_VALUE_I32), VALUE::from(frame.cme).into());
+    // ep[-1]: block_handler or prev EP
+    // block_handler is not supported for now
+    asm.store(Opnd::mem(64, SP, (ep_offset - 1) * SIZEOF_VALUE_I32), VM_BLOCK_HANDLER_NONE.into());
+    // ep[0]: ENV_FLAGS
+    asm.store(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32), frame.frame_type.into());
+
+    // Write to the 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");
+    // cfp_opnd(RUBY_OFFSET_CFP_PC): written by the callee frame on side-exits or non-leaf calls
+    // cfp_opnd(RUBY_OFFSET_CFP_SP): written by the callee frame on side-exits or non-leaf calls
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), VALUE::from(frame.iseq).into());
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);
+    let ep = asm.lea(Opnd::mem(64, SP, ep_offset * SIZEOF_VALUE_I32));
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_EP), ep);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
+}
+
 /// 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.
@@ -744,10 +859,13 @@ fn param_reg(idx: usize) -> Reg {
 
 /// Inverse of ep_offset_to_local_idx(). See ep_offset_to_local_idx() for details.
 fn local_idx_to_ep_offset(iseq: IseqPtr, local_idx: usize) -> i32 {
-    let local_table_size: i32 = unsafe { get_iseq_body_local_table_size(iseq) }
-        .try_into()
-        .unwrap();
-    local_table_size - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
+    let local_size = unsafe { get_iseq_body_local_table_size(iseq) };
+    local_size_and_idx_to_ep_offset(local_size as usize, local_idx)
+}
+
+/// Convert the number of locals and a local index to an offset in the EP
+pub fn local_size_and_idx_to_ep_offset(local_size: usize, local_idx: usize) -> i32 {
+    local_size as i32 - local_idx as i32 - 1 + VM_ENV_DATA_SIZE as i32
 }
 
 /// Convert ISEQ into High-level IR
@@ -796,9 +914,8 @@ impl Assembler {
             move |code_ptr, _| {
                 start_branch.start_addr.set(Some(code_ptr));
             },
-            move |code_ptr, cb| {
+            move |code_ptr, _| {
                 end_branch.end_addr.set(Some(code_ptr));
-                ZJITState::add_iseq_return_addr(code_ptr.raw_ptr(cb));
             },
         )
     }
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index d5be47e026..de1c86e8d6 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -485,18 +485,6 @@ impl VALUE {
         unsafe { rb_obj_shape_id(self) }
     }
 
-    pub fn shape_of(self) -> *mut rb_shape {
-        unsafe {
-            let shape = rb_shape_lookup(self.shape_id_of());
-
-            if shape.is_null() {
-                panic!("Shape should not be null");
-            } else {
-                shape
-            }
-        }
-    }
-
     pub fn embedded_p(self) -> bool {
         unsafe {
             FL_TEST_RAW(self, VALUE(ROBJECT_EMBED as usize)) != VALUE(0)
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index 0447f46fd0..5fb5c2ec02 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -101,10 +101,11 @@ pub const RUBY_FL_PROMOTED: ruby_fl_type = 32;
 pub const RUBY_FL_UNUSED6: ruby_fl_type = 64;
 pub const RUBY_FL_FINALIZE: ruby_fl_type = 128;
 pub const RUBY_FL_TAINT: ruby_fl_type = 0;
+pub const RUBY_FL_EXIVAR: ruby_fl_type = 0;
 pub const RUBY_FL_SHAREABLE: ruby_fl_type = 256;
 pub const RUBY_FL_UNTRUSTED: ruby_fl_type = 0;
 pub const RUBY_FL_UNUSED9: ruby_fl_type = 512;
-pub const RUBY_FL_EXIVAR: ruby_fl_type = 1024;
+pub const RUBY_FL_UNUSED10: ruby_fl_type = 1024;
 pub const RUBY_FL_FREEZE: ruby_fl_type = 2048;
 pub const RUBY_FL_USER0: ruby_fl_type = 4096;
 pub const RUBY_FL_USER1: ruby_fl_type = 8192;
@@ -226,6 +227,7 @@ pub const imemo_parser_strterm: imemo_type = 10;
 pub const imemo_callinfo: imemo_type = 11;
 pub const imemo_callcache: imemo_type = 12;
 pub const imemo_constcache: imemo_type = 13;
+pub const imemo_class_fields: imemo_type = 14;
 pub type imemo_type = u32;
 pub const METHOD_VISI_UNDEF: rb_method_visibility_t = 0;
 pub const METHOD_VISI_PUBLIC: rb_method_visibility_t = 1;
@@ -395,26 +397,6 @@ pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 pub type attr_index_t = u16;
 pub type shape_id_t = u32;
-pub type redblack_id_t = u32;
-pub type redblack_node_t = redblack_node;
-#[repr(C)]
-pub struct rb_shape {
-    pub edges: VALUE,
-    pub edge_name: ID,
-    pub ancestor_index: *mut redblack_node_t,
-    pub parent_id: shape_id_t,
-    pub next_field_index: attr_index_t,
-    pub capacity: attr_index_t,
-    pub type_: u8,
-}
-pub type rb_shape_t = rb_shape;
-#[repr(C)]
-pub struct redblack_node {
-    pub key: ID,
-    pub value: *mut rb_shape_t,
-    pub l: redblack_id_t,
-    pub r: redblack_id_t,
-}
 #[repr(C)]
 pub struct rb_cvar_class_tbl_entry {
     pub index: u32,
@@ -865,7 +847,6 @@ unsafe extern "C" {
     pub fn rb_obj_info(obj: VALUE) -> *const ::std::os::raw::c_char;
     pub fn rb_ec_stack_check(ec: *mut rb_execution_context_struct) -> ::std::os::raw::c_int;
     pub fn rb_shape_id_offset() -> i32;
-    pub fn rb_shape_lookup(shape_id: shape_id_t) -> *mut rb_shape_t;
     pub fn rb_obj_shape_id(obj: VALUE) -> shape_id_t;
     pub fn rb_shape_get_iv_index(shape_id: shape_id_t, id: ID, value: *mut attr_index_t) -> bool;
     pub fn rb_shape_transition_add_ivar_no_warnings(obj: VALUE, id: ID) -> shape_id_t;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index f988b629d9..45a9024ca9 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -392,6 +392,11 @@ pub enum Insn {
     Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId },
     GetConstantPath { ic: *const iseq_inline_constant_cache },
 
+    /// Get a global variable named `id`
+    GetGlobal { id: ID, state: InsnId },
+    /// Set a global variable named `id` to `val`
+    SetGlobal { id: ID, val: InsnId, state: InsnId },
+
     //NewObject?
     /// Get an instance variable `id` from `self_val`
     GetIvar { self_val: InsnId, id: ID, state: InsnId },
@@ -421,7 +426,15 @@ pub enum Insn {
     /// Ignoring keyword arguments etc for now
     SendWithoutBlock { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
     Send { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
-    SendWithoutBlockDirect { self_val: InsnId, call_info: CallInfo, cd: *const rb_call_data, iseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+    SendWithoutBlockDirect {
+        self_val: InsnId,
+        call_info: CallInfo,
+        cd: *const rb_call_data,
+        cme: *const rb_callable_method_entry_t,
+        iseq: IseqPtr,
+        args: Vec<InsnId>,
+        state: InsnId,
+    },
 
     /// Control flow instructions
     Return { val: InsnId },
@@ -459,7 +472,7 @@ impl Insn {
             Insn::ArraySet { .. } | Insn::Snapshot { .. } | Insn::Jump(_)
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
-            | Insn::ArrayPush { .. } | Insn::SideExit { .. } => false,
+            | Insn::ArrayPush { .. } | Insn::SideExit { .. } | Insn::SetGlobal { .. } => false,
             _ => true,
         }
     }
@@ -625,6 +638,8 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::DefinedIvar { self_val, id, .. } => write!(f, "DefinedIvar {self_val}, :{}", id.contents_lossy().into_owned()),
             Insn::GetIvar { self_val, id, .. } => write!(f, "GetIvar {self_val}, :{}", id.contents_lossy().into_owned()),
             Insn::SetIvar { self_val, id, val, .. } => write!(f, "SetIvar {self_val}, :{}, {val}", id.contents_lossy().into_owned()),
+            Insn::GetGlobal { id, .. } => write!(f, "GetGlobal :{}", id.contents_lossy().into_owned()),
+            Insn::SetGlobal { id, val, .. } => write!(f, "SetGlobal :{}, {val}", id.contents_lossy().into_owned()),
             Insn::ToArray { val, .. } => write!(f, "ToArray {val}"),
             Insn::ToNewArray { val, .. } => write!(f, "ToNewArray {val}"),
             Insn::ArrayExtend { left, right, .. } => write!(f, "ArrayExtend {left}, {right}"),
@@ -950,10 +965,11 @@ impl Function {
                 args: args.iter().map(|arg| find!(*arg)).collect(),
                 state: *state,
             },
-            SendWithoutBlockDirect { self_val, call_info, cd, iseq, args, state } => SendWithoutBlockDirect {
+            SendWithoutBlockDirect { self_val, call_info, cd, cme, iseq, args, state } => SendWithoutBlockDirect {
                 self_val: find!(*self_val),
                 call_info: call_info.clone(),
                 cd: *cd,
+                cme: *cme,
                 iseq: *iseq,
                 args: args.iter().map(|arg| find!(*arg)).collect(),
                 state: *state,
@@ -982,6 +998,8 @@ impl Function {
             }
             &NewRange { low, high, flag, state } => NewRange { low: find!(low), high: find!(high), flag, state: find!(state) },
             ArrayMax { elements, state } => ArrayMax { elements: find_vec!(*elements), state: find!(*state) },
+            &GetGlobal { id, state } => GetGlobal { id, state },
+            &SetGlobal { id, val, state } => SetGlobal { id, val: find!(val), state },
             &GetIvar { self_val, id, state } => GetIvar { self_val: find!(self_val), id, state },
             &SetIvar { self_val, id, val, state } => SetIvar { self_val: find!(self_val), id, val, state },
             &ToArray { val, state } => ToArray { val: find!(val), state },
@@ -1012,7 +1030,7 @@ impl Function {
         assert!(self.insns[insn.0].has_output());
         match &self.insns[insn.0] {
             Insn::Param { .. } => unimplemented!("params should not be present in block.insns"),
-            Insn::ArraySet { .. } | Insn::Snapshot { .. } | Insn::Jump(_)
+            Insn::SetGlobal { .. } | Insn::ArraySet { .. } | Insn::Snapshot { .. } | Insn::Jump(_)
             | Insn::IfTrue { .. } | Insn::IfFalse { .. } | Insn::Return { .. }
             | Insn::PatchPoint { .. } | Insn::SetIvar { .. } | Insn::ArrayExtend { .. }
             | Insn::ArrayPush { .. } | Insn::SideExit { .. } =>
@@ -1063,6 +1081,7 @@ impl Function {
             Insn::DefinedIvar { .. } => types::BasicObject,
             Insn::GetConstantPath { .. } => types::BasicObject,
             Insn::ArrayMax { .. } => types::BasicObject,
+            Insn::GetGlobal { .. } => types::BasicObject,
             Insn::GetIvar { .. } => types::BasicObject,
             Insn::ToNewArray { .. } => types::ArrayExact,
             Insn::ToArray { .. } => types::ArrayExact,
@@ -1251,7 +1270,7 @@ impl Function {
                         if let Some(expected) = guard_equal_to {
                             self_val = self.push_insn(block, Insn::GuardBitEquals { val: self_val, expected, state });
                         }
-                        let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { self_val, call_info, cd, iseq, args, state });
+                        let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { self_val, call_info, cd, cme, iseq, args, state });
                         self.make_equal_to(insn_id, send_direct);
                     }
                     Insn::GetConstantPath { ic } => {
@@ -1568,7 +1587,8 @@ impl Function {
                 | Insn::Test { val }
                 | Insn::IsNil { val } =>
                     worklist.push_back(val),
-                Insn::GuardType { val, state, .. }
+                Insn::SetGlobal { val, state, .. }
+                | Insn::GuardType { val, state, .. }
                 | Insn::GuardBitEquals { val, state, .. }
                 | Insn::ToArray { val, state }
                 | Insn::ToNewArray { val, state } => {
@@ -1635,6 +1655,7 @@ impl Function {
                     worklist.push_back(val);
                     worklist.push_back(state);
                 }
+                Insn::GetGlobal { state, .. } |
                 Insn::SideExit { state } => worklist.push_back(state),
             }
         }
@@ -2380,6 +2401,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 YARVINSN_opt_and |
                 YARVINSN_opt_or |
                 YARVINSN_opt_not |
+                YARVINSN_opt_regexpmatch2 |
                 YARVINSN_opt_send_without_block => {
                     let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
                     let call_info = unsafe { rb_get_call_data_ci(cd) };
@@ -2434,6 +2456,18 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let send = fun.push_insn(block, Insn::Send { self_val: recv, call_info: CallInfo { method_name }, cd, blockiseq, args, state: exit_id });
                     state.stack_push(send);
                 }
+                YARVINSN_getglobal => {
+                    let id = ID(get_arg(pc, 0).as_u64());
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let result = fun.push_insn(block, Insn::GetGlobal { id, state: exit_id });
+                    state.stack_push(result);
+                }
+                YARVINSN_setglobal => {
+                    let id = ID(get_arg(pc, 0).as_u64());
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let val = state.stack_pop()?;
+                    fun.push_insn(block, Insn::SetGlobal { id, val, state: exit_id });
+                }
                 YARVINSN_getinstancevariable => {
                     let id = ID(get_arg(pc, 0).as_u64());
                     // ic is in arg 1
@@ -3711,6 +3745,35 @@ mod tests {
     }
 
     #[test]
+    fn test_setglobal() {
+        eval("
+            def test = $foo = 1
+            test
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_setglobal, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v2:Fixnum[1] = Const Value(1)
+              SetGlobal :$foo, v2
+              Return v2
+        "#]]);
+    }
+
+    #[test]
+    fn test_getglobal() {
+        eval("
+            def test = $foo
+            test
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_getglobal, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              v3:BasicObject = GetGlobal :$foo
+              Return v3
+        "#]]);
+    }
+
+    #[test]
     fn test_splatarray_mut() {
         eval("
             def test(a) = [*a]
@@ -3879,6 +3942,20 @@ mod tests {
               Return v4
         "#]]);
     }
+
+    #[test]
+    fn opt_regexpmatch2() {
+        eval("
+            def test(regexp, matchee) = regexp =~ matchee
+        ");
+        assert_method_hir_with_opcode("test", YARVINSN_opt_regexpmatch2, expect![[r#"
+            fn test:
+            bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+              v5:BasicObject = SendWithoutBlock v1, :=~, v2
+              Return v5
+        "#]]);
+    }
+
     #[test]
     fn test_branchnil() {
         eval("
diff --git a/zjit/src/lib.rs b/zjit/src/lib.rs
index 8ccb6ae4c1..9d139b9801 100644
--- a/zjit/src/lib.rs
+++ b/zjit/src/lib.rs
@@ -21,3 +21,5 @@ mod disasm;
 mod options;
 mod profile;
 mod invariants;
+#[cfg(test)]
+mod assertions;
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index e8c389a5f8..acaac850c3 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -1,10 +1,10 @@
-use std::collections::HashSet;
-
 use crate::cruby::{self, rb_bug_panic_hook, EcPtr, Qnil, VALUE};
 use crate::cruby_methods;
 use crate::invariants::Invariants;
 use crate::options::Options;
 use crate::asm::CodeBlock;
+use crate::backend::lir::{Assembler, C_RET_OPND};
+use crate::virtualmem::CodePtr;
 
 #[allow(non_upper_case_globals)]
 #[unsafe(no_mangle)]
@@ -32,8 +32,8 @@ pub struct ZJITState {
     /// Properties of core library methods
     method_annotations: cruby_methods::Annotations,
 
-    /// The address of the instruction that JIT-to-JIT calls return to
-    iseq_return_addrs: HashSet<*const u8>,
+    /// Trampoline to propagate a callee's side exit to the caller
+    exit_trampoline: Option<CodePtr>,
 }
 
 /// Private singleton instance of the codegen globals
@@ -88,9 +88,14 @@ impl ZJITState {
             invariants: Invariants::default(),
             assert_compiles: false,
             method_annotations: cruby_methods::init(),
-            iseq_return_addrs: HashSet::new(),
+            exit_trampoline: None,
         };
         unsafe { ZJIT_STATE = Some(zjit_state); }
+
+        // Generate trampolines after initializing ZJITState, which Assembler will use
+        let cb = ZJITState::get_code_block();
+        let exit_trampoline = Self::gen_exit_trampoline(cb).unwrap();
+        ZJITState::get_instance().exit_trampoline = Some(exit_trampoline);
     }
 
     /// Return true if zjit_state has been initialized
@@ -133,14 +138,17 @@ impl ZJITState {
         instance.assert_compiles = true;
     }
 
-    /// Record an address that a JIT-to-JIT call returns to
-    pub fn add_iseq_return_addr(addr: *const u8) {
-        ZJITState::get_instance().iseq_return_addrs.insert(addr);
+    /// Generate a trampoline to propagate a callee's side exit to the caller
+    fn gen_exit_trampoline(cb: &mut CodeBlock) -> Option<CodePtr> {
+        let mut asm = Assembler::new();
+        asm.frame_teardown();
+        asm.cret(C_RET_OPND);
+        asm.compile(cb).map(|(start_ptr, _)| start_ptr)
     }
 
-    /// Returns true if a JIT-to-JIT call returns to a given address
-    pub fn is_iseq_return_addr(addr: *const u8) -> bool {
-        ZJITState::get_instance().iseq_return_addrs.contains(&addr)
+    /// Get the trampoline to propagate a callee's side exit to the caller
+    pub fn get_exit_trampoline() -> CodePtr {
+        ZJITState::get_instance().exit_trampoline.unwrap()
     }
 }