ZJIT: Improve asm comments for side exits (#13853)

* ZJIT: Improve asm comments for side exits

* Use GuardType(Type) and GuardBitEquals(VALUE)

4 files changed, 41 insertions, 27 deletions

diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index 1a2210c561..5d608f1e63 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -3,6 +3,7 @@ use std::fmt;
 use std::mem::take;
 use crate::codegen::local_size_and_idx_to_ep_offset;
 use crate::cruby::{Qundef, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32};
+use crate::hir::SideExitReason;
 use crate::options::{debug, get_option};
 use crate::{cruby::VALUE};
 use crate::backend::current::*;
@@ -284,6 +285,7 @@ pub enum Target
         stack: Vec<Opnd>,
         locals: Vec<Opnd>,
         c_stack_bytes: usize,
+        reason: SideExitReason,
         // Some if the side exit should write this label. We use it for patch points.
         label: Option<Label>,
     },
@@ -1803,8 +1805,8 @@ impl Assembler
         for (idx, target) in targets {
             // Compile a side exit. Note that this is past the split pass and alloc_regs(),
             // so you can't use a VReg or an instruction that needs to be split.
-            if let Target::SideExit { pc, stack, locals, c_stack_bytes, label } = target {
-                asm_comment!(self, "side exit to the interpreter");
+            if let Target::SideExit { pc, stack, locals, c_stack_bytes, reason, label } = target {
+                asm_comment!(self, "Exit: {reason}");
                 let side_exit_label = if let Some(label) = label {
                     Target::Label(label)
                 } else {
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index e231dcaac0..c19a745197 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -8,7 +8,7 @@ use crate::gc::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, Invariant, RangeType, SpecialObjectType, SELF_PARAM_IDX};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, Invariant, RangeType, SideExitReason, SideExitReason::*, SpecialObjectType, SELF_PARAM_IDX};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
 use crate::options::get_option;
@@ -308,7 +308,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SetLocal { val, ep_offset, level } => return gen_setlocal_with_ep(asm, opnd!(val), *ep_offset, *level),
         Insn::GetConstantPath { ic, state } => gen_get_constant_path(asm, *ic, &function.frame_state(*state)),
         Insn::SetIvar { self_val, id, val, state: _ } => return gen_setivar(asm, opnd!(self_val), *id, opnd!(val)),
-        Insn::SideExit { state, reason: _ } => return gen_side_exit(jit, asm, &function.frame_state(*state)),
+        Insn::SideExit { state, reason } => return gen_side_exit(jit, asm, reason, &function.frame_state(*state)),
         Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state))?,
         Insn::Defined { op_type, obj, pushval, v } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v))?,
