ZJIT: Add new ZJIT types for Set

Author: st0012

zjit.c

@@ -32,6 +32,8 @@
 
 #include <errno.h>
 
+RUBY_EXTERN VALUE rb_cSet; // defined in set.c and it's not exposed yet
+
 uint32_t
 rb_zjit_get_page_size(void)
 {

zjit/bindgen/src/main.rs

@@ -186,6 +186,7 @@ fn main() {
         .allowlist_var("rb_cThread")
         .allowlist_var("rb_cArray")
         .allowlist_var("rb_cHash")
+        .allowlist_var("rb_cSet")
         .allowlist_var("rb_cClass")
         .allowlist_var("rb_cISeq")
 

zjit/src/cruby_bindings.inc.rs

@@ -6235,4 +6235,24 @@ mod opt_tests {
               Return v11
         "#]]);
     }
+
+    #[test]
+    fn test_set_type_from_constant() {
+        eval("
+            MY_SET = Set.new
+
+            def test = MY_SET
+
+            test
+            test
+        ");
+        assert_optimized_method_hir("test", expect![[r#"
+            fn test:
+            bb0(v0:BasicObject):
+              PatchPoint SingleRactorMode
+              PatchPoint StableConstantNames(0x1000, MY_SET)
+              v7:SetExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              Return v7
+        "#]]);
+    }
 }

zjit/src/hir.rs

@@ -72,6 +72,7 @@ def base_type name
 base_type "Array"
 base_type "Hash"
 base_type "Range"
+base_type "Set"
 
 (integer, integer_exact) = base_type "Integer"
 # CRuby partitions Integer into immediate and non-immediate variants.

zjit/src/hir_type/gen_hir_type.rb

@@ -1,6 +1,6 @@
 #![allow(non_upper_case_globals)]
 use crate::cruby::{Qfalse, Qnil, Qtrue, VALUE, RUBY_T_ARRAY, RUBY_T_STRING, RUBY_T_HASH};
-use crate::cruby::{rb_cInteger, rb_cFloat, rb_cArray, rb_cHash, rb_cString, rb_cSymbol, rb_cObject, rb_cTrueClass, rb_cFalseClass, rb_cNilClass, rb_cRange};
+use crate::cruby::{rb_cInteger, rb_cFloat, rb_cArray, rb_cHash, rb_cString, rb_cSymbol, rb_cObject, rb_cTrueClass, rb_cFalseClass, rb_cNilClass, rb_cRange, rb_cSet};
 use crate::cruby::ClassRelationship;
 use crate::cruby::get_class_name;
 use crate::cruby::rb_mRubyVMFrozenCore;
@@ -195,6 +195,9 @@ impl Type {
         else if is_string_exact(val) {
             Type { bits: bits::StringExact, spec: Specialization::Object(val) }
         }
+        else if val.class_of() == unsafe { rb_cSet } {
+            Type { bits: bits::SetExact, spec: Specialization::Object(val) }
+        }
         else if val.class_of() == unsafe { rb_cObject } {
             Type { bits: bits::ObjectExact, spec: Specialization::Object(val) }
         }
@@ -394,6 +397,7 @@ impl Type {
         if self.is_subtype(types::NilClassExact) { return Some(unsafe { rb_cNilClass }); }
         if self.is_subtype(types::ObjectExact) { return Some(unsafe { rb_cObject }); }
         if self.is_subtype(types::RangeExact) { return Some(unsafe { rb_cRange }); }
+        if self.is_subtype(types::SetExact) { return Some(unsafe { rb_cSet }); }
         if self.is_subtype(types::StringExact) { return Some(unsafe { rb_cString }); }
         if self.is_subtype(types::SymbolExact) { return Some(unsafe { rb_cSymbol }); }
         if self.is_subtype(types::TrueClassExact) { return Some(unsafe { rb_cTrueClass }); }
@@ -585,6 +589,21 @@ mod tests {
         assert_eq!(types::Integer.inexact_ruby_class(), None);
     }
 
+    #[test]
+    fn set() {
+        assert_subtype(types::SetExact, types::Set);
+        assert_subtype(types::SetSubclass, types::Set);
+    }
+
+    #[test]
+    fn set_has_ruby_class() {
+        crate::cruby::with_rubyvm(|| {
+            assert_eq!(types::SetExact.runtime_exact_ruby_class(), Some(unsafe { rb_cSet }));
+            assert_eq!(types::Set.runtime_exact_ruby_class(), None);
+            assert_eq!(types::SetSubclass.runtime_exact_ruby_class(), None);
+        });
+    }
+
     #[test]
     fn display_exact_bits_match() {
         assert_eq!(format!("{}", Type::fixnum(4)), "Fixnum[4]");