3 files changed, 43 insertions, 15 deletions
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index 72a4bc711b..4e7e9a1542 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -933,8 +933,36 @@ pub struct SideExitContext {
     /// PC of the instruction being compiled
     pub pc: *mut VALUE,
 
-    /// Context when it started to compile the instruction
-    pub ctx: Context,
+    /// Context fields used by get_generic_ctx()
+    pub stack_size: u8,
+    pub sp_offset: i8,
+    pub reg_temps: RegTemps,
+}
+
+impl SideExitContext {
+    /// Convert PC and Context into SideExitContext
+    pub fn new(pc: *mut VALUE, ctx: Context) -> Self {
+        let exit_ctx = SideExitContext {
+            pc,
+            stack_size: ctx.get_stack_size(),
+            sp_offset: ctx.get_sp_offset(),
+            reg_temps: ctx.get_reg_temps(),
+        };
+        if cfg!(debug_assertions) {
+            // Assert that we're not losing any mandatory metadata
+            assert_eq!(exit_ctx.get_ctx(), ctx.get_generic_ctx());
+        }
+        exit_ctx
+    }
+
+    /// Convert SideExitContext to Context
+    fn get_ctx(&self) -> Context {
+        let mut ctx = Context::default();
+        ctx.set_stack_size(self.stack_size);
+        ctx.set_sp_offset(self.sp_offset);
+        ctx.set_reg_temps(self.reg_temps);
+        ctx
+    }
 }
 
 /// Object into which we assemble instructions to be
@@ -1039,10 +1067,10 @@ impl Assembler
         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()),
-                });
+                *context = Some(SideExitContext::new(
+                    self.side_exit_pc.unwrap(),
+                    self.ctx.with_stack_size(self.side_exit_stack_size.unwrap()),
+                ));
             }
         }
 
@@ -1052,14 +1080,10 @@ impl Assembler
 
     /// 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);
+                let exit_code = gen_outlined_exit(side_exit_context.pc, &side_exit_context.get_ctx(), ocb);
                 self.side_exits.insert(side_exit_context.clone(), exit_code);
                 exit_code
             }
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index cc2f33ef63..f6f0225de4 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -546,7 +546,7 @@ pub fn jit_ensure_block_entry_exit(jit: &mut JITState, asm: &mut Assembler, ocb:
     // 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 Assembler.
-        let side_exit_context = SideExitContext { pc: jit.pc, ctx: block_starting_context.clone() };
+        let side_exit_context = SideExitContext::new(jit.pc, block_starting_context.clone());
         let entry_exit = asm.get_side_exit(&side_exit_context, None, ocb);
         jit.block_entry_exit = Some(entry_exit);
     } else {
@@ -800,7 +800,7 @@ pub fn gen_single_block(
         let blockid_idx = blockid.idx;
         let chain_depth = if asm.ctx.get_chain_depth() > 0 { format!("(chain_depth: {})", asm.ctx.get_chain_depth()) } else { "".to_string() };
         asm.comment(&format!("Block: {} {}", iseq_get_location(blockid.iseq, blockid_idx), chain_depth));
-        asm.comment(&format!("reg_temps: {:08b}", ctx.get_reg_temps().as_u8()));
+        asm.comment(&format!("reg_temps: {:08b}", asm.ctx.get_reg_temps().as_u8()));
     }
 
     // For each instruction to compile
@@ -8555,7 +8555,7 @@ mod tests {
     #[test]
     fn test_get_side_exit() {
         let (_jit, ctx, mut asm, _, mut ocb) = setup_codegen();
-        let side_exit_context = SideExitContext { pc: 0 as _, ctx };
+        let side_exit_context = SideExitContext::new(0 as _, ctx);
         asm.get_side_exit(&side_exit_context, None, &mut ocb);
         assert!(ocb.unwrap().get_write_pos() > 0);
     }
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index b5a9d69205..97a84c6306 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -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, Copy, Default, Eq, Hash, PartialEq, Debug)]
+#[derive(Clone, Default, Eq, Hash, PartialEq, Debug)]
 pub struct Context {
     // Number of values currently on the temporary stack
     stack_size: u8,
@@ -1567,6 +1567,10 @@ impl Context {
         self.stack_size
     }
 
+    pub fn set_stack_size(&mut self, stack_size: u8) {
+        self.stack_size = 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();