ZJIT: Replace Option<i32> with Recompile enum for side exit recompilation (ruby#16665)

Author: k0kubun

zjit/src/backend/lir.rs

@@ -558,14 +558,11 @@ pub struct SideExit {
 }
 
 /// Arguments for the recompile callback on side exit.
-/// Used for both no-profile sends (argc >= 0) and shape guard failures (argc = -1).
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 pub struct SideExitRecompile {
     pub iseq: Opnd,
     pub insn_idx: u32,
-    /// Number of arguments (not including receiver) for send profiling.
-    /// -1 means profile self from CFP for shape guard exits.
-    pub argc: i32,
+    pub strategy: hir::Recompile,
 }
 
 /// Branch target (something that we can jump to)
@@ -2687,7 +2684,10 @@ impl Assembler
                     EC,
                     recompile.iseq,
                     Opnd::UImm(recompile.insn_idx as u64),
-                    Opnd::Imm(recompile.argc as i64)
+                    Opnd::Imm(match recompile.strategy {
+                        hir::Recompile::ProfileSend { argc } => argc as i64,
+                        hir::Recompile::ProfileSelf => -1,
+                    })
                 );
             }
             compile_exit_return(asm);

zjit/src/codegen.rs

@@ -21,7 +21,7 @@ use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, Assembler, C_ARG_OPNDS, C_RET_OPND, CFP, EC, NATIVE_STACK_PTR, Opnd, SP, SideExit, SideExitRecompile, Target, asm_ccall, asm_comment};
 use crate::hir::{iseq_to_hir, BlockId, Invariant, RangeType, SideExitReason::{self, *}, SpecialBackrefSymbol, SpecialObjectType};
-use crate::hir::{BlockHandler, Const, FrameState, Function, Insn, InsnId, SendFallbackReason};
+use crate::hir::{BlockHandler, Const, FrameState, Function, Insn, InsnId, Recompile, SendFallbackReason};
 use crate::hir_type::{types, Type};
 use crate::options::{get_option, PerfMap};
 use crate::cast::IntoUsize;
@@ -1238,7 +1238,7 @@ fn gen_setglobal(jit: &mut JITState, asm: &mut Assembler, id: ID, val: Opnd, sta
 }
 
 /// Side-exit into the interpreter
-fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, reason: &SideExitReason, recompile: Option<i32>, state: &FrameState) {
+fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, reason: &SideExitReason, recompile: Option<Recompile>, state: &FrameState) {
     asm.jmp(side_exit_with_recompile(jit, state, *reason, recompile));
 }
 
@@ -2631,7 +2631,7 @@ fn gen_guard_type_not(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, g
 }
 
 /// Compile an identity check with a side exit
-fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: crate::hir::Const, reason: SideExitReason, recompile: Option<i32>, state: &FrameState) -> lir::Opnd {
+fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: crate::hir::Const, reason: SideExitReason, recompile: Option<Recompile>, state: &FrameState) -> lir::Opnd {
     if matches!(reason, SideExitReason::GuardShape(_) ) {
         gen_incr_counter(asm, Counter::guard_shape_count);
     }
@@ -2977,13 +2977,12 @@ fn side_exit(jit: &JITState, state: &FrameState, reason: SideExitReason) -> Targ
 }
 
 /// Build a Target::SideExit that optionally triggers exit_recompile on the exit path.
