ZJIT: Add Type::has_value method (#15867)HEADmaster

Resolves TODO added in #15863 (See https://github.com/ruby/ruby/pull/15863#discussion_r2687769112)

Adds a method `Type::has_value` for comparing value specialized types with a `Const`.

2 files changed, 100 insertions, 8 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 8c7c838900..7a7d03d75d 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -3812,14 +3812,11 @@ impl Function {
                         }
                     }
                     Insn::GuardBitEquals { val, expected, .. } => {
-                        match self.find(val) {
-                            // TODO: Refactor this into a more general method like
-                            // has_value(Const) that can check on the value specialization
-                            // of the Type instead
-                            Insn::Const { val: const_val } if const_val == expected => {
-                                continue;
-                            }
-                            _ => insn_id
+                        let recv_type = self.type_of(val);
+                        if recv_type.has_value(expected) {
+                            continue;
+                        } else {
+                            insn_id
                         }
                     }
                     Insn::AnyToString { str, .. } if self.is_a(str, types::String) => {
diff --git a/zjit/src/hir_type/mod.rs b/zjit/src/hir_type/mod.rs
index c87f1313b5..061fe16578 100644
--- a/zjit/src/hir_type/mod.rs
+++ b/zjit/src/hir_type/mod.rs
@@ -335,6 +335,25 @@ impl Type {
         self.is_subtype(types::NilClass) || self.is_subtype(types::FalseClass)
     }
 
+    pub fn has_value(&self, val: Const) -> bool {
+        match (self.spec, val) {
+            (Specialization::Object(v1), Const::Value(v2)) => v1 == v2,
+            (Specialization::Int(v1), Const::CBool(v2)) if self.is_subtype(types::CBool) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CInt8(v2)) if self.is_subtype(types::CInt8) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CInt16(v2)) if self.is_subtype(types::CInt16) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CInt32(v2)) if self.is_subtype(types::CInt32) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CInt64(v2)) if self.is_subtype(types::CInt64) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CUInt8(v2)) if self.is_subtype(types::CUInt8) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CUInt16(v2)) if self.is_subtype(types::CUInt16) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CUInt32(v2)) if self.is_subtype(types::CUInt32) => v1 == (v2 as u64),
+            (Specialization::Int(v1), Const::CShape(v2)) if self.is_subtype(types::CShape) => v1 == (v2.0 as u64),
+            (Specialization::Int(v1), Const::CUInt64(v2)) if self.is_subtype(types::CUInt64) => v1 == v2,
+            (Specialization::Int(v1), Const::CPtr(v2)) if self.is_subtype(types::CPtr) => v1 == (v2 as u64),
+            (Specialization::Double(v1), Const::CDouble(v2)) => v1.to_bits() == v2.to_bits(),
+            _ => false,
+        }
+    }
+
     /// Return the object specialization, if any.
     pub fn ruby_object(&self) -> Option<VALUE> {
         match self.spec {
@@ -958,4 +977,80 @@ mod tests {
             assert_bit_equal(d_instance.union(c_instance), Type { bits: bits::ObjectSubclass, spec: Specialization::Type(c_class)});
         });
     }
+
+    #[test]
+    fn has_value() {
+        // With known values
+        crate::cruby::with_rubyvm(|| {
+            let a = rust_str_to_sym("a");
+            let b = rust_str_to_sym("b");
+            let ty = Type::from_value(a);
+            assert!(ty.has_value(Const::Value(a)));
+            assert!(!ty.has_value(Const::Value(b)));
+        });
+
+        let true_ty = Type::from_cbool(true);
+        assert!(true_ty.has_value(Const::CBool(true)));
+        assert!(!true_ty.has_value(Const::CBool(false)));
+
+        let int8_ty = Type::from_cint(types::CInt8, 42);
+        assert!(int8_ty.has_value(Const::CInt8(42)));
+        assert!(!int8_ty.has_value(Const::CInt8(-1)));
+        let neg_int8_ty = Type::from_cint(types::CInt8, -1);
+        assert!(neg_int8_ty.has_value(Const::CInt8(-1)));
+
+        let int16_ty = Type::from_cint(types::CInt16, 1000);
+        assert!(int16_ty.has_value(Const::CInt16(1000)));
+        assert!(!int16_ty.has_value(Const::CInt16(2000)));
+
+        let int32_ty = Type::from_cint(types::CInt32, 100000);
+        assert!(int32_ty.has_value(Const::CInt32(100000)));
+        assert!(!int32_ty.has_value(Const::CInt32(-100000)));
+
+        let int64_ty = Type::from_cint(types::CInt64, i64::MAX);
+        assert!(int64_ty.has_value(Const::CInt64(i64::MAX)));
+        assert!(!int64_ty.has_value(Const::CInt64(0)));
+
+        let uint8_ty = Type::from_cint(types::CUInt8, u8::MAX as i64);
+        assert!(uint8_ty.has_value(Const::CUInt8(u8::MAX)));
+        assert!(!uint8_ty.has_value(Const::CUInt8(0)));
+
+        let uint16_ty = Type::from_cint(types::CUInt16, u16::MAX as i64);
+        assert!(uint16_ty.has_value(Const::CUInt16(u16::MAX)));
+        assert!(!uint16_ty.has_value(Const::CUInt16(1)));
+
+        let uint32_ty = Type::from_cint(types::CUInt32, u32::MAX as i64);
+        assert!(uint32_ty.has_value(Const::CUInt32(u32::MAX)));
+        assert!(!uint32_ty.has_value(Const::CUInt32(42)));
+
+        let uint64_ty = Type::from_cint(types::CUInt64, i64::MAX);
+        assert!(uint64_ty.has_value(Const::CUInt64(i64::MAX as u64)));
+        assert!(!uint64_ty.has_value(Const::CUInt64(123)));
+
+        let shape_ty = Type::from_cint(types::CShape, 0x1234);
+        assert!(shape_ty.has_value(Const::CShape(crate::cruby::ShapeId(0x1234))));
+        assert!(!shape_ty.has_value(Const::CShape(crate::cruby::ShapeId(0x5678))));
+
+        let ptr = 0x1000 as *const u8;
+        let ptr_ty = Type::from_cptr(ptr);
+        assert!(ptr_ty.has_value(Const::CPtr(ptr)));
+        assert!(!ptr_ty.has_value(Const::CPtr(0x2000 as *const u8)));
+
+        let double_ty = Type::from_double(std::f64::consts::PI);
+        assert!(double_ty.has_value(Const::CDouble(std::f64::consts::PI)));
+        assert!(!double_ty.has_value(Const::CDouble(3.123)));
+
+        let nan_ty = Type::from_double(f64::NAN);
+        assert!(nan_ty.has_value(Const::CDouble(f64::NAN)));
+
+        // Mismatched types
+        assert!(!int8_ty.has_value(Const::CInt16(42)));
+        assert!(!int16_ty.has_value(Const::CInt32(1000)));
+        assert!(!uint8_ty.has_value(Const::CInt8(-1i8)));
+
+        // Wrong specialization (unknown value)
+        assert!(!types::CInt8.has_value(Const::CInt8(42)));
+        assert!(!types::CBool.has_value(Const::CBool(true)));
+        assert!(!types::CShape.has_value(Const::CShape(crate::cruby::ShapeId(0x1234))));
+    }
 }