YJIT: Introduce Target::SideExit (#7712)

* YJIT: Introduce Target::SideExit

* YJIT: Obviate Insn::SideExitContext

* YJIT: Avoid cloning a Context for each insn

9 files changed, 374 insertions, 291 deletions

diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 4a052bd0ef..8aa8124089 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -2,8 +2,10 @@
 #![allow(unused_variables)]
 #![allow(unused_imports)]
 
+use std::mem::take;
+
 use crate::asm::x86_64::jmp_ptr;
-use crate::asm::{CodeBlock};
+use crate::asm::{CodeBlock, OutlinedCb};
 use crate::asm::arm64::*;
 use crate::codegen::{JITState, CodegenGlobals};
 use crate::core::Context;
@@ -374,7 +376,7 @@ impl Assembler
             }
         }
 
-        let mut asm_local = Assembler::new_with_label_names(std::mem::take(&mut self.label_names));
+        let mut asm_local = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
         let asm = &mut asm_local;
         let mut iterator = self.into_draining_iter();
 
@@ -675,7 +677,7 @@ impl Assembler
 
     /// Emit platform-specific machine code
     /// Returns a list of GC offsets. Can return failure to signal caller to retry.
-    fn arm64_emit(&mut self, cb: &mut CodeBlock) -> Result<Vec<u32>, ()> {
+    fn arm64_emit(&mut self, cb: &mut CodeBlock, ocb: &mut Option<&mut OutlinedCb>) -> Result<Vec<u32>, ()> {
         /// Determine how many instructions it will take to represent moving
         /// this value into a register. Note that the return value of this
         /// function must correspond to how many instructions are used to
@@ -765,6 +767,9 @@ impl Assembler
                         bcond(cb, CONDITION, InstructionOffset::from_bytes(bytes));
                     });
                 },
+                Target::SideExit { .. } => {
+                    unreachable!("Target::SideExit should have been compiled by compile_side_exit")
+                },
             };
         }
 
@@ -780,6 +785,20 @@ impl Assembler
             ldr_post(cb, opnd, A64Opnd::new_mem(64, C_SP_REG, C_SP_STEP));
         }
 
+        /// Compile a side exit if Target::SideExit is given.
+        fn compile_side_exit(
+            target: Target,
+            asm: &mut Assembler,
+            ocb: &mut Option<&mut OutlinedCb>,
+        ) -> Target {
+            if let Target::SideExit { counter, context } = target {
+                let side_exit = asm.get_side_exit(&context.unwrap(), counter, ocb.as_mut().unwrap());
+                Target::SideExitPtr(side_exit)
+            } else {
+                target
+            }
+        }
+
         // dbg!(&self.insns);
 
         // List of GC offsets