-/// When `recompile` is Some(argc), the side exit calls exit_recompile with that argc.
-fn side_exit_with_recompile(jit: &JITState, state: &FrameState, reason: SideExitReason, recompile: Option<i32>) -> Target {
+fn side_exit_with_recompile(jit: &JITState, state: &FrameState, reason: SideExitReason, recompile: Option<Recompile>) -> Target {
     let mut exit = build_side_exit(jit, state);
-    exit.recompile = recompile.map(|argc| SideExitRecompile {
+    exit.recompile = recompile.map(|strategy| SideExitRecompile {
         iseq: Opnd::Value(VALUE::from(jit.iseq)),
         insn_idx: state.insn_idx() as u32,
-        argc,
+        strategy,
     });
     Target::SideExit { exit, reason }
 }

zjit/src/hir.rs

@@ -540,6 +540,15 @@ pub enum SideExitReason {
     InvokeBlockNotIfunc,
 }
 
+/// Controls how a side exit triggers recompilation.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub enum Recompile {
+    /// Profile receiver + arguments from the stack (for sends without profile data).
+    ProfileSend { argc: i32 },
+    /// Profile self from the CFP (for shape guard failures).
+    ProfileSelf,
+}
+
 #[derive(Debug, Clone, Copy)]
 pub enum MethodType {
     Iseq,
@@ -1090,9 +1099,7 @@ pub enum Insn {
     GuardType { val: InsnId, guard_type: Type, state: InsnId },
     GuardTypeNot { val: InsnId, guard_type: Type, state: InsnId },
     /// Side-exit if val is not the expected Const.
-    /// `recompile`: if Some(argc), the side exit triggers exit_recompile.
-    /// argc >= 0 profiles receiver + args from stack; argc == -1 profiles self from CFP.
-    GuardBitEquals { val: InsnId, expected: Const, reason: SideExitReason, state: InsnId, recompile: Option<i32> },
+    GuardBitEquals { val: InsnId, expected: Const, reason: SideExitReason, state: InsnId, recompile: Option<Recompile> },
     /// Side-exit if (val & mask) == 0
     GuardAnyBitSet { val: InsnId, mask: Const, mask_name: Option<ID>, reason: SideExitReason, state: InsnId },
     /// Side-exit if (val & mask) != 0
@@ -1107,9 +1114,9 @@ pub enum Insn {
     PatchPoint { invariant: Invariant, state: InsnId },
 
     /// Side-exit into the interpreter.
-    /// If `recompile` is set, the side exit will profile the send and invalidate the ISEQ
+    /// If recompile is not None, the side exit will profile and invalidate the ISEQ
     /// so that it gets recompiled with the new profile data.
-    SideExit { state: InsnId, reason: SideExitReason, recompile: Option<i32> },
+    SideExit { state: InsnId, reason: SideExitReason, recompile: Option<Recompile> },
 
     /// Increment a counter in ZJIT stats
     IncrCounter(Counter),
@@ -4392,7 +4399,7 @@ impl Function {
         })
     }
 
-    fn guard_shape(&mut self, block: BlockId, val: InsnId, expected: ShapeId, state: InsnId, recompile: Option<i32>) -> InsnId {
+    fn guard_shape(&mut self, block: BlockId, val: InsnId, expected: ShapeId, state: InsnId, recompile: Option<Recompile>) -> InsnId {
         self.push_insn(block, Insn::GuardBitEquals {
             val,
             expected: Const::CShape(expected),
@@ -4543,7 +4550,7 @@ impl Function {
                         }
                         let self_val = self.load_ivar_guard_type(block, self_val, recv_type, state);
                         let shape = self.load_shape(block, self_val);
-                        self.guard_shape(block, shape, recv_type.shape(), state, Some(-1));
+                        self.guard_shape(block, shape, recv_type.shape(), state, Some(Recompile::ProfileSelf));
                         let replacement = self.load_ivar(block, self_val, recv_type, id, state);
                         self.make_equal_to(insn_id, replacement);
                     }
@@ -5178,7 +5185,7 @@ impl Function {
                 match self.find(insn_id) {
                     Insn::Send { cd, state, reason: SendFallbackReason::SendWithoutBlockNoProfiles | SendFallbackReason::SendNoProfiles, .. } => {
                         let argc = unsafe { vm_ci_argc((*cd).ci) } as i32;
-                        self.push_insn(block, Insn::SideExit { state, reason: SideExitReason::NoProfileSend, recompile: Some(argc) });
+                        self.push_insn(block, Insn::SideExit { state, reason: SideExitReason::NoProfileSend, recompile: Some(Recompile::ProfileSend { argc }) });
                         // SideExit is a terminator; don't add remaining instructions
                         break;
                     }