ZJIT: Add a YarvInsnIdx type for clarity (#16783)

3 files changed, 19 insertions, 16 deletions

diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index 40df1fab03..6728d884f9 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -261,6 +261,9 @@ pub struct ID(pub ::std::os::raw::c_ulong);
 /// Pointer to an ISEQ
 pub type IseqPtr = *const rb_iseq_t;
 
+/// Index of a YARV instruction within an ISEQ's bytecode array.
+pub type YarvInsnIdx = usize;
+
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct ShapeId(pub u32);
 
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 671249f701..935072af40 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -3323,7 +3323,7 @@ impl Function {
     /// Return the profiled type of the HIR instruction at the given ISEQ instruction
     /// index, if it is known to be monomorphic or skewed polymorphic. This historical type
     /// record is not a guarantee and must be checked with a GuardType or similar.
-    fn profiled_type_of_at(&self, insn: InsnId, iseq_insn_idx: usize) -> Option<ProfiledType> {
+    fn profiled_type_of_at(&self, insn: InsnId, iseq_insn_idx: YarvInsnIdx) -> Option<ProfiledType> {
         match self.resolve_receiver_type_from_profile(insn, iseq_insn_idx) {
             ReceiverTypeResolution::Monomorphic { profiled_type }
             | ReceiverTypeResolution::SkewedPolymorphic { profiled_type } => Some(profiled_type),
@@ -3492,14 +3492,14 @@ impl Function {
     /// Returns:
     /// - `StaticallyKnown` if the receiver's exact class is known at compile-time
     /// - Result of [`Self::resolve_receiver_type_from_profile`] if we need to check profile data
-    fn resolve_receiver_type(&self, recv: InsnId, recv_type: Type, insn_idx: usize) -> ReceiverTypeResolution {
+    fn resolve_receiver_type(&self, recv: InsnId, recv_type: Type, insn_idx: YarvInsnIdx) -> ReceiverTypeResolution {
         if let Some(class) = recv_type.runtime_exact_ruby_class() {
             return ReceiverTypeResolution::StaticallyKnown { class };
         }
         self.resolve_receiver_type_from_profile(recv, insn_idx)
     }
 
-    fn polymorphic_summary(&self, profiles: &ProfileOracle, recv: InsnId, insn_idx: usize) -> Option<TypeDistributionSummary> {
+    fn polymorphic_summary(&self, profiles: &ProfileOracle, recv: InsnId, insn_idx: YarvInsnIdx) -> Option<TypeDistributionSummary> {
         let Some(entries) = profiles.types.get(&insn_idx) else {
             return None;
         };
@@ -3523,7 +3523,7 @@ impl Function {
     /// - `Megamorphic`/`SkewedMegamorphic` if the receiver has too many types to optimize
     ///   (SkewedMegamorphic may be optimized in the future, but for now we don't)
     /// - `NoProfile` if we have no type information
-    fn resolve_receiver_type_from_profile(&self, recv: InsnId, insn_idx: usize) -> ReceiverTypeResolution {
+    fn resolve_receiver_type_from_profile(&self, recv: InsnId, insn_idx: YarvInsnIdx) -> ReceiverTypeResolution {
         let Some(profiles) = self.profiles.as_ref() else {
             return ReceiverTypeResolution::NoProfile;
         };
@@ -6412,7 +6412,7 @@ impl<'a> std::fmt::Display for FunctionPrinter<'a> {
 #[derive(Debug, Clone, PartialEq)]
 pub struct FrameState {
     pub iseq: IseqPtr,
-    insn_idx: usize,
+    insn_idx: YarvInsnIdx,
     // Ruby bytecode instruction pointer
     pub pc: *const VALUE,
 
@@ -6422,7 +6422,7 @@ pub struct FrameState {
 
 impl FrameState {
     /// Get the YARV instruction index for the current instruction
-    pub fn insn_idx(&self) -> usize {
+    pub fn insn_idx(&self) -> YarvInsnIdx {
         self.insn_idx
     }
 
@@ -6702,7 +6702,7 @@ struct ProfileOracle {
     /// instruction index. At a given ISEQ instruction, the interpreter has profiled the stack
     /// operands to a given ISEQ instruction, and this list of pairs of (InsnId, Type) map that
     /// profiling information into HIR instructions.
-    types: HashMap<usize, Vec<(InsnId, TypeDistributionSummary)>>,
+    types: HashMap<YarvInsnIdx, Vec<(InsnId, TypeDistributionSummary)>>,
 }
 
 impl ProfileOracle {
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
index ae496745bc..d8bcd50d96 100644
--- a/zjit/src/profile.rs
+++ b/zjit/src/profile.rs
@@ -13,7 +13,7 @@ use crate::stats::with_time_stat;
 struct Profiler {
     cfp: CfpPtr,
     iseq: IseqPtr,
-    insn_idx: usize,
+    insn_idx: YarvInsnIdx,
 }
 
 impl Profiler {
@@ -424,7 +424,7 @@ pub struct IseqProfile {
     entries: Vec<ProfileEntry>,
 
     /// Method entries for `super` calls (stored as VALUE to be GC-safe)
-    super_cme: HashMap<usize, TypeDistribution>
+    super_cme: HashMap<YarvInsnIdx, TypeDistribution>
 }
 
 impl IseqProfile {
@@ -436,7 +436,7 @@ impl IseqProfile {
     }
 
     /// Get or create a mutable profile entry for the given instruction index.
-    fn entry_mut(&mut self, insn_idx: usize) -> &mut ProfileEntry {
+    fn entry_mut(&mut self, insn_idx: YarvInsnIdx) -> &mut ProfileEntry {
         let idx = insn_idx as u32;
         match self.entries.binary_search_by_key(&idx, |e| e.insn_idx) {
             Ok(i) => &mut self.entries[i],
@@ -452,14 +452,14 @@ impl IseqProfile {
     }
 
     /// Get a profile entry for the given instruction index (read-only).
-    fn entry(&self, insn_idx: usize) -> Option<&ProfileEntry> {
+    fn entry(&self, insn_idx: YarvInsnIdx) -> Option<&ProfileEntry> {
         let idx = insn_idx as u32;
         self.entries.binary_search_by_key(&idx, |e| e.insn_idx)
             .ok().map(|i| &self.entries[i])
     }
 
     /// Check if enough profiles have been gathered for this instruction.
-    pub fn done_profiling_at(&self, insn_idx: usize) -> bool {
+    pub fn done_profiling_at(&self, insn_idx: YarvInsnIdx) -> bool {
         self.entry(insn_idx).map_or(false, |e| e.profiles_remaining == 0)
     }
 
@@ -467,7 +467,7 @@ impl IseqProfile {
     /// `sp` is the current stack pointer (after the args and receiver).
     /// `argc` is the number of arguments (not counting receiver).
     /// Returns true if enough profiles have been gathered and the ISEQ should be recompiled.
-    pub fn profile_send_at(&mut self, iseq: IseqPtr, insn_idx: usize, sp: *const VALUE, argc: usize) -> bool {
+    pub fn profile_send_at(&mut self, iseq: IseqPtr, insn_idx: YarvInsnIdx, sp: *const VALUE, argc: usize) -> bool {
         let n = argc + 1; // args + receiver
         let entry = self.entry_mut(insn_idx);
         if entry.opnd_types.is_empty() {
@@ -487,7 +487,7 @@ impl IseqProfile {
     /// This may be called on an instruction that was already profiled by YARV,
     /// so we reset the counter to re-profile with the new shapes seen at runtime.
     /// Returns true if enough profiles have been gathered and the ISEQ should be recompiled.
-    pub fn profile_self_at(&mut self, iseq: IseqPtr, insn_idx: usize, self_val: VALUE) -> bool {
+    pub fn profile_self_at(&mut self, iseq: IseqPtr, insn_idx: YarvInsnIdx, self_val: VALUE) -> bool {
         let entry = self.entry_mut(insn_idx);
         // Reset profiling if the previous round already finished (stale YARV profiles).
         // This ensures we collect num_profiles samples of the new shapes before recompiling.
@@ -505,11 +505,11 @@ impl IseqProfile {
     }
 
     /// Get profiled operand types for a given instruction index
-    pub fn get_operand_types(&self, insn_idx: usize) -> Option<&[TypeDistribution]> {
+    pub fn get_operand_types(&self, insn_idx: YarvInsnIdx) -> Option<&[TypeDistribution]> {
         self.entry(insn_idx).map(|e| e.opnd_types.as_slice()).filter(|s| !s.is_empty())
     }
 
-    pub fn get_super_method_entry(&self, insn_idx: usize) -> Option<*const rb_callable_method_entry_t> {
+    pub fn get_super_method_entry(&self, insn_idx: YarvInsnIdx) -> Option<*const rb_callable_method_entry_t> {
         let Some(entry) = self.super_cme.get(&insn_idx) else { return None };
         let summary = TypeDistributionSummary::new(entry);