@@ -1016,12 +1035,12 @@ impl Assembler
                     br(cb, opnd.into());
                 },
                 Insn::Jmp(target) => {
-                    match target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(dst_ptr) => {
-                            emit_jmp_ptr(cb, *dst_ptr, true);
+                            emit_jmp_ptr(cb, dst_ptr, true);
                         },
                         Target::SideExitPtr(dst_ptr) => {
-                            emit_jmp_ptr(cb, *dst_ptr, false);
+                            emit_jmp_ptr(cb, dst_ptr, false);
                         },
                         Target::Label(label_idx) => {
                             // Here we're going to save enough space for
@@ -1029,27 +1048,30 @@ impl Assembler
                             // instruction once we know the offset. We're going
                             // to assume we can fit into a single b instruction.
                             // It will panic otherwise.
-                            cb.label_ref(*label_idx, 4, |cb, src_addr, dst_addr| {
+                            cb.label_ref(label_idx, 4, |cb, src_addr, dst_addr| {
                                 let bytes: i32 = (dst_addr - (src_addr - 4)).try_into().unwrap();
                                 b(cb, InstructionOffset::from_bytes(bytes));
                             });
                         },
+                        Target::SideExit { .. } => {
+                            unreachable!("Target::SideExit should have been compiled by compile_side_exit")
+                        },
                     };
                 },
                 Insn::Je(target) | Insn::Jz(target) => {
-                    emit_conditional_jump::<{Condition::EQ}>(cb, *target);
+                    emit_conditional_jump::<{Condition::EQ}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::Jne(target) | Insn::Jnz(target) => {
-                    emit_conditional_jump::<{Condition::NE}>(cb, *target);
+                    emit_conditional_jump::<{Condition::NE}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::Jl(target) => {
-                    emit_conditional_jump::<{Condition::LT}>(cb, *target);
+                    emit_conditional_jump::<{Condition::LT}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::Jbe(target) => {
-                    emit_conditional_jump::<{Condition::LS}>(cb, *target);
+                    emit_conditional_jump::<{Condition::LS}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::Jo(target) => {
-                    emit_conditional_jump::<{Condition::VS}>(cb, *target);
+                    emit_conditional_jump::<{Condition::VS}>(cb, compile_side_exit(*target, self, ocb));
                 },
                 Insn::IncrCounter { mem, value } => {
                     let label = cb.new_label("incr_counter_loop".to_string());
@@ -1121,7 +1143,7 @@ impl Assembler
     }
 
     /// Optimize and compile the stored instructions
-    pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
+    pub fn compile_with_regs(self, cb: &mut CodeBlock, ocb: Option<&mut OutlinedCb>, regs: Vec<Reg>) -> Vec<u32>
     {
         let asm = self.lower_stack();
         let asm = asm.arm64_split();
@@ -1135,14 +1157,15 @@ impl Assembler
 
         let start_ptr = cb.get_write_ptr();
         let starting_label_state = cb.get_label_state();
-        let gc_offsets = asm.arm64_emit(cb)
+        let mut ocb = ocb; // for &mut
+        let gc_offsets = asm.arm64_emit(cb, &mut ocb)
             .unwrap_or_else(|_err| {
                 // we want to lower jumps to labels to b.cond instructions, which have a 1 MiB
                 // range limit. We can easily exceed the limit in case the jump straddles two pages.
                 // In this case, we retry with a fresh page.
                 cb.set_label_state(starting_label_state);
                 cb.next_page(start_ptr, emit_jmp_ptr_with_invalidation);
-                asm.arm64_emit(cb).expect("should not fail when writing to a fresh code page")
+                asm.arm64_emit(cb, &mut ocb).expect("should not fail when writing to a fresh code page")
             });
 
         if cb.has_dropped_bytes() {
@@ -1180,7 +1203,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, vec![X3_REG]);
+        asm.compile_with_regs(&mut cb, None, vec![X3_REG]);
 
         // Assert that only 2 instructions were written.
         assert_eq!(8, cb.get_write_pos());
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index d22a5ff55b..1bf2dca04e 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -3,13 +3,15 @@
 #![allow(unused_imports)]
 
 use std::cell::Cell;
+use std::collections::HashMap;
 use std::fmt;
 use std::convert::From;
 use std::io::Write;
 use std::mem::take;
+use crate::codegen::{gen_outlined_exit, gen_counted_exit};
 use crate::cruby::{VALUE, SIZEOF_VALUE_I32};
 use crate::virtualmem::{CodePtr};
-use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits};
+use crate::asm::{CodeBlock, uimm_num_bits, imm_num_bits, OutlinedCb};
 use crate::core::{Context, Type, TempMapping, RegTemps, MAX_REG_TEMPS, MAX_TEMP_TYPES};
 use crate::options::*;
 use crate::stats::*;
@@ -280,13 +282,22 @@ impl From<VALUE> for Opnd {
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub enum Target
 {
-    CodePtr(CodePtr),     // Pointer to a piece of YJIT-generated code
-    SideExitPtr(CodePtr), // Pointer to a side exit code
-    Label(usize),         // A label within the generated code
+    /// Pointer to a piece of YJIT-generated code
+    CodePtr(CodePtr),
+    /// Side exit with a counter
+    SideExit { counter: Option<Counter>, context: Option<SideExitContext> },
+    /// Pointer to a side exit code
+    SideExitPtr(CodePtr),
+    /// A label within the generated code
+    Label(usize),
 }
 
 impl Target
 {
+    pub fn side_exit(counter: Option<Counter>) -> Target {
+        Target::SideExit { counter, context: None }
+    }
+
     pub fn unwrap_label_idx(&self) -> usize {
         match self {
             Target::Label(idx) => *idx,
@@ -500,6 +511,25 @@ impl Insn {
         InsnOpndMutIterator::new(self)
     }
 
+    /// Get a mutable reference to a Target if it exists.
+    pub(super) fn target_mut(&mut self) -> Option<&mut Target> {
+        match self {
+            Insn::Jbe(target) |
+            Insn::Je(target) |
+            Insn::Jl(target) |
+            Insn::Jmp(target) |
+            Insn::Jne(target) |
+            Insn::Jnz(target) |
+            Insn::Jo(target) |
+            Insn::Jz(target) |
+            Insn::Label(target) |
+            Insn::LeaLabel { target, .. } => {
+                Some(target)
+            }
+            _ => None,
+        }
+    }
+
     /// Returns a string that describes which operation this instruction is
     /// performing. This is used for debugging.
     fn op(&self) -> &'static str {
@@ -880,10 +910,19 @@ impl fmt::Debug for Insn {
     }
 }
 
+/// Set of variables used for generating side exits
+#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
+pub struct SideExitContext {
+    /// PC of the instruction being compiled
+    pub pc: *mut VALUE,
+
+    /// Context when it started to compile the instruction
+    pub ctx: Context,
+}
+
 /// Object into which we assemble instructions to be
 /// optimized and lowered
-pub struct Assembler
-{
+pub struct Assembler {
     pub(super) insns: Vec<Insn>,
 
     /// Parallel vec with insns
@@ -899,21 +938,33 @@ pub struct Assembler
 
     /// Context for generating the current insn
     pub ctx: Context,
+
+    /// Side exit caches for each SideExitContext
+    pub(super) side_exits: HashMap<SideExitContext, CodePtr>,
+
+    /// PC for Target::SideExit
+    side_exit_pc: Option<*mut VALUE>,
+
+    /// Stack size for Target::SideExit
+    side_exit_stack_size: Option<u8>,
 }
 
 impl Assembler
 {
     pub fn new() -> Self {
-        Self::new_with_label_names(Vec::default())
+        Self::new_with_label_names(Vec::default(), HashMap::default())
     }
 
-    pub fn new_with_label_names(label_names: Vec<String>) -> Self {
+    pub fn new_with_label_names(label_names: Vec<String>, side_exits: HashMap<SideExitContext, CodePtr>) -> Self {
         Self {
             insns: Vec::default(),
             live_ranges: Vec::default(),
             reg_temps: Vec::default(),
             label_names,
             ctx: Context::default(),
+            side_exits,
+            side_exit_pc: None,
+            side_exit_stack_size: None,
         }
     }
 
@@ -924,6 +975,12 @@ impl Assembler
         regs.drain(0..num_regs).collect()
     }
 
+    /// Set a context for generating side exits
+    pub fn set_side_exit_context(&mut self, pc: *mut VALUE, stack_size: u8) {
+        self.side_exit_pc = Some(pc);
+        self.side_exit_stack_size = Some(stack_size);
+    }
+
     /// Build an Opnd::InsnOut from the current index of the assembler and the
     /// given number of bits.
     pub(super) fn next_opnd_out(&self, num_bits: u8) -> Opnd {
@@ -973,6 +1030,18 @@ impl Assembler
             }
         }
 
+        // Set a side exit context to Target::SideExit
+        let mut insn = insn;
+        if let Some(Target::SideExit { context, .. }) = insn.target_mut() {
+            // We should skip this when this instruction is being copied from another Assembler.
+            if context.is_none() {
+                *context = Some(SideExitContext {
+                    pc: self.side_exit_pc.unwrap(),
+                    ctx: self.ctx.with_stack_size(self.side_exit_stack_size.unwrap()),
+                });
+            }
+        }
+
         self.insns.push(insn);
         self.live_ranges.push(insn_idx);
         self.reg_temps.push(reg_temps);
@@ -983,6 +1052,26 @@ impl Assembler
         *self.reg_temps.last().unwrap_or(&RegTemps::default())
     }
 
+    /// Get a cached side exit, wrapping a counter if specified
+    pub fn get_side_exit(&mut self, side_exit_context: &SideExitContext, counter: Option<Counter>, ocb: &mut OutlinedCb) -> CodePtr {
+        // Drop type information from a cache key
+        let mut side_exit_context = side_exit_context.clone();
+        side_exit_context.ctx = side_exit_context.ctx.get_generic_ctx();
+
+        // Get a cached side exit
+        let side_exit = match self.side_exits.get(&side_exit_context) {
+            None => {
+                let exit_code = gen_outlined_exit(side_exit_context.pc, &side_exit_context.ctx, ocb);
+                self.side_exits.insert(side_exit_context.clone(), exit_code);
+                exit_code
+            }
+            Some(code_ptr) => *code_ptr,
+        };
+
+        // Wrap a counter if needed
+        gen_counted_exit(side_exit, ocb, counter)
+    }
+
     /// Create a new label instance that we can jump to
     pub fn new_label(&mut self, name: &str) -> Target
     {
@@ -1016,7 +1105,7 @@ impl Assembler
             }
         }
 
-        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
         let regs = Assembler::get_temp_regs();
         let reg_temps = take(&mut self.reg_temps);
         let mut iterator = self.into_draining_iter();
@@ -1172,7 +1261,7 @@ impl Assembler
         }
 
         let live_ranges: Vec<usize> = take(&mut self.live_ranges);
-        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
         let mut iterator = self.into_draining_iter();
 
         while let Some((index, mut insn)) = iterator.next_unmapped() {
@@ -1305,13 +1394,13 @@ impl Assembler
     /// Compile the instructions down to machine code
     /// NOTE: should compile return a list of block labels to enable
     ///       compiling multiple blocks at a time?
-    pub fn compile(self, cb: &mut CodeBlock) -> Vec<u32>
+    pub fn compile(self, cb: &mut CodeBlock, ocb: Option<&mut OutlinedCb>) -> Vec<u32>
     {
         #[cfg(feature = "disasm")]
         let start_addr = cb.get_write_ptr();
 
         let alloc_regs = Self::get_alloc_regs();
-        let gc_offsets = self.compile_with_regs(cb, alloc_regs);
+        let gc_offsets = self.compile_with_regs(cb, ocb, alloc_regs);
 
         #[cfg(feature = "disasm")]
         if let Some(dump_disasm) = get_option_ref!(dump_disasm) {
@@ -1327,7 +1416,7 @@ impl Assembler
     {
         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)
+        self.compile_with_regs(cb, None, alloc_regs)
     }
 
     /// Consume the assembler by creating a new draining iterator.
diff --git a/yjit/src/backend/tests.rs b/yjit/src/backend/tests.rs
index 3098c7e3b0..8ba9f61d25 100644
--- a/yjit/src/backend/tests.rs
+++ b/yjit/src/backend/tests.rs
@@ -199,7 +199,7 @@ fn test_alloc_ccall_regs() {
     let out2 = asm.ccall(0 as *const u8, vec![out1]);
     asm.mov(EC, out2);
     let mut cb = CodeBlock::new_dummy(1024);
-    asm.compile_with_regs(&mut cb, Assembler::get_alloc_regs());
+    asm.compile_with_regs(&mut cb, None, Assembler::get_alloc_regs());
 }
 
 #[test]
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index 170fc888e0..de38cd98d5 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -7,6 +7,7 @@ use std::mem::take;
 use crate::asm::*;
 use crate::asm::x86_64::*;
 use crate::codegen::{JITState};
+use crate::core::Context;
 use crate::cruby::*;
 use crate::backend::ir::*;
 use crate::codegen::CodegenGlobals;
@@ -115,7 +116,7 @@ impl Assembler
     fn x86_split(mut self) -> Assembler
     {
         let live_ranges: Vec<usize> = take(&mut self.live_ranges);
-        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names));
+        let mut asm = Assembler::new_with_label_names(take(&mut self.label_names), take(&mut self.side_exits));
         let mut iterator = self.into_draining_iter();
 
         while let Some((index, mut insn)) = iterator.next_unmapped() {
@@ -381,7 +382,7 @@ impl Assembler
     }
 
     /// Emit platform-specific machine code
-    pub fn x86_emit(&mut self, cb: &mut CodeBlock) -> Vec<u32>
+    pub fn x86_emit(&mut self, cb: &mut CodeBlock, ocb: &mut Option<&mut OutlinedCb>) -> Vec<u32>
     {
         /// For some instructions, we want to be able to lower a 64-bit operand
         /// without requiring more registers to be available in the register
@@ -411,6 +412,19 @@ impl Assembler
             }
         }
 
+        /// Compile a side exit if Target::SideExit is given.
+        fn compile_side_exit(
+            target: Target,
+            asm: &mut Assembler,
+            ocb: &mut Option<&mut OutlinedCb>,
+        ) -> Target {
+            if let Target::SideExit { counter, context } = target {
+                let side_exit = asm.get_side_exit(&context.unwrap(), counter, ocb.as_mut().unwrap());
+                Target::SideExitPtr(side_exit)
+            } else {
+                target
+            }
+        }
 
         fn emit_csel(cb: &mut CodeBlock, truthy: Opnd, falsy: Opnd, out: Opnd, cmov_fn: fn(&mut CodeBlock, X86Opnd, X86Opnd)) {
             if out != truthy {
@@ -426,8 +440,8 @@ impl Assembler
 
         // For each instruction
         let start_write_pos = cb.get_write_pos();
-        let mut insns_idx: usize = 0;
-        while let Some(insn) = self.insns.get(insns_idx) {
+        let mut insn_idx: usize = 0;
+        while let Some(insn) = self.insns.get(insn_idx) {
             let src_ptr = cb.get_write_ptr();
             let had_dropped_bytes = cb.has_dropped_bytes();
             let old_label_state = cb.get_label_state();
@@ -626,58 +640,66 @@ impl Assembler
 
                 // Conditional jump to a label
                 Insn::Jmp(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jmp_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jmp_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
                 Insn::Je(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => je_ptr(cb, code_ptr),
                         Target::Label(label_idx) => je_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
                 Insn::Jne(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jne_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jne_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
                 Insn::Jl(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jl_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jl_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 },
 
                 Insn::Jbe(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jbe_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jbe_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 },
 
                 Insn::Jz(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jz_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jz_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
                 Insn::Jnz(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jnz_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jnz_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
                 Insn::Jo(target) => {
-                    match *target {
+                    match compile_side_exit(*target, self, ocb) {
                         Target::CodePtr(code_ptr) | Target::SideExitPtr(code_ptr) => jo_ptr(cb, code_ptr),
                         Target::Label(label_idx) => jo_label(cb, label_idx),
+                        Target::SideExit { .. } => unreachable!("Target::SideExit should have been compiled by compile_side_exit"),
                     }
                 }
 
@@ -724,13 +746,6 @@ impl Assembler
                         nop(cb, (cb.jmp_ptr_bytes() - code_size) as u32);
                     }
                 }
-
-                // We want to keep the panic here because some instructions that
-                // we feed to the backend could get lowered into other
-                // instructions. So it's possible that some of our backend
-                // instructions can never make it to the emit stage.
-                #[allow(unreachable_patterns)]
-                _ => panic!("unsupported instruction passed to x86 backend: {:?}", insn)
             };
 
             // On failure, jump to the next page and retry the current insn
@@ -738,7 +753,7 @@ impl Assembler
                 // Reset cb states before retrying the current Insn
                 cb.set_label_state(old_label_state);
             } else {
-                insns_idx += 1;
+                insn_idx += 1;
                 gc_offsets.append(&mut insn_gc_offsets);
             }
         }
@@ -747,8 +762,7 @@ impl Assembler
     }
 
     /// Optimize and compile the stored instructions
-    pub fn compile_with_regs(self, cb: &mut CodeBlock, regs: Vec<Reg>) -> Vec<u32>
-    {
+    pub fn compile_with_regs(self, cb: &mut CodeBlock, ocb: Option<&mut OutlinedCb>, regs: Vec<Reg>) -> Vec<u32> {
         let asm = self.lower_stack();
         let asm = asm.x86_split();
         let mut asm = asm.alloc_regs(regs);
@@ -759,7 +773,8 @@ impl Assembler
             assert!(label_idx == idx);
         }
 
-        let gc_offsets = asm.x86_emit(cb);
+        let mut ocb = ocb; // for &mut
+        let gc_offsets = asm.x86_emit(cb, &mut ocb);
 
         if cb.has_dropped_bytes() {
             cb.clear_labels();
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 2debb2e270..e135389ecf 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -39,13 +39,6 @@ type InsnGenFn = fn(
     ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus>;
 
-/// Subset of Context that matters for generating a side exit.
-#[derive(Eq, Hash, PartialEq)]
-struct SideExitContext {
-    sp_offset: i8,
-    reg_temps: RegTemps,
-}
-
 /// Ephemeral code generation state.
 /// Represents a [core::Block] while we build it.
 pub struct JITState {
@@ -73,10 +66,6 @@ pub struct JITState {
     /// stack_size when it started to compile the current instruction.
     stack_size_for_pc: u8,
 
-    /// Side exit to the instruction being compiled. See :side-exit:.
-    /// For the current PC, it's cached for each (sp_offset, reg_temps).
-    side_exit_for_pc: HashMap<SideExitContext, CodePtr>,
-
     /// Execution context when compilation started
     /// This allows us to peek at run-time values
     ec: EcPtr,
@@ -121,7 +110,6 @@ impl JITState {
             opcode: 0,
             pc: ptr::null_mut::<VALUE>(),
             stack_size_for_pc: starting_ctx.get_stack_size(),
-            side_exit_for_pc: HashMap::new(),
             pending_outgoing: vec![],
             ec,
             record_boundary_patch_point: false,
@@ -226,8 +214,8 @@ impl JITState {
         }
     }
 
-    pub fn assume_method_lookup_stable(&mut self, ocb: &mut OutlinedCb, cme: CmePtr) {
-        jit_ensure_block_entry_exit(self, ocb);
+    pub fn assume_method_lookup_stable(&mut self, asm: &mut Assembler, ocb: &mut OutlinedCb, cme: CmePtr) {
+        jit_ensure_block_entry_exit(self, asm, ocb);
         self.method_lookup_assumptions.push(cme);
     }
 
@@ -235,8 +223,8 @@ impl JITState {
         unsafe { get_ec_cfp(self.ec) }
     }
 
-    pub fn assume_stable_constant_names(&mut self, ocb: &mut OutlinedCb, id: *const ID) {
-        jit_ensure_block_entry_exit(self, ocb);
+    pub fn assume_stable_constant_names(&mut self, asm: &mut Assembler, ocb: &mut OutlinedCb, id: *const ID) {
+        jit_ensure_block_entry_exit(self, asm, ocb);
         self.stable_constant_names_assumption = Some(id);
     }
 
@@ -274,12 +262,6 @@ macro_rules! gen_counter_incr {
     };
 }
 
-macro_rules! counted_exit {
-    ($jit:tt, $ctx:expr, $ocb:tt, $counter_name:ident) => {
-        counted_exit($jit, $ctx, $ocb, Some(Counter::$counter_name))
-    };
-}
-
 // Save the incremented PC on the CFP
 // This is necessary when callees can raise or allocate
 fn jit_save_pc(jit: &JITState, asm: &mut Assembler) {
@@ -441,7 +423,7 @@ fn gen_code_for_exit_from_stub(ocb: &mut OutlinedCb) -> CodePtr {
 
     asm.cret(Qundef.into());
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     code_ptr
 }
@@ -499,8 +481,19 @@ fn gen_exit(exit_pc: *mut VALUE, asm: &mut Assembler) {
     asm.cret(Qundef.into());
 }
 
-/// Generate an exit to the interpreter in the outlined code block
-fn gen_outlined_exit(exit_pc: *mut VALUE, ctx: &Context, ocb: &mut OutlinedCb) -> CodePtr {
+/// :side-exit:
+/// Get an exit for the current instruction in the outlined block. The code
+/// for each instruction often begins with several guards before proceeding
+/// to do work. When guards fail, an option we have is to exit to the
+/// interpreter at an instruction boundary. The piece of code that takes
+/// care of reconstructing interpreter state and exiting out of generated
+/// code is called the side exit.
+///
+/// No guards change the logic for reconstructing interpreter state at the
+/// moment, so there is one unique side exit for each context. Note that
+/// it's incorrect to jump to the side exit after any ctx stack push operations
+/// since they change the logic required for reconstructing interpreter state.
+pub fn gen_outlined_exit(exit_pc: *mut VALUE, ctx: &Context, ocb: &mut OutlinedCb) -> CodePtr {
     let mut cb = ocb.unwrap();
     let exit_code = cb.get_write_ptr();
     let mut asm = Assembler::new();
@@ -509,46 +502,13 @@ fn gen_outlined_exit(exit_pc: *mut VALUE, ctx: &Context, ocb: &mut OutlinedCb) -
 
     gen_exit(exit_pc, &mut asm);
 
-    asm.compile(&mut cb);
+    asm.compile(&mut cb, None);
 
     exit_code
 }
 
-// :side-exit:
-// Get an exit for the current instruction in the outlined block. The code
-// for each instruction often begins with several guards before proceeding
-// to do work. When guards fail, an option we have is to exit to the
-// interpreter at an instruction boundary. The piece of code that takes
-// care of reconstructing interpreter state and exiting out of generated
-// code is called the side exit.
-//
-// No guards change the logic for reconstructing interpreter state at the
-// moment, so there is one unique side exit for each context. Note that
-// it's incorrect to jump to the side exit after any ctx stack push operations
-// since they change the logic required for reconstructing interpreter state.
-fn side_exit(jit: &mut JITState, ctx: &Context, ocb: &mut OutlinedCb) -> Target {
-    // We use the latest ctx.sp_offset to generate a side exit to tolerate sp_offset changes by gen_save_sp.
-    // We also need to use the latest ctx when we implement stack temp register allocation in the future.
-    // However, we want to simulate an old stack_size when we take a side exit. We do that by adjusting the
-    // sp_offset because gen_outlined_exit uses ctx.sp_offset to move SP.
-    let ctx = ctx.with_stack_size(jit.stack_size_for_pc);
-
-    // Cache a side exit for each (sp_offset, reg_temps).
-    let exit_ctx = SideExitContext { sp_offset: ctx.get_sp_offset(), reg_temps: ctx.get_reg_temps() };
-    match jit.side_exit_for_pc.get(&exit_ctx) {
-        None => {
-            let exit_code = gen_outlined_exit(jit.pc, &ctx, ocb);
-            jit.side_exit_for_pc.insert(exit_ctx, exit_code);
-            exit_code.as_side_exit()
-        }
-        Some(code_ptr) => code_ptr.as_side_exit()
-    }
-}
-
 /// Get a side exit. Increment a counter in it if --yjit-stats is enabled.
-fn counted_exit(jit: &mut JITState, ctx: &Context, ocb: &mut OutlinedCb, counter: Option<Counter>) -> Target {
-    let side_exit = side_exit(jit, ctx, ocb);
-
+pub fn gen_counted_exit(side_exit: CodePtr, ocb: &mut OutlinedCb, counter: Option<Counter>) -> CodePtr {
     // The counter is only incremented when stats are enabled
     if !get_option!(gen_stats) {
         return side_exit;
@@ -572,16 +532,15 @@ fn counted_exit(jit: &mut JITState, ctx: &Context, ocb: &mut OutlinedCb, counter
     asm.incr_counter(counter_opnd, Opnd::UImm(1));
 
     // Jump to the existing side exit
-    asm.jmp(side_exit);
-    asm.compile(ocb);
+    asm.jmp(Target::CodePtr(side_exit));
+    asm.compile(ocb, None);
 
-    // Pointer to the side-exit code
-    code_ptr.as_side_exit()
+    code_ptr
 }
 
 // Ensure that there is an exit for the start of the block being compiled.
 // Block invalidation uses this exit.
-pub fn jit_ensure_block_entry_exit(jit: &mut JITState, ocb: &mut OutlinedCb) {
+pub fn jit_ensure_block_entry_exit(jit: &mut JITState, asm: &mut Assembler, ocb: &mut OutlinedCb) {
     if jit.block_entry_exit.is_some() {
         return;
     }
@@ -590,8 +549,9 @@ pub fn jit_ensure_block_entry_exit(jit: &mut JITState, ocb: &mut OutlinedCb) {
 
     // If we're compiling the first instruction in the block.
     if jit.insn_idx == jit.starting_insn_idx {
-        // Generate the exit with the cache in jitstate.
-        let entry_exit = side_exit(jit, block_starting_context, ocb).unwrap_code_ptr();
+        // Generate the exit with the cache in Assembler.
+        let side_exit_context = SideExitContext { pc: jit.pc, ctx: block_starting_context.clone() };
+        let entry_exit = asm.get_side_exit(&side_exit_context, None, ocb);
         jit.block_entry_exit = Some(entry_exit);
     } else {
         let block_entry_pc = unsafe { rb_iseq_pc_at_idx(jit.iseq, jit.starting_insn_idx.into()) };
@@ -626,7 +586,7 @@ fn gen_full_cfunc_return(ocb: &mut OutlinedCb) -> CodePtr {
 
     asm.cret(Qundef.into());
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     return code_ptr;
 }
@@ -654,7 +614,7 @@ fn gen_leave_exit(ocb: &mut OutlinedCb) -> CodePtr {
 
     asm.cret(ret_opnd);
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     return code_ptr;
 }
@@ -732,7 +692,7 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, ocb: &mut OutlinedCb, iseq: IseqPt
         None
     };
 
-    asm.compile(cb);
+    asm.compile(cb, Some(ocb));
 
     if cb.has_dropped_bytes() {
         None
@@ -755,9 +715,7 @@ pub fn gen_entry_prologue(cb: &mut CodeBlock, ocb: &mut OutlinedCb, iseq: IseqPt
 // Generate code to check for interrupts and take a side-exit.
 // Warning: this function clobbers REG0
 fn gen_check_ints(
-    jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     counter: Option<Counter>,
 ) {
     // Check for interrupts
@@ -769,7 +727,7 @@ fn gen_check_ints(
     let interrupt_flag = asm.load(Opnd::mem(32, EC, RUBY_OFFSET_EC_INTERRUPT_FLAG));
     asm.test(interrupt_flag, interrupt_flag);
 
-    asm.jnz(counted_exit(jit, &asm.ctx, ocb, counter));
+    asm.jnz(Target::side_exit(counter));
 }
 
 // Generate a stubbed unconditional jump to the next bytecode instruction.
@@ -872,7 +830,7 @@ pub fn gen_single_block(
         jit.opcode = opcode;
         jit.pc = pc;
         jit.stack_size_for_pc = asm.ctx.get_stack_size();
-        jit.side_exit_for_pc.clear();
+        asm.set_side_exit_context(pc, asm.ctx.get_stack_size());
 
         // stack_pop doesn't immediately deallocate a register for stack temps,
         // but it's safe to do so at this instruction boundary.
@@ -961,7 +919,7 @@ pub fn gen_single_block(
     }
 
     // Compile code into the code block
-    let gc_offsets = asm.compile(cb);
+    let gc_offsets = asm.compile(cb, Some(ocb));
     let end_addr = cb.get_write_ptr();
 
     // If code for the block doesn't fit, fail
@@ -1203,7 +1161,7 @@ fn gen_opt_plus(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_PLUS) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_PLUS) {
             return None;
         }
 
@@ -1217,7 +1175,7 @@ fn gen_opt_plus(
         // Add arg0 + arg1 and test for overflow
         let arg0_untag = asm.sub(arg0, Opnd::Imm(1));
         let out_val = asm.add(arg0_untag, arg1);
-        asm.jo(side_exit(jit, &asm.ctx, ocb));
+        asm.jo(Target::side_exit(None));
 
         // Push the output on the stack
         let dst = asm.stack_push(Type::Fixnum);
@@ -1386,9 +1344,7 @@ fn gen_newrange(
 }
 
 fn guard_object_is_heap(
-    jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     object: Opnd,
     object_opnd: YARVOpnd,
     counter: Option<Counter>,
@@ -1402,11 +1358,11 @@ fn guard_object_is_heap(
 
     // Test that the object is not an immediate
     asm.test(object, (RUBY_IMMEDIATE_MASK as u64).into());
-    asm.jnz(counted_exit(jit, &asm.ctx, ocb, counter));
+    asm.jnz(Target::side_exit(counter));
 
     // Test that the object is not false
     asm.cmp(object, Qfalse.into());
-    asm.je(counted_exit(jit, &asm.ctx, ocb, counter));
+    asm.je(Target::side_exit(counter));
 
     if object_type.diff(Type::UnknownHeap) != TypeDiff::Incompatible {
         asm.ctx.upgrade_opnd_type(object_opnd, Type::UnknownHeap);
@@ -1414,9 +1370,7 @@ fn guard_object_is_heap(
 }
 
 fn guard_object_is_array(
-    jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     object: Opnd,
     object_opnd: YARVOpnd,
     counter: Option<Counter>,
@@ -1430,7 +1384,7 @@ fn guard_object_is_array(
         Opnd::Reg(_) => object,
         _ => asm.load(object),
     };
-    guard_object_is_heap(jit, asm, ocb, object_reg, object_opnd, counter);
+    guard_object_is_heap(asm, object_reg, object_opnd, counter);
 
     asm.comment("guard object is array");
 
@@ -1440,7 +1394,7 @@ fn guard_object_is_array(
 
     // Compare the result with T_ARRAY
     asm.cmp(flags_opnd, (RUBY_T_ARRAY as u64).into());
-    asm.jne(counted_exit(jit, &asm.ctx, ocb, counter));
+    asm.jne(Target::side_exit(counter));
 
     if object_type.diff(Type::TArray) != TypeDiff::Incompatible {
         asm.ctx.upgrade_opnd_type(object_opnd, Type::TArray);
@@ -1448,9 +1402,7 @@ fn guard_object_is_array(
 }
 
 fn guard_object_is_string(
-    jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     object: Opnd,
     object_opnd: YARVOpnd,
     counter: Option<Counter>,
@@ -1464,7 +1416,7 @@ fn guard_object_is_string(
         Opnd::Reg(_) => object,
         _ => asm.load(object),
     };
-    guard_object_is_heap(jit, asm, ocb, object_reg, object_opnd, counter);
+    guard_object_is_heap(asm, object_reg, object_opnd, counter);
 
     asm.comment("guard object is string");
 
@@ -1474,7 +1426,7 @@ fn guard_object_is_string(
 
     // Compare the result with T_STRING
     asm.cmp(flags_reg, Opnd::UImm(RUBY_T_STRING as u64));
-    asm.jne(counted_exit(jit, &asm.ctx, ocb, counter));
+    asm.jne(Target::side_exit(counter));
 
     if object_type.diff(Type::TString) != TypeDiff::Incompatible {
         asm.ctx.upgrade_opnd_type(object_opnd, Type::TString);
@@ -1488,7 +1440,7 @@ fn guard_object_is_string(
 fn guard_object_is_not_ruby2_keyword_hash(
     asm: &mut Assembler,
     object_opnd: Opnd,
-    side_exit: Target,
+    counter: Option<Counter>,
 ) {
     asm.comment("guard object is not ruby2 keyword hash");
 
@@ -1510,7 +1462,7 @@ fn guard_object_is_not_ruby2_keyword_hash(
     asm.jne(not_ruby2_keyword);
 
     asm.test(flags_opnd, (RHASH_PASS_AS_KEYWORDS as u64).into());
-    asm.jnz(side_exit);
+    asm.jnz(Target::side_exit(counter));
 
     asm.write_label(not_ruby2_keyword);
 }
@@ -1519,7 +1471,7 @@ fn guard_object_is_not_ruby2_keyword_hash(
 fn gen_expandarray(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     // Both arguments are rb_num_t which is unsigned
     let num = jit.get_arg(0).as_usize();
@@ -1554,9 +1506,7 @@ fn gen_expandarray(
 
     // Move the array from the stack and check that it's an array.
     guard_object_is_array(
-        jit,
         asm,
-        ocb,
         array_opnd,
         array_opnd.into(),
         Some(Counter::expandarray_not_array),
@@ -1574,7 +1524,7 @@ fn gen_expandarray(
     // Only handle the case where the number of values in the array is greater
     // than or equal to the number of values requested.
     asm.cmp(array_len_opnd, num.into());
-    asm.jl(counted_exit!(jit, &asm.ctx, ocb, expandarray_rhs_too_small));
+    asm.jl(Target::side_exit(Some(Counter::expandarray_rhs_too_small)));
 
     // Load the address of the embedded array into REG1.
     // (struct RArray *)(obj)->as.ary
@@ -1720,7 +1670,6 @@ fn gen_getlocal_wc1(
 fn gen_setlocal_generic(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     ep_offset: u32,
     level: u32,
 ) -> Option<CodegenStatus> {
@@ -1742,7 +1691,7 @@ fn gen_setlocal_generic(
         asm.test(flags_opnd, VM_ENV_FLAG_WB_REQUIRED.into());
 
         // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
-        asm.jnz(side_exit(jit, &asm.ctx, ocb));
+        asm.jnz(Target::side_exit(None));
     }
 
     if level == 0 {
@@ -1763,29 +1712,29 @@ fn gen_setlocal_generic(
 fn gen_setlocal(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     let idx = jit.get_arg(0).as_u32();
     let level = jit.get_arg(1).as_u32();
-    gen_setlocal_generic(jit, asm, ocb, idx, level)
+    gen_setlocal_generic(jit, asm, idx, level)
 }
 
 fn gen_setlocal_wc0(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     let idx = jit.get_arg(0).as_u32();
-    gen_setlocal_generic(jit, asm, ocb, idx, 0)
+    gen_setlocal_generic(jit, asm, idx, 0)
 }
 
 fn gen_setlocal_wc1(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     let idx = jit.get_arg(0).as_u32();
-    gen_setlocal_generic(jit, asm, ocb, idx, 1)
+    gen_setlocal_generic(jit, asm, idx, 1)
 }
 
 // new hash initialized from top N values
@@ -1925,7 +1874,7 @@ fn jit_chain_guard(
 
         gen_branch(jit, asm, ocb, bid, &deeper, None, None, target0_gen_fn);
     } else {
-        target0_gen_fn.call(asm, counted_exit(jit, &asm.ctx, ocb, counter).unwrap_code_ptr(), None);
+        target0_gen_fn.call(asm, Target::side_exit(counter), None);
     }
 }
 
@@ -2069,7 +2018,7 @@ fn gen_get_ivar(
     };
 
     // Guard heap object (recv_opnd must be used before stack_pop)
-    guard_object_is_heap(jit, asm, ocb, recv, recv_opnd, None);
+    guard_object_is_heap(asm, recv, recv_opnd, None);
 
     // Compile time self is embedded and the ivar index lands within the object
     let embed_test_result = unsafe { FL_TEST_RAW(comptime_receiver, VALUE(ROBJECT_EMBED.as_usize())) != VALUE(0) };
@@ -2281,7 +2230,7 @@ fn gen_setinstancevariable(
         let recv_opnd = SelfOpnd;
 
         // Upgrade type
-        guard_object_is_heap(jit, asm, ocb, recv, recv_opnd, None);
+        guard_object_is_heap(asm, recv, recv_opnd, None);
 
         let expected_shape = unsafe { rb_shape_get_shape_id(comptime_receiver) };
         let shape_id_offset = unsafe { rb_shape_id_offset() };
@@ -2496,7 +2445,7 @@ fn gen_definedivar(
     };
 
     // Guard heap object (recv_opnd must be used before stack_pop)
-    guard_object_is_heap(jit, asm, ocb, recv, SelfOpnd, None);
+    guard_object_is_heap(asm, recv, SelfOpnd, None);
 
     let shape_id_offset = unsafe { rb_shape_id_offset() };
     let shape_opnd = Opnd::mem(SHAPE_ID_NUM_BITS as u8, recv, shape_id_offset);
@@ -2616,19 +2565,19 @@ fn guard_two_fixnums(
 
     if arg0_type.is_heap() || arg1_type.is_heap() {
         asm.comment("arg is heap object");
-        asm.jmp(side_exit(jit, &asm.ctx, ocb));
+        asm.jmp(Target::side_exit(None));
         return;
     }
 
     if arg0_type != Type::Fixnum && arg0_type.is_specific() {
         asm.comment("arg0 not fixnum");
-        asm.jmp(side_exit(jit, &asm.ctx, ocb));
+        asm.jmp(Target::side_exit(None));
         return;
     }
 
     if arg1_type != Type::Fixnum && arg1_type.is_specific() {
         asm.comment("arg1 not fixnum");
-        asm.jmp(side_exit(jit, &asm.ctx, ocb));
+        asm.jmp(Target::side_exit(None));
         return;
     }
 
@@ -2690,7 +2639,7 @@ fn gen_fixnum_cmp(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, bop) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, bop) {
             return None;
         }
 
@@ -2765,7 +2714,7 @@ fn gen_equality_specialized(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_EQ) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_EQ) {
             // if overridden, emit the generic version
             return Some(false);
         }
@@ -2794,7 +2743,7 @@ fn gen_equality_specialized(
     let comptime_b = jit.peek_at_stack(&asm.ctx, 0);
 
     if unsafe { comptime_a.class_of() == rb_cString && comptime_b.class_of() == rb_cString } {
-        if !assume_bop_not_redefined(jit, ocb, STRING_REDEFINED_OP_FLAG, BOP_EQ) {
+        if !assume_bop_not_redefined(jit, asm, ocb, STRING_REDEFINED_OP_FLAG, BOP_EQ) {
             // if overridden, emit the generic version
             return Some(false);
         }
@@ -2921,7 +2870,7 @@ fn gen_opt_aref(
     let comptime_recv = jit.peek_at_stack(&asm.ctx, 1);
 
     if comptime_recv.class_of() == unsafe { rb_cArray } && comptime_idx.fixnum_p() {
-        if !assume_bop_not_redefined(jit, ocb, ARRAY_REDEFINED_OP_FLAG, BOP_AREF) {
+        if !assume_bop_not_redefined(jit, asm, ocb, ARRAY_REDEFINED_OP_FLAG, BOP_AREF) {
             return None;
         }
 
@@ -2946,7 +2895,7 @@ fn gen_opt_aref(
         // Bail if idx is not a FIXNUM
         let idx_reg = asm.load(idx_opnd);
         asm.test(idx_reg, (RUBY_FIXNUM_FLAG as u64).into());
-        asm.jz(counted_exit!(jit, &asm.ctx, ocb, oaref_arg_not_fixnum));
+        asm.jz(Target::side_exit(Some(Counter::oaref_arg_not_fixnum)));
 
         // Call VALUE rb_ary_entry_internal(VALUE ary, long offset).
         // It never raises or allocates, so we don't need to write to cfp->pc.
@@ -2967,7 +2916,7 @@ fn gen_opt_aref(
         jump_to_next_insn(jit, asm, ocb);
         return Some(EndBlock);
     } else if comptime_recv.class_of() == unsafe { rb_cHash } {
-        if !assume_bop_not_redefined(jit, ocb, HASH_REDEFINED_OP_FLAG, BOP_AREF) {
+        if !assume_bop_not_redefined(jit, asm, ocb, HASH_REDEFINED_OP_FLAG, BOP_AREF) {
             return None;
         }
 
@@ -3127,7 +3076,7 @@ fn gen_opt_and(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_AND) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_AND) {
             return None;
         }
 
@@ -3167,7 +3116,7 @@ fn gen_opt_or(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_OR) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_OR) {
             return None;
         }
 
@@ -3207,7 +3156,7 @@ fn gen_opt_minus(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MINUS) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MINUS) {
             return None;
         }
 
@@ -3220,7 +3169,7 @@ fn gen_opt_minus(
 
         // Subtract arg0 - arg1 and test for overflow
         let val_untag = asm.sub(arg0, arg1);
-        asm.jo(side_exit(jit, &asm.ctx, ocb));
+        asm.jo(Target::side_exit(None));
         let val = asm.add(val_untag, Opnd::Imm(1));
 
         // Push the output on the stack
@@ -3267,7 +3216,7 @@ fn gen_opt_mod(
     };
 
     if two_fixnums {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MOD) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_MOD) {
             return None;
         }
 
@@ -3281,7 +3230,7 @@ fn gen_opt_mod(
 
         // Check for arg0 % 0
         asm.cmp(arg1, Opnd::Imm(VALUE::fixnum_from_usize(0).as_i64()));
-        asm.je(side_exit(jit, &asm.ctx, ocb));
+        asm.je(Target::side_exit(None));
 
         // Call rb_fix_mod_fix(VALUE recv, VALUE obj)
         let ret = asm.ccall(rb_fix_mod_fix as *const u8, vec![arg0, arg1]);
@@ -3339,7 +3288,7 @@ fn gen_opt_str_freeze(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
-    if !assume_bop_not_redefined(jit, ocb, STRING_REDEFINED_OP_FLAG, BOP_FREEZE) {
+    if !assume_bop_not_redefined(jit, asm, ocb, STRING_REDEFINED_OP_FLAG, BOP_FREEZE) {
         return None;
     }
 
@@ -3357,7 +3306,7 @@ fn gen_opt_str_uminus(
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
-    if !assume_bop_not_redefined(jit, ocb, STRING_REDEFINED_OP_FLAG, BOP_UMINUS) {
+    if !assume_bop_not_redefined(jit, asm, ocb, STRING_REDEFINED_OP_FLAG, BOP_UMINUS) {
         return None;
     }
 
@@ -3516,7 +3465,7 @@ fn gen_opt_case_dispatch(
     }
 
     if comptime_key.fixnum_p() && comptime_key.0 <= u32::MAX.as_usize() && case_hash_all_fixnum_p(case_hash) {
-        if !assume_bop_not_redefined(jit, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_EQQ) {
+        if !assume_bop_not_redefined(jit, asm, ocb, INTEGER_REDEFINED_OP_FLAG, BOP_EQQ) {
             return None;
         }
 
@@ -3562,7 +3511,7 @@ fn gen_branchif(
 
     // Check for interrupts, but only on backward branches that may create loops
     if jump_offset < 0 {
-        gen_check_ints(jit, asm, ocb, None);
+        gen_check_ints(asm, None);
     }
 
     // Get the branch target instruction offsets
@@ -3617,7 +3566,7 @@ fn gen_branchunless(
 
     // Check for interrupts, but only on backward branches that may create loops
     if jump_offset < 0 {
-        gen_check_ints(jit, asm, ocb, None);
+        gen_check_ints(asm, None);
     }
 
     // Get the branch target instruction offsets
@@ -3673,7 +3622,7 @@ fn gen_branchnil(
 
     // Check for interrupts, but only on backward branches that may create loops
     if jump_offset < 0 {
-        gen_check_ints(jit, asm, ocb, None);
+        gen_check_ints(asm, None);
     }
 
     // Get the branch target instruction offsets
@@ -3754,13 +3703,13 @@ fn gen_throw(
 fn gen_jump(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     let jump_offset = jit.get_arg(0).as_i32();
 
     // Check for interrupts, but only on backward branches that may create loops
     if jump_offset < 0 {
-        gen_check_ints(jit, asm, ocb, None);
+        gen_check_ints(asm, None);
     }
 
     // Get the branch target instruction offsets
@@ -3923,9 +3872,7 @@ fn jit_guard_known_klass(
 // Generate ancestry guard for protected callee.
 // Calls to protected callees only go through when self.is_a?(klass_that_defines_the_callee).
 fn jit_protected_callee_ancestry_guard(
-    jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
     cme: *const rb_callable_method_entry_t,
 ) {
     // See vm_call_method().
@@ -3942,7 +3889,7 @@ fn jit_protected_callee_ancestry_guard(
         ],
     );
     asm.test(val, val);
-    asm.jz(counted_exit!(jit, &asm.ctx, ocb, send_se_protected_check_failed))
+    asm.jz(Target::side_exit(Some(Counter::send_se_protected_check_failed)))
 }
 
 // Codegen for rb_obj_not().
@@ -4022,7 +3969,7 @@ fn jit_rb_false(
 fn jit_rb_kernel_is_a(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
     _ci: *const rb_callinfo,
     _cme: *const rb_callable_method_entry_t,
     _block: Option<IseqPtr>,
@@ -4057,7 +4004,7 @@ fn jit_rb_kernel_is_a(
 
     asm.comment("Kernel#is_a?");
     asm.cmp(asm.stack_opnd(0), sample_rhs.into());
-    asm.jne(counted_exit!(jit, &asm.ctx, ocb, send_is_a_class_mismatch));
+    asm.jne(Target::side_exit(Some(Counter::send_is_a_class_mismatch)));
 
     asm.stack_pop(2);
 
@@ -4075,7 +4022,7 @@ fn jit_rb_kernel_is_a(
 fn jit_rb_kernel_instance_of(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
     _ci: *const rb_callinfo,
     _cme: *const rb_callable_method_entry_t,
     _block: Option<IseqPtr>,
@@ -4116,7 +4063,7 @@ fn jit_rb_kernel_instance_of(
 
     asm.comment("Kernel#instance_of?");
     asm.cmp(asm.stack_opnd(0), sample_rhs.into());
-    asm.jne(counted_exit!(jit, &asm.ctx, ocb, send_instance_of_class_mismatch));
+    asm.jne(Target::side_exit(Some(Counter::send_instance_of_class_mismatch)));
 
     asm.stack_pop(2);
 
@@ -4279,7 +4226,7 @@ fn jit_rb_int_div(
 
     // Check for arg0 % 0
     asm.cmp(obj, VALUE::fixnum_from_usize(0).as_i64().into());
-    asm.je(side_exit(jit, &asm.ctx, ocb));
+    asm.je(Target::side_exit(None));
 
     let ret = asm.ccall(rb_fix_div_fix as *const u8, vec![recv, obj]);
 
@@ -4446,7 +4393,7 @@ fn jit_rb_str_empty_p(
 fn jit_rb_str_concat(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
     _ci: *const rb_callinfo,
     _cme: *const rb_callable_method_entry_t,
     _block: Option<IseqPtr>,
@@ -4463,7 +4410,7 @@ fn jit_rb_str_concat(
     }
 
     // Guard that the concat argument is a string
-    guard_object_is_string(jit, asm, ocb, asm.stack_opnd(0), StackOpnd(0), None);
+    guard_object_is_string(asm, asm.stack_opnd(0), StackOpnd(0), None);
 
     // Guard buffers from GC since rb_str_buf_append may allocate. During the VM lock on GC,
     // other Ractors may trigger global invalidation, so we need ctx.clear_local_types().
@@ -4625,7 +4572,7 @@ fn jit_obj_respond_to(
     if result != Qtrue {
         // Only if respond_to_missing? hasn't been overridden
         // In the future, we might want to jit the call to respond_to_missing?
-        if !assume_method_basic_definition(jit, ocb, recv_class, idRespond_to_missing.into()) {
+        if !assume_method_basic_definition(jit, asm, ocb, recv_class, idRespond_to_missing.into()) {
             return false;
         }
     }
@@ -4633,7 +4580,7 @@ fn jit_obj_respond_to(
     // Invalidate this block if method lookup changes for the method being queried. This works
     // both for the case where a method does or does not exist, as for the latter we asked for a
     // "negative CME" earlier.
-    jit.assume_method_lookup_stable(ocb, target_cme);
+    jit.assume_method_lookup_stable(asm, ocb, target_cme);
 
     if argc == 2 {
         // pop include_all argument (we only use its type info)
@@ -4646,7 +4593,7 @@ fn jit_obj_respond_to(
     // This is necessary because we have no guarantee that sym_opnd is a constant
     asm.comment("guard known mid");
     asm.cmp(sym_opnd, mid_sym.into());
-    asm.jne(side_exit(jit, &asm.ctx, ocb));
+    asm.jne(Target::side_exit(None));
 
     jit_putobject(asm, result);
 
@@ -4988,7 +4935,7 @@ fn gen_send_cfunc(
     }
 
     // Check for interrupts
-    gen_check_ints(jit, asm, ocb, None);
+    gen_check_ints(asm, None);
 
     // Stack overflow check
     // #define CHECK_VM_STACK_OVERFLOW0(cfp, sp, margin)
@@ -4996,7 +4943,7 @@ fn gen_send_cfunc(
     asm.comment("stack overflow check");
     let stack_limit = asm.lea(asm.ctx.sp_opnd((SIZEOF_VALUE * 4 + 2 * RUBY_SIZEOF_CONTROL_FRAME) as isize));
     asm.cmp(CFP, stack_limit);
-    asm.jbe(counted_exit!(jit, &asm.ctx, ocb, send_se_cf_overflow));
+    asm.jbe(Target::side_exit(Some(Counter::send_se_cf_overflow)));
 
     // Number of args which will be passed through to the callee
     // This is adjusted by the kwargs being combined into a hash.
@@ -5073,7 +5020,7 @@ fn gen_send_cfunc(
         // and if not side exit.
         argc = cfunc_argc;
         passed_argc = argc;
-        push_splat_args(required_args, jit, asm, ocb)
+        push_splat_args(required_args, asm)
     }
 
     // This is a .send call and we need to adjust the stack
@@ -5254,14 +5201,14 @@ fn get_array_ptr(asm: &mut Assembler, array_reg: Opnd) -> Opnd {
 
 /// Pushes arguments from an array to the stack. Differs from push splat because
 /// the array can have items left over.
-fn move_rest_args_to_stack(array: Opnd, num_args: u32, jit: &mut JITState, asm: &mut Assembler, ocb: &mut OutlinedCb) {
+fn move_rest_args_to_stack(array: Opnd, num_args: u32, asm: &mut Assembler) {
     asm.comment("move_rest_args_to_stack");
 
     let array_len_opnd = get_array_len(asm, array);
 
     asm.comment("Side exit if length is less than required");
     asm.cmp(array_len_opnd, num_args.into());
-    asm.jl(counted_exit!(jit, &asm.ctx, ocb, send_iseq_has_rest_and_splat_not_equal));
+    asm.jl(Target::side_exit(Some(Counter::send_iseq_has_rest_and_splat_not_equal)));
 
     asm.comment("Push arguments from array");
 
@@ -5292,16 +5239,14 @@ fn move_rest_args_to_stack(array: Opnd, num_args: u32, jit: &mut JITState, asm:
 /// Pushes arguments from an array to the stack that are passed with a splat (i.e. *args)
 /// It optimistically compiles to a static size that is the exact number of arguments
 /// needed for the function.
-fn push_splat_args(required_args: u32, jit: &mut JITState, asm: &mut Assembler, ocb: &mut OutlinedCb) {
+fn push_splat_args(required_args: u32, asm: &mut Assembler) {
     asm.comment("push_splat_args");
 
     let array_opnd = asm.stack_opnd(0);
     let array_reg = asm.load(array_opnd);
 
     guard_object_is_array(
-        jit,
         asm,
-        ocb,
         array_reg,
         array_opnd.into(),
         Some(Counter::send_splat_not_array),
@@ -5333,7 +5278,7 @@ fn push_splat_args(required_args: u32, jit: &mut JITState, asm: &mut Assembler,
 
     asm.comment("Side exit if length doesn't not equal remaining args");
     asm.cmp(array_len_opnd, required_args.into());
-    asm.jne(counted_exit!(jit, &asm.ctx, ocb, send_splatarray_length_not_equal));
+    asm.jne(Target::side_exit(Some(Counter::send_splatarray_length_not_equal)));
 
     asm.comment("Check last argument is not ruby2keyword hash");
 
@@ -5347,7 +5292,7 @@ fn push_splat_args(required_args: u32, jit: &mut JITState, asm: &mut Assembler,
     guard_object_is_not_ruby2_keyword_hash(
         asm,
         last_array_value,
-        counted_exit!(jit, &asm.ctx, ocb, send_splatarray_last_ruby_2_keywords),
+        Some(Counter::send_splatarray_last_ruby_2_keywords),
     );
 
     asm.comment("Push arguments from array");
@@ -5405,7 +5350,7 @@ fn gen_send_bmethod(
 
     // Optimize for single ractor mode and avoid runtime check for
     // "defined with an un-shareable Proc in a different Ractor"
-    if !assume_single_ractor_mode(jit, ocb) {
+    if !assume_single_ractor_mode(jit, asm, ocb) {
         gen_counter_incr!(asm, send_bmethod_ractor);
         return None;
     }
@@ -5737,8 +5682,7 @@ fn gen_send_iseq(
             asm.comment("Side exit if length doesn't not equal compile time length");
             let array_len_opnd = get_array_len(asm, asm.stack_opnd(if block_arg { 1 } else { 0 }));
             asm.cmp(array_len_opnd, array_length.into());
-            let exit = counted_exit!(jit, &mut asm.ctx, ocb, send_splatarray_length_not_equal);
-            asm.jne(exit);
+            asm.jne(Target::side_exit(Some(Counter::send_splatarray_length_not_equal)));
         }
 
         Some(array_length)
@@ -5852,7 +5796,7 @@ fn gen_send_iseq(
         SIZEOF_VALUE_I32 * (num_locals + stack_max) + 2 * (RUBY_SIZEOF_CONTROL_FRAME as i32);
     let stack_limit = asm.lea(asm.ctx.sp_opnd(locals_offs as isize));
     asm.cmp(CFP, stack_limit);
-    asm.jbe(counted_exit!(jit, &asm.ctx, ocb, send_se_cf_overflow));
+    asm.jbe(Target::side_exit(Some(Counter::send_se_cf_overflow)));
 
     // push_splat_args does stack manipulation so we can no longer side exit
     if let Some(array_length) = splat_array_length {
@@ -5873,7 +5817,7 @@ fn gen_send_iseq(
             // all the remaining arguments. In the generated code
             // we test if this is true and if not side exit.
             argc = argc - 1 + array_length as i32 + remaining_opt as i32;
-            push_splat_args(array_length, jit, asm, ocb);
+            push_splat_args(array_length, asm);
 
             for _ in 0..remaining_opt {
                 // We need to push nil for the optional arguments
@@ -5932,7 +5876,7 @@ fn gen_send_iseq(
 
                 let diff = (required_num - non_rest_arg_count + opts_filled_with_splat.unwrap_or(0)) as u32;
                 // This moves the arguments onto the stack. But it doesn't modify the array.
-                move_rest_args_to_stack(array, diff, jit, asm, ocb);
+                move_rest_args_to_stack(array, diff, asm);
 
                 // We will now slice the array to give us a new array of the correct size
                 asm.spill_temps(); // for ccall
@@ -6125,13 +6069,13 @@ fn gen_send_iseq(
 
         // Only handle the case that you don't need to_ary conversion
         let not_array_counter = Some(Counter::invokeblock_iseq_arg0_not_array);
-        guard_object_is_array(jit, asm, ocb, arg0_opnd, arg0_opnd.into(), not_array_counter);
+        guard_object_is_array(asm, arg0_opnd, arg0_opnd.into(), not_array_counter);
 
         // Only handle the same that the array length == ISEQ's lead_num (most common)
         let arg0_len_opnd = get_array_len(asm, arg0_opnd);
         let lead_num = unsafe { rb_get_iseq_body_param_lead_num(iseq) };
         asm.cmp(arg0_len_opnd, lead_num.into());
-        asm.jne(counted_exit!(jit, &asm.ctx, ocb, invokeblock_iseq_arg0_wrong_len));
+        asm.jne(Target::side_exit(Some(Counter::invokeblock_iseq_arg0_wrong_len)));
 
         let arg0_reg = asm.load(arg0_opnd);
         let array_opnd = get_array_ptr(asm, arg0_reg);
@@ -6481,7 +6425,7 @@ fn gen_send_general(
             if flags & VM_CALL_FCALL == 0 {
                 // otherwise we need an ancestry check to ensure the receiver is valid to be called
                 // as protected
-                jit_protected_callee_ancestry_guard(jit, asm, ocb, cme);
+                jit_protected_callee_ancestry_guard(asm, cme);
             }
         }
         _ => {
@@ -6491,7 +6435,7 @@ fn gen_send_general(
 
     // Register block for invalidation
     //assert!(cme->called_id == mid);
-    jit.assume_method_lookup_stable(ocb, cme);
+    jit.assume_method_lookup_stable(asm, ocb, cme);
 
     // To handle the aliased method case (VM_METHOD_TYPE_ALIAS)
     loop {
@@ -6656,7 +6600,7 @@ fn gen_send_general(
 
                         flags |= VM_CALL_FCALL | VM_CALL_OPT_SEND;
 
-                        jit.assume_method_lookup_stable(ocb, cme);
+                        jit.assume_method_lookup_stable(asm, ocb, cme);
 
                         let (known_class, type_mismatch_counter) = {
                             if compile_time_name.string_p() {
@@ -6727,7 +6671,7 @@ fn gen_send_general(
 
                         // Optimize for single ractor mode and avoid runtime check for
                         // "defined with an un-shareable Proc in a different Ractor"
-                        if !assume_single_ractor_mode(jit, ocb) {
+                        if !assume_single_ractor_mode(jit, asm, ocb) {
                             gen_counter_incr!(asm, send_call_multi_ractor);
                             return None;
                         }
@@ -7109,7 +7053,7 @@ fn gen_invokesuper(
 
     let me_as_value = VALUE(me as usize);
     asm.cmp(ep_me_opnd, me_as_value.into());
-    asm.jne(counted_exit!(jit, &asm.ctx, ocb, invokesuper_me_changed));
+    asm.jne(Target::side_exit(Some(Counter::invokesuper_me_changed)));
 
     if block.is_none() {
         // Guard no block passed
@@ -7125,13 +7069,13 @@ fn gen_invokesuper(
             SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL,
         );
         asm.cmp(ep_specval_opnd, VM_BLOCK_HANDLER_NONE.into());
-        asm.jne(counted_exit!(jit, &asm.ctx, ocb, invokesuper_block));
+        asm.jne(Target::side_exit(Some(Counter::invokesuper_block)));
     }
 
     // We need to assume that both our current method entry and the super
     // method entry we invoke remain stable
-    jit.assume_method_lookup_stable(ocb, me);
-    jit.assume_method_lookup_stable(ocb, cme);
+    jit.assume_method_lookup_stable(asm, ocb, me);
+    jit.assume_method_lookup_stable(asm, ocb, cme);
 
     // Method calls may corrupt types
     asm.ctx.clear_local_types();
@@ -7150,7 +7094,7 @@ fn gen_invokesuper(
 }
 
 fn gen_leave(
-    jit: &mut JITState,
+    _jit: &mut JITState,
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
@@ -7160,8 +7104,8 @@ fn gen_leave(
     let ocb_asm = Assembler::new();
 
     // Check for interrupts
-    gen_check_ints(jit, asm, ocb, Some(Counter::leave_se_interrupt));
-    ocb_asm.compile(ocb.unwrap());
+    gen_check_ints(asm, Some(Counter::leave_se_interrupt));
+    ocb_asm.compile(ocb.unwrap(), None);
 
     // Pop the current frame (ec->cfp++)
     // Note: the return PC is already in the previous CFP
@@ -7533,7 +7477,7 @@ fn gen_opt_getconstant_path(
 
     // Make sure there is an exit for this block as the interpreter might want
     // to invalidate this block from yjit_constant_ic_update().
-    jit_ensure_block_entry_exit(jit, ocb);
+    jit_ensure_block_entry_exit(jit, asm, ocb);
 
     if !unsafe { (*ice).ic_cref }.is_null() {
         // Cache is keyed on a certain lexical scope. Use the interpreter's cache.
@@ -7549,7 +7493,7 @@ fn gen_opt_getconstant_path(
         // Check the result. SysV only specifies one byte for _Bool return values,
         // so it's important we only check one bit to ignore the higher bits in the register.
         asm.test(ret_val, 1.into());
-        asm.jz(counted_exit!(jit, &asm.ctx, ocb, opt_getinlinecache_miss));
+        asm.jz(Target::side_exit(Some(Counter::opt_getinlinecache_miss)));
 
         let inline_cache = asm.load(Opnd::const_ptr(ic as *const u8));
 
@@ -7570,13 +7514,13 @@ fn gen_opt_getconstant_path(
         asm.store(stack_top, ic_entry_val);
     } else {
         // Optimize for single ractor mode.
-        if !assume_single_ractor_mode(jit, ocb) {
+        if !assume_single_ractor_mode(jit, asm, ocb) {
             return None;
         }
 
         // Invalidate output code on any constant writes associated with
         // constants referenced within the current block.
-        jit.assume_stable_constant_names(ocb, idlist);
+        jit.assume_stable_constant_names(asm, ocb, idlist);
 
         jit_putobject(asm, unsafe { (*ice).value });
     }
@@ -7623,7 +7567,7 @@ fn gen_getblockparamproxy(
         SIZEOF_VALUE_I32 * (VM_ENV_DATA_INDEX_FLAGS as i32),
     );
     asm.test(flag_check, VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM.into());
-    asm.jnz(counted_exit!(jit, &asm.ctx, ocb, gbpp_block_param_modified));
+    asm.jnz(Target::side_exit(Some(Counter::gbpp_block_param_modified)));
 
     // Load the block handler for the current frame
     // note, VM_ASSERT(VM_ENV_LOCAL_P(ep))
@@ -7677,7 +7621,7 @@ fn gen_getblockparamproxy(
 fn gen_getblockparam(
     jit: &mut JITState,
     asm: &mut Assembler,
-    ocb: &mut OutlinedCb,
+    _ocb: &mut OutlinedCb,
 ) -> Option<CodegenStatus> {
     // EP level
     let level = jit.get_arg(1).as_u32();
@@ -7714,7 +7658,7 @@ fn gen_getblockparam(
     asm.test(flags_opnd, VM_ENV_FLAG_WB_REQUIRED.into());
 
     // if (flags & VM_ENV_FLAG_WB_REQUIRED) != 0
-    asm.jnz(side_exit(jit, &asm.ctx, ocb));
+    asm.jnz(Target::side_exit(None));
 
     // Convert the block handler in to a proc
     // call rb_vm_bh_to_procval(const rb_execution_context_t *ec, VALUE block_handler)
@@ -8282,28 +8226,30 @@ mod tests {
     fn test_gen_exit() {
         let (_, _ctx, mut asm, mut cb, _) = setup_codegen();
         gen_exit(0 as *mut VALUE, &mut asm);
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
     #[test]
     fn test_get_side_exit() {
-        let (mut jit, ctx, _, _, mut ocb) = setup_codegen();
-        side_exit(&mut jit, &ctx, &mut ocb);
+        let (_jit, ctx, mut asm, _, mut ocb) = setup_codegen();
+        let side_exit_context = SideExitContext { pc: 0 as _, ctx };
+        asm.get_side_exit(&side_exit_context, None, &mut ocb);
         assert!(ocb.unwrap().get_write_pos() > 0);
     }
 
     #[test]
     fn test_gen_check_ints() {
-        let (mut jit, _ctx, mut asm, _cb, mut ocb) = setup_codegen();
-        gen_check_ints(&mut jit, &mut asm, &mut ocb, None);
+        let (_jit, _ctx, mut asm, _cb, _ocb) = setup_codegen();
+        asm.set_side_exit_context(0 as _, 0);
+        gen_check_ints(&mut asm, None);
     }
 
     #[test]
     fn test_gen_nop() {
         let (mut jit, context, mut asm, mut cb, mut ocb) = setup_codegen();
         let status = gen_nop(&mut jit, &mut asm, &mut ocb);
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(context.diff(&Context::default()), TypeDiff::Compatible(0));
@@ -8335,7 +8281,7 @@ mod tests {
         assert_eq!(Type::Fixnum, asm.ctx.get_opnd_type(StackOpnd(0)));
         assert_eq!(Type::Fixnum, asm.ctx.get_opnd_type(StackOpnd(1)));
 
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0); // Write some movs
     }
 
@@ -8359,7 +8305,7 @@ mod tests {
         assert_eq!(Type::Flonum, asm.ctx.get_opnd_type(StackOpnd(0)));
 
         // TODO: this is writing zero bytes on x86. Why?
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0); // Write some movs
     }
 
@@ -8388,7 +8334,7 @@ mod tests {
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::Nil);
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
@@ -8407,7 +8353,7 @@ mod tests {
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::True);
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
@@ -8427,7 +8373,7 @@ mod tests {
 
         assert_eq!(status, Some(KeepCompiling));
         assert_eq!(tmp_type_top, Type::Fixnum);
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
@@ -8450,7 +8396,7 @@ mod tests {
         let status = gen_putself(&mut jit, &mut asm, &mut ocb);
 
         assert_eq!(status, Some(KeepCompiling));
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
@@ -8473,7 +8419,7 @@ mod tests {
         assert_eq!(Type::Flonum, asm.ctx.get_opnd_type(StackOpnd(1)));
         assert_eq!(Type::CString, asm.ctx.get_opnd_type(StackOpnd(0)));
 
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0);
     }
 
@@ -8495,7 +8441,7 @@ mod tests {
         assert_eq!(Type::CString, asm.ctx.get_opnd_type(StackOpnd(1)));
         assert_eq!(Type::Flonum, asm.ctx.get_opnd_type(StackOpnd(0)));
 
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() > 0); // Write some movs
     }
 
@@ -8516,7 +8462,7 @@ mod tests {
 
         assert_eq!(Type::Flonum, asm.ctx.get_opnd_type(StackOpnd(0)));
 
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
         assert!(cb.get_write_pos() == 0); // No instructions written
     }
 
@@ -8525,6 +8471,7 @@ mod tests {
         let (mut jit, _context, mut asm, _cb, mut ocb) = setup_codegen();
         // Push return value
         asm.stack_push(Type::Fixnum);
+        asm.set_side_exit_context(0 as _, 0);
         gen_leave(&mut jit, &mut asm, &mut ocb);
     }
 }
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 377ab7c8bf..5362618a28 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -38,7 +38,7 @@ const MAX_LOCAL_TYPES: usize = 8;
 pub type IseqIdx = u16;
 
 // Represent the type of a value (local/stack/self) in YJIT
-#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+#[derive(Copy, Clone, Hash, PartialEq, Eq, Debug)]
 pub enum Type {
     Unknown,
     UnknownImm,
@@ -298,7 +298,7 @@ pub enum TypeDiff {
 
 // Potential mapping of a value on the temporary stack to
 // self, a local variable or constant so that we can track its type
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, Hash, PartialEq, Debug)]
 pub enum TempMapping {
     MapToStack, // Normal stack value
     MapToSelf,  // Temp maps to the self operand
@@ -307,7 +307,7 @@ pub enum TempMapping {
 }
 
 // Index used by MapToLocal. Using this instead of u8 makes TempMapping 1 byte.
-#[derive(Copy, Clone, Eq, PartialEq, Debug)]
+#[derive(Copy, Clone, Eq, Hash, PartialEq, Debug)]
 pub enum LocalIndex {
     Local0,
     Local1,
@@ -417,7 +417,7 @@ impl RegTemps {
 /// Code generation context
 /// Contains information we can use to specialize/optimize code
 /// There are a lot of context objects so we try to keep the size small.
-#[derive(Clone, Default, PartialEq, Debug)]
+#[derive(Clone, Copy, Default, Eq, Hash, PartialEq, Debug)]
 pub struct Context {
     // Number of values currently on the temporary stack
     stack_size: u8,
@@ -478,12 +478,12 @@ pub enum BranchGenFn {
 }
 
 impl BranchGenFn {
-    pub fn call(&self, asm: &mut Assembler, target0: CodePtr, target1: Option<CodePtr>) {
+    pub fn call(&self, asm: &mut Assembler, target0: Target, target1: Option<Target>) {
         match self {
             BranchGenFn::BranchIf(shape) => {
                 match shape.get() {
-                    BranchShape::Next0 => asm.jz(target1.unwrap().into()),
-                    BranchShape::Next1 => asm.jnz(target0.into()),
+                    BranchShape::Next0 => asm.jz(target1.unwrap()),
+                    BranchShape::Next1 => asm.jnz(target0),
                     BranchShape::Default => {
                         asm.jnz(target0.into());
                         asm.jmp(target1.unwrap().into());
@@ -492,21 +492,21 @@ impl BranchGenFn {
             }
             BranchGenFn::BranchNil(shape) => {
                 match shape.get() {
-                    BranchShape::Next0 => asm.jne(target1.unwrap().into()),
-                    BranchShape::Next1 => asm.je(target0.into()),
+                    BranchShape::Next0 => asm.jne(target1.unwrap()),
+                    BranchShape::Next1 => asm.je(target0),
                     BranchShape::Default => {
-                        asm.je(target0.into());
-                        asm.jmp(target1.unwrap().into());
+                        asm.je(target0);
+                        asm.jmp(target1.unwrap());
                     }
                 }
             }
             BranchGenFn::BranchUnless(shape) => {
                 match shape.get() {
-                    BranchShape::Next0 => asm.jnz(target1.unwrap().into()),
-                    BranchShape::Next1 => asm.jz(target0.into()),
+                    BranchShape::Next0 => asm.jnz(target1.unwrap()),
+                    BranchShape::Next1 => asm.jz(target0),
                     BranchShape::Default => {
-                        asm.jz(target0.into());
-                        asm.jmp(target1.unwrap().into());
+                        asm.jz(target0);
+                        asm.jmp(target1.unwrap());
                     }
                 }
             }
@@ -522,14 +522,14 @@ impl BranchGenFn {
                 asm.jnz(target0.into())
             }
             BranchGenFn::JZToTarget0 => {
-                asm.jz(Target::CodePtr(target0))
+                asm.jz(target0)
             }
             BranchGenFn::JBEToTarget0 => {
-                asm.jbe(Target::CodePtr(target0))
+                asm.jbe(target0)
             }
             BranchGenFn::JITReturn => {
                 asm.comment("update cfp->jit_return");
-                asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_JIT_RETURN), Opnd::const_ptr(target0.raw_ptr()));
+                asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_JIT_RETURN), Opnd::const_ptr(target0.unwrap_code_ptr().raw_ptr()));
             }
         }
     }
@@ -1378,10 +1378,7 @@ pub fn limit_block_versions(blockid: BlockId, ctx: &Context) -> Context {
         // Produce a generic context that stores no type information,
         // but still respects the stack_size and sp_offset constraints.
         // This new context will then match all future requests.
-        let mut generic_ctx = Context::default();
-        generic_ctx.stack_size = ctx.stack_size;
-        generic_ctx.sp_offset = ctx.sp_offset;
-        generic_ctx.reg_temps = ctx.reg_temps;
+        let generic_ctx = ctx.get_generic_ctx();
 
         debug_assert_ne!(
             TypeDiff::Incompatible,
@@ -1570,6 +1567,15 @@ impl Context {
         self.stack_size
     }
 
+    /// Create a new Context that is compatible with self but doesn't have type information.
+    pub fn get_generic_ctx(&self) -> Context {
+        let mut generic_ctx = Context::default();
+        generic_ctx.stack_size = self.stack_size;
+        generic_ctx.sp_offset = self.sp_offset;
+        generic_ctx.reg_temps = self.reg_temps;
+        generic_ctx
+    }
+
     /// Create a new Context instance with a given stack_size and sp_offset adjusted
     /// accordingly. This is useful when you want to virtually rewind a stack_size for
     /// generating a side exit while considering past sp_offset changes on gen_save_sp.
@@ -2170,7 +2176,7 @@ pub fn regenerate_entry(cb: &mut CodeBlock, entryref: &EntryRef, next_entry: Cod
     let old_dropped_bytes = cb.has_dropped_bytes();
     cb.set_write_ptr(unsafe { entryref.as_ref() }.start_addr);
     cb.set_dropped_bytes(false);
-    asm.compile(cb);
+    asm.compile(cb, None);
 
     // Rewind write_pos to the original one
     assert_eq!(cb.get_write_ptr(), unsafe { entryref.as_ref() }.end_addr);
@@ -2219,7 +2225,7 @@ fn entry_stub_hit_body(entry_ptr: *const c_void, ec: EcPtr) -> Option<*const u8>
     let next_entry = cb.get_write_ptr();
     let mut asm = Assembler::new();
     let pending_entry = gen_entry_chain_guard(&mut asm, ocb, iseq, insn_idx)?;
-    asm.compile(cb);
+    asm.compile(cb, Some(ocb));
 
     // Try to find an existing compiled version of this block
     let blockid = BlockId { iseq, idx: insn_idx };
@@ -2229,7 +2235,7 @@ fn entry_stub_hit_body(entry_ptr: *const c_void, ec: EcPtr) -> Option<*const u8>
         Some(blockref) => {
             let mut asm = Assembler::new();
             asm.jmp(unsafe { blockref.as_ref() }.start_addr.into());
-            asm.compile(cb);
+            asm.compile(cb, Some(ocb));
             blockref
         }
         // If this block hasn't yet been compiled, generate blocks after the entry guard.
@@ -2273,7 +2279,7 @@ pub fn gen_entry_stub(entry_address: usize, ocb: &mut OutlinedCb) -> Option<Code
     // Not really a side exit, just don't need a padded jump here.
     asm.jmp(CodegenGlobals::get_entry_stub_hit_trampoline().as_side_exit());
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     if ocb.has_dropped_bytes() {
         return None; // No space
@@ -2296,7 +2302,7 @@ pub fn gen_entry_stub_hit_trampoline(ocb: &mut OutlinedCb) -> CodePtr {
     // Jump to the address returned by the entry_stub_hit() call
     asm.jmp_opnd(jump_addr);
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     code_ptr
 }
@@ -2316,8 +2322,8 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &Branch) {
     asm.comment("regenerate_branch");
     branch.gen_fn.call(
         &mut asm,
-        branch.get_target_address(0).unwrap(),
-        branch.get_target_address(1),
+        Target::CodePtr(branch.get_target_address(0).unwrap()),
+        branch.get_target_address(1).map(|addr| Target::CodePtr(addr)),
     );
 
     // Rewrite the branch
@@ -2325,7 +2331,7 @@ fn regenerate_branch(cb: &mut CodeBlock, branch: &Branch) {
     let old_dropped_bytes = cb.has_dropped_bytes();
     cb.set_write_ptr(branch.start_addr);
     cb.set_dropped_bytes(false);
-    asm.compile(cb);
+    asm.compile(cb, None);
     let new_end_addr = cb.get_write_ptr();
 
     branch.end_addr.set(new_end_addr);
@@ -2584,7 +2590,7 @@ fn gen_branch_stub(
     // Not really a side exit, just don't need a padded jump here.
     asm.jmp(CodegenGlobals::get_branch_stub_hit_trampoline().as_side_exit());
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     if ocb.has_dropped_bytes() {
         // No space
@@ -2623,7 +2629,7 @@ pub fn gen_branch_stub_hit_trampoline(ocb: &mut OutlinedCb) -> CodePtr {
     // Jump to the address returned by the branch_stub_hit() call
     asm.jmp_opnd(jump_addr);
 
-    asm.compile(ocb);
+    asm.compile(ocb, None);
 
     code_ptr
 }
@@ -2715,7 +2721,7 @@ pub fn gen_branch(
     // Call the branch generation function
     asm.mark_branch_start(&branch);
     if let Some(dst_addr) = target0_addr {
-        branch.gen_fn.call(asm, dst_addr, target1_addr);
+        branch.gen_fn.call(asm, Target::CodePtr(dst_addr), target1_addr.map(|addr| Target::CodePtr(addr)));
     }
     asm.mark_branch_end(&branch);
 }
@@ -2732,7 +2738,7 @@ pub fn gen_direct_jump(jit: &mut JITState, ctx: &Context, target0: BlockId, asm:
         // Call the branch generation function
         asm.comment("gen_direct_jmp: existing block");
         asm.mark_branch_start(&branch);
-        branch.gen_fn.call(asm, block_addr, None);
+        branch.gen_fn.call(asm, Target::CodePtr(block_addr), None);
         asm.mark_branch_end(&branch);
 
         BranchTarget::Block(blockref)
@@ -2787,7 +2793,7 @@ pub fn defer_compilation(
     asm.comment("defer_compilation");
     asm.mark_branch_start(&branch);
     if let Some(dst_addr) = target0_address {
-        branch.gen_fn.call(asm, dst_addr, None);
+        branch.gen_fn.call(asm, Target::CodePtr(dst_addr), None);
     }
     asm.mark_branch_end(&branch);
 
@@ -2962,7 +2968,7 @@ pub fn invalidate_block_version(blockref: &BlockRef) {
             let mut asm = Assembler::new();
             asm.jmp(block_entry_exit.as_side_exit());
             cb.set_dropped_bytes(false);
-            asm.compile(&mut cb);
+            asm.compile(&mut cb, Some(ocb));
 
             assert!(
                 cb.get_write_ptr() <= block_end,
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index c93213b484..854ef6cf14 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -2,6 +2,7 @@
 //! generated code if and when these assumptions are invalidated.
 
 use crate::asm::OutlinedCb;
+use crate::backend::ir::Assembler;
 use crate::codegen::*;
 use crate::core::*;
 use crate::cruby::*;
@@ -82,12 +83,13 @@ impl Invariants {
 #[must_use]
 pub fn assume_bop_not_redefined(
     jit: &mut JITState,
+    asm: &mut Assembler,
     ocb: &mut OutlinedCb,
     klass: RedefinitionFlag,
     bop: ruby_basic_operators,
 ) -> bool {
     if unsafe { BASIC_OP_UNREDEFINED_P(bop, klass) } {
-        jit_ensure_block_entry_exit(jit, ocb);
+        jit_ensure_block_entry_exit(jit, asm, ocb);
         jit.bop_assumptions.push((klass, bop));
 
         return true;
@@ -131,13 +133,14 @@ pub fn track_method_lookup_stability_assumption(
 // default behavior.
 pub fn assume_method_basic_definition(
     jit: &mut JITState,
+    asm: &mut Assembler,
     ocb: &mut OutlinedCb,
     klass: VALUE,
     mid: ID
 ) -> bool {
     if unsafe { rb_method_basic_definition_p(klass, mid) } != 0 {
         let cme = unsafe { rb_callable_method_entry(klass, mid) };
-        jit.assume_method_lookup_stable(ocb, cme);
+        jit.assume_method_lookup_stable(asm, ocb, cme);
         true
     } else {
         false
@@ -146,11 +149,11 @@ pub fn assume_method_basic_definition(
 
 /// Tracks that a block is assuming it is operating in single-ractor mode.
 #[must_use]
-pub fn assume_single_ractor_mode(jit: &mut JITState, ocb: &mut OutlinedCb) -> bool {
+pub fn assume_single_ractor_mode(jit: &mut JITState, asm: &mut Assembler, ocb: &mut OutlinedCb) -> bool {
     if unsafe { rb_yjit_multi_ractor_p() } {
         false
     } else {
-        jit_ensure_block_entry_exit(jit, ocb);
+        jit_ensure_block_entry_exit(jit, asm, ocb);
         jit.block_assumes_single_ractor = true;
 
         true
@@ -524,7 +527,7 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
 
             cb.set_write_ptr(patch.inline_patch_pos);
             cb.set_dropped_bytes(false);
-            asm.compile(cb);
+            asm.compile(cb, None);
             last_patch_end = cb.get_write_ptr().raw_ptr();
         }
         cb.set_pos(old_pos);
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 51ea050c23..f624cf8f0a 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -135,7 +135,7 @@ macro_rules! make_counters {
 
         /// Enum to represent a counter
         #[allow(non_camel_case_types)]
-        #[derive(Clone, Copy)]
+        #[derive(Clone, Copy, PartialEq, Eq, Debug)]
         pub enum Counter { $($counter_name),+ }
 
         impl Counter {
diff --git a/yjit/src/utils.rs b/yjit/src/utils.rs
index d9a75b5302..396b330abf 100644
--- a/yjit/src/utils.rs
+++ b/yjit/src/utils.rs
@@ -263,7 +263,7 @@ mod tests {
         let mut cb = CodeBlock::new_dummy(1024);
 
         print_int(&mut asm, Opnd::Imm(42));
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
     }
 
     #[test]
@@ -272,6 +272,6 @@ mod tests {
         let mut cb = CodeBlock::new_dummy(1024);
 
         print_str(&mut asm, "Hello, world!");
-        asm.compile(&mut cb);
+        asm.compile(&mut cb, None);
     }
 }