ZJIT: Combine two IFUNC tag guards into one

Use IntAnd + GuardBitEquals to check (block_handler & 0x3) == 0x3,
matching VM_BH_IFUNC_P() in the interpreter. Add the IntAnd HIR
instruction (placed next to IntOr) for bitwise AND on raw C integers.

3 files changed, 28 insertions, 25 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index b4b66e5240..b21d4070d2 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -644,6 +644,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::FixnumAnd { left, right } => gen_fixnum_and(asm, opnd!(left), opnd!(right)),
         Insn::FixnumOr { left, right } => gen_fixnum_or(asm, opnd!(left), opnd!(right)),
         Insn::FixnumXor { left, right } => gen_fixnum_xor(asm, opnd!(left), opnd!(right)),
+        Insn::IntAnd { left, right } => asm.and(opnd!(left), opnd!(right)),
         Insn::IntOr { left, right } => gen_int_or(asm, opnd!(left), opnd!(right)),
         &Insn::FixnumLShift { left, right, state } => {
             // We only create FixnumLShift when we know the shift amount statically and it's in [0,
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 90d6682af8..097e1b824b 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -1072,6 +1072,7 @@ pub enum Insn {
     FixnumAnd  { left: InsnId, right: InsnId },
     FixnumOr   { left: InsnId, right: InsnId },
     FixnumXor  { left: InsnId, right: InsnId },
+    IntAnd     { left: InsnId, right: InsnId },
     IntOr      { left: InsnId, right: InsnId },
     FixnumLShift { left: InsnId, right: InsnId, state: InsnId },
     FixnumRShift { left: InsnId, right: InsnId },
@@ -1275,6 +1276,7 @@ macro_rules! for_each_operand_impl {
             | Insn::FixnumAnd { left, right }
             | Insn::FixnumOr { left, right }
             | Insn::FixnumXor { left, right }
+            | Insn::IntAnd { left, right }
             | Insn::IntOr { left, right }
             | Insn::FixnumRShift { left, right }
             | Insn::IsBitEqual { left, right }
@@ -1617,6 +1619,7 @@ impl Insn {
             Insn::FixnumAnd { .. } => effects::Empty,
             Insn::FixnumOr { .. } => effects::Empty,
             Insn::FixnumXor { .. } => effects::Empty,
+            Insn::IntAnd { .. } => effects::Empty,
             Insn::IntOr { .. } => effects::Empty,
             Insn::FixnumLShift { .. } => effects::Empty,
             Insn::FixnumRShift { .. } => effects::Empty,
@@ -2006,6 +2009,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::FixnumAnd  { left, right, .. } => { write!(f, "FixnumAnd {left}, {right}") },
             Insn::FixnumOr   { left, right, .. } => { write!(f, "FixnumOr {left}, {right}") },
             Insn::FixnumXor  { left, right, .. } => { write!(f, "FixnumXor {left}, {right}") },
+            Insn::IntAnd     { left, right } => { write!(f, "IntAnd {left}, {right}") },
             Insn::IntOr      { left, right } => { write!(f, "IntOr {left}, {right}") },
             Insn::FixnumLShift { left, right, .. } => { write!(f, "FixnumLShift {left}, {right}") },
             Insn::FixnumRShift { left, right, .. } => { write!(f, "FixnumRShift {left}, {right}") },
@@ -2795,6 +2799,7 @@ impl Function {
             &FixnumAnd { left, right } => FixnumAnd { left: find!(left), right: find!(right) },
             &FixnumOr { left, right } => FixnumOr { left: find!(left), right: find!(right) },
             &FixnumXor { left, right } => FixnumXor { left: find!(left), right: find!(right) },
+            &IntAnd { left, right } => IntAnd { left: find!(left), right: find!(right) },
             &IntOr { left, right } => IntOr { left: find!(left), right: find!(right) },
             &FixnumLShift { left, right, state } => FixnumLShift { left: find!(left), right: find!(right), state },
             &FixnumRShift { left, right } => FixnumRShift { left: find!(left), right: find!(right) },
@@ -3057,6 +3062,7 @@ impl Function {
             Insn::FixnumAnd  { .. } => types::Fixnum,
             Insn::FixnumOr   { .. } => types::Fixnum,
             Insn::FixnumXor  { .. } => types::Fixnum,
+            Insn::IntAnd { .. } => types::CInt64,
             Insn::IntOr { left, .. } => self.type_of(*left).unspecialized(),
             Insn::FixnumLShift { .. } => types::Fixnum,
             Insn::FixnumRShift { .. } => types::Fixnum,
@@ -6248,7 +6254,8 @@ impl Function {
                     Err(ValidationError::MiscValidationError(insn_id, "IsBitEqual can only compare CInt/CInt or RubyValue/RubyValue".to_string()))
                 }
             }
-            Insn::IntOr { left, right } => {
+            Insn::IntAnd { left, right }
+            | Insn::IntOr { left, right } => {
                 // TODO: Expand this to other matching C integer sizes when we need them.
                 self.assert_subtype(insn_id, left, types::CInt64)?;
                 self.assert_subtype(insn_id, right, types::CInt64)
@@ -8161,20 +8168,13 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                             return_type: types::CInt64,
                         });
 
-                        // Guard that the block handler is an IFUNC (tag bits & 0x3 == 0x3)
-                        // bit 0x1 must be set (iseq or ifunc)
-                        fun.push_insn(block, Insn::GuardAnyBitSet {
-                            val: block_handler,
-                            mask: Const::CUInt64(0x1),
-                            mask_name: None,
-                            reason: SideExitReason::InvokeBlockNotIfunc,
-                            state: exit_id,
-                        });
-                        // bit 0x2 must be set (ifunc specifically, not iseq)
-                        fun.push_insn(block, Insn::GuardAnyBitSet {
-                            val: block_handler,
-                            mask: Const::CUInt64(0x2),
-                            mask_name: None,
+                        // Guard that the block handler is an IFUNC (tag bits & 0x3 == 0x3),
+                        // matching VM_BH_IFUNC_P() in the interpreter.
+                        let tag_mask = fun.push_insn(block, Insn::Const { val: Const::CInt64(0x3) });
+                        let tag_bits = fun.push_insn(block, Insn::IntAnd { left: block_handler, right: tag_mask });
+                        fun.push_insn(block, Insn::GuardBitEquals {
+                            val: tag_bits,
+                            expected: Const::CInt64(0x3),
                             reason: SideExitReason::InvokeBlockNotIfunc,
                             state: exit_id,
                         });
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 315467ca72..c9f2a45b04 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -15102,17 +15102,19 @@ mod hir_opt_tests {
           v10:Fixnum[1] = Const Value(1)
           v12:CPtr = GetEP 0
           v13:CInt64 = LoadField v12, :_env_data_index_specval@0x1000
-          v14:CInt64 = GuardAnyBitSet v13, CUInt64(1)
-          v15:CInt64 = GuardAnyBitSet v13, CUInt64(2)
-          v16:BasicObject = InvokeBlockIfunc v13, v10
-          v20:Fixnum[2] = Const Value(2)
-          v22:CPtr = GetEP 0
-          v23:CInt64 = LoadField v22, :_env_data_index_specval@0x1000
-          v24:CInt64 = GuardAnyBitSet v23, CUInt64(1)
-          v25:CInt64 = GuardAnyBitSet v23, CUInt64(2)
-          v26:BasicObject = InvokeBlockIfunc v23, v20
+          v14:CInt64[3] = Const CInt64(3)
+          v15:CInt64 = IntAnd v13, v14
+          v16:CInt64[3] = GuardBitEquals v15, CInt64(3)
+          v17:BasicObject = InvokeBlockIfunc v13, v10
+          v21:Fixnum[2] = Const Value(2)
+          v23:CPtr = GetEP 0
+          v24:CInt64 = LoadField v23, :_env_data_index_specval@0x1000
+          v25:CInt64[3] = Const CInt64(3)
+          v26:CInt64 = IntAnd v24, v25
+          v27:CInt64[3] = GuardBitEquals v26, CInt64(3)
+          v28:BasicObject = InvokeBlockIfunc v24, v21
           CheckInterrupts
-          Return v26
+          Return v28
         ");
     }
 }