@@ -442,7 +442,8 @@ fn gen_patch_point(jit: &mut JITState, asm: &mut Assembler, invariant: &Invarian
     let invariant = invariant.clone();
 
     // Compile a side exit. Fill nop instructions if the last patch point is too close.
-    asm.patch_point(build_side_exit(jit, state, Some(label))?);
+    asm.patch_point(build_side_exit(jit, state, PatchPoint(invariant), Some(label))?);
+
     // Remember the current address as a patch point
     asm.pos_marker(move |code_ptr, cb| {
         match invariant {
@@ -500,8 +501,8 @@ fn gen_setglobal(asm: &mut Assembler, id: ID, val: Opnd) {
 }
 
 /// Side-exit into the interpreter
-fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, state: &FrameState) -> Option<()> {
-    asm.jmp(side_exit(jit, state)?);
+fn gen_side_exit(jit: &mut JITState, asm: &mut Assembler, reason: &SideExitReason, state: &FrameState) -> Option<()> {
+    asm.jmp(side_exit(jit, state, *reason)?);
     Some(())
 }
 
@@ -886,7 +887,7 @@ fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, righ
     // Add left + right and test for overflow
     let left_untag = asm.sub(left, Opnd::Imm(1));
     let out_val = asm.add(left_untag, right);
-    asm.jo(side_exit(jit, state)?);
+    asm.jo(side_exit(jit, state, FixnumAddOverflow)?);
 
     Some(out_val)
 }
@@ -895,7 +896,7 @@ fn gen_fixnum_add(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, righ
 fn gen_fixnum_sub(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
     // Subtract left - right and test for overflow
     let val_untag = asm.sub(left, right);
-    asm.jo(side_exit(jit, state)?);
+    asm.jo(side_exit(jit, state, FixnumSubOverflow)?);
     let out_val = asm.add(val_untag, Opnd::Imm(1));
 
     Some(out_val)
@@ -910,7 +911,7 @@ fn gen_fixnum_mult(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, rig
     let out_val = asm.mul(left_untag, right_untag);
 
     // Test for overflow
-    asm.jo_mul(side_exit(jit, state)?);
+    asm.jo_mul(side_exit(jit, state, FixnumMultOverflow)?);
     let out_val = asm.add(out_val, Opnd::UImm(1));
 
     Some(out_val)
@@ -996,7 +997,7 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
     if guard_type.is_subtype(Fixnum) {
         // Check if opnd is Fixnum
         asm.test(val, Opnd::UImm(RUBY_FIXNUM_FLAG as u64));
-        asm.jz(side_exit(jit, state)?);
+        asm.jz(side_exit(jit, state, GuardType(guard_type))?);
     } else if let Some(expected_class) = guard_type.runtime_exact_ruby_class() {
         asm_comment!(asm, "guard exact class");
 
@@ -1004,7 +1005,7 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
         let klass = asm.ccall(rb_yarv_class_of as *const u8, vec![val]);
 
         asm.cmp(klass, Opnd::Value(expected_class));
-        asm.jne(side_exit(jit, state)?);
+        asm.jne(side_exit(jit, state, GuardType(guard_type))?);
     } else {
         unimplemented!("unsupported type: {guard_type}");
     }
@@ -1014,7 +1015,7 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
 /// Compile an identity check with a side exit
 fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: VALUE, state: &FrameState) -> Option<lir::Opnd> {
     asm.cmp(val, Opnd::UImm(expected.into()));
-    asm.jnz(side_exit(jit, state)?);
+    asm.jnz(side_exit(jit, state, GuardBitEquals(expected))?);
     Some(val)
 }
 
@@ -1119,12 +1120,12 @@ fn compile_iseq(iseq: IseqPtr) -> Option<Function> {
 }
 
 /// Build a Target::SideExit for non-PatchPoint instructions
-fn side_exit(jit: &mut JITState, state: &FrameState) -> Option<Target> {
-    build_side_exit(jit, state, None)
+fn side_exit(jit: &mut JITState, state: &FrameState, reason: SideExitReason) -> Option<Target> {
+    build_side_exit(jit, state, reason, None)
 }
 
 /// Build a Target::SideExit out of a FrameState
-fn build_side_exit(jit: &mut JITState, state: &FrameState, label: Option<Label>) -> Option<Target> {
+fn build_side_exit(jit: &mut JITState, state: &FrameState, reason: SideExitReason, label: Option<Label>) -> Option<Target> {
     let mut stack = Vec::new();
     for &insn_id in state.stack() {
         stack.push(jit.get_opnd(insn_id)?);
@@ -1140,6 +1141,7 @@ fn build_side_exit(jit: &mut JITState, state: &FrameState, label: Option<Label>)
         stack,
         locals,
         c_stack_bytes: jit.c_stack_bytes,
+        reason,
         label,
     };
     Some(target)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index dfa9ea1e17..07ea06fcf0 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -7,12 +7,7 @@ use crate::{
     cast::IntoUsize, cruby::*, options::{get_option, DumpHIR}, gc::{get_or_create_iseq_payload, IseqPayload}, state::ZJITState
 };
 use std::{
-    cell::RefCell,
-    collections::{HashMap, HashSet, VecDeque},
-    ffi::{c_int, c_void, CStr},
-    mem::{align_of, size_of},
-    ptr,
-    slice::Iter
+    cell::RefCell, collections::{HashMap, HashSet, VecDeque}, ffi::{c_int, c_void, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
 };
 use crate::hir_type::{Type, types};
 use crate::bitset::BitSet;
@@ -148,6 +143,12 @@ impl Invariant {
     }
 }
 
+impl Display for Invariant {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.print(&PtrPrintMap::identity()).fmt(f)
+    }
+}
+
 #[derive(Debug, Clone, Copy, PartialEq)]
 pub enum SpecialObjectType {
     VMCore = 1,
@@ -402,11 +403,17 @@ impl PtrPrintMap {
     }
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Copy)]
 pub enum SideExitReason {
     UnknownNewarraySend(vm_opt_newarray_send_type),
     UnknownCallType,
     UnknownOpcode(u32),
+    FixnumAddOverflow,
+    FixnumSubOverflow,
+    FixnumMultOverflow,
+    GuardType(Type),
+    GuardBitEquals(VALUE),
+    PatchPoint(Invariant),
 }
 
 impl std::fmt::Display for SideExitReason {
@@ -419,6 +426,9 @@ impl std::fmt::Display for SideExitReason {
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_PACK) => write!(f, "UnknownNewarraySend(PACK)"),
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_PACK_BUFFER) => write!(f, "UnknownNewarraySend(PACK_BUFFER)"),
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_INCLUDE_P) => write!(f, "UnknownNewarraySend(INCLUDE_P)"),
+            SideExitReason::GuardType(guard_type) => write!(f, "GuardType({guard_type})"),
+            SideExitReason::GuardBitEquals(value) => write!(f, "GuardBitEquals({})", value.print(&PtrPrintMap::identity())),
+            SideExitReason::PatchPoint(invariant) => write!(f, "PatchPoint({invariant})"),
             _ => write!(f, "{self:?}"),
         }
     }
diff --git a/zjit/src/invariants.rs b/zjit/src/invariants.rs
index cd8da28617..62b3805485 100644
--- a/zjit/src/invariants.rs
+++ b/zjit/src/invariants.rs
@@ -1,6 +1,6 @@
 use std::{collections::{HashMap, HashSet}};
 
-use crate::{backend::lir::{asm_comment, Assembler}, cruby::{ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag}, hir::{Invariant, PtrPrintMap}, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
+use crate::{backend::lir::{asm_comment, Assembler}, cruby::{ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag}, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
 
 #[derive(Debug, Eq, Hash, PartialEq)]
 struct Jump {
@@ -36,14 +36,14 @@ pub extern "C" fn rb_zjit_bop_redefined(klass: RedefinitionFlag, bop: ruby_basic
         let invariants = ZJITState::get_invariants();
         if let Some(jumps) = invariants.bop_patch_points.get(&(klass, bop)) {
             let cb = ZJITState::get_code_block();
+            let bop = Invariant::BOPRedefined { klass, bop };
+            debug!("BOP is redefined: {}", bop);
 
             // Invalidate all patch points for this BOP
-            let bop = Invariant::BOPRedefined { klass, bop };
-            debug!("BOP is redefined: {}", bop.print(&PtrPrintMap::identity()));
             for jump in jumps {
                 cb.with_write_ptr(jump.from, |cb| {
                     let mut asm = Assembler::new();
-                    asm_comment!(asm, "BOP redefined: {}", bop.print(&PtrPrintMap::identity()));
+                    asm_comment!(asm, "BOP is redefined: {}", bop);
                     asm.jmp(jump.to.into());
                     asm.compile(cb).expect("can write existing code");
                 });