YJIT: Propagate Array, Hash, and String classes

bootstraptest/test_yjit.rb

@@ -4725,3 +4725,31 @@ def test_cases(file, chain)
 
   test_cases(File, Enumerator::Chain)
 }
+
+# singleton class should invalidate Type::CString assumption
+assert_equal 'foo', %q{
+  def define_singleton(str, define)
+    if define
+      # Wrap a C method frame to avoid exiting JIT code on defineclass
+      [nil].reverse_each do
+        class << str
+          def +(_)
+            "foo"
+          end
+        end
+      end
+    end
+    "bar"
+  end
+
+  def entry(define)
+    str = ""
+    # When `define` is false, #+ compiles to rb_str_plus() without a class guard.
+    # When the code is reused with `define` is true, the class of `str` is changed
+    # to a singleton class, so the block should be invalidated.
+    str + define_singleton(str, define)
+  end
+
+  entry(false)
+  entry(true)
+}

class.c

@@ -29,6 +29,7 @@
 #include "internal/variable.h"
 #include "ruby/st.h"
 #include "vm_core.h"
+#include "yjit.h"
 
 /* Flags of T_CLASS
  *
@@ -805,6 +806,7 @@ make_singleton_class(VALUE obj)
     FL_SET(klass, FL_SINGLETON);
     RBASIC_SET_CLASS(obj, klass);
     rb_singleton_class_attached(klass, obj);
+    rb_yjit_invalidate_no_singleton_class(orig_class);
 
     SET_METACLASS_OF(klass, METACLASS_OF(rb_class_real(orig_class)));
     return klass;

common.mk

@@ -3108,6 +3108,7 @@ class.$(OBJEXT): {$(VPATH)}vm_core.h
 class.$(OBJEXT): {$(VPATH)}vm_debug.h
 class.$(OBJEXT): {$(VPATH)}vm_opts.h
 class.$(OBJEXT): {$(VPATH)}vm_sync.h
+class.$(OBJEXT): {$(VPATH)}yjit.h
 compar.$(OBJEXT): $(hdrdir)/ruby/ruby.h
 compar.$(OBJEXT): $(hdrdir)/ruby/version.h
 compar.$(OBJEXT): $(top_srcdir)/internal/basic_operators.h

yjit.h

@@ -46,6 +46,7 @@ void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned ins
 void rb_yjit_tracing_invalidate_all(void);
 void rb_yjit_show_usage(int help, int highlight, unsigned int width, int columns);
 void rb_yjit_lazy_push_frame(const VALUE *pc);
+void rb_yjit_invalidate_no_singleton_class(VALUE klass);
 
 #else
 // !USE_YJIT
@@ -68,6 +69,7 @@ static inline void rb_yjit_before_ractor_spawn(void) {}
 static inline void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx) {}
 static inline void rb_yjit_tracing_invalidate_all(void) {}
 static inline void rb_yjit_lazy_push_frame(const VALUE *pc) {}
+static inline void rb_yjit_invalidate_no_singleton_class(VALUE klass) {}
 
 #endif // #if USE_YJIT
 

yjit/src/codegen.rs

@@ -97,6 +97,9 @@ pub struct JITState {
     /// not been written to for the block to be valid.
     pub stable_constant_names_assumption: Option<*const ID>,
 
+    /// A list of classes that are not supposed to have a singleton class.
+    pub no_singleton_class_assumptions: Vec<VALUE>,
+
     /// When true, the block is valid only when there is a total of one ractor running
     pub block_assumes_single_ractor: bool,
 
@@ -125,6 +128,7 @@ impl JITState {
             method_lookup_assumptions: vec![],
             bop_assumptions: vec![],
             stable_constant_names_assumption: None,
+            no_singleton_class_assumptions: vec![],
             block_assumes_single_ractor: false,
             perf_map: Rc::default(),
             perf_stack: vec![],
@@ -231,6 +235,20 @@ impl JITState {
         Some(())
     }
 
+    /// Assume that objects of a given class will have no singleton class.
+    /// Return true if there has been no such singleton class since boot
+    /// and we can safely invalidate it.
+    pub fn assume_no_singleton_class(&mut self, asm: &mut Assembler, ocb: &mut OutlinedCb, klass: VALUE) -> bool {
+        if jit_ensure_block_entry_exit(self, asm, ocb).is_none() {
+            return false; // out of space, give up
+        }
+        if has_singleton_class_of(klass) {
+            return false; // we've seen a singleton class. disable the optimization to avoid an invalidation loop.
+        }
+        self.no_singleton_class_assumptions.push(klass);
+        true
+    }
+
     fn get_cfp(&self) -> *mut rb_control_frame_struct {
         unsafe { get_ec_cfp(self.ec) }
     }
@@ -1504,7 +1522,7 @@ fn gen_newarray(
     );
 
     asm.stack_pop(n.as_usize());
-    let stack_ret = asm.stack_push(Type::TArray);
+    let stack_ret = asm.stack_push(Type::CArray);
     asm.mov(stack_ret, new_ary);
 
     Some(KeepCompiling)
@@ -1527,7 +1545,7 @@ fn gen_duparray(
         vec![ary.into()],
     );
 
-    let stack_ret = asm.stack_push(Type::TArray);
+    let stack_ret = asm.stack_push(Type::CArray);
     asm.mov(stack_ret, new_ary);
 
     Some(KeepCompiling)
@@ -1547,7 +1565,7 @@ fn gen_duphash(
     // call rb_hash_resurrect(VALUE hash);
     let hash = asm.ccall(rb_hash_resurrect as *const u8, vec![hash.into()]);
 
-    let stack_ret = asm.stack_push(Type::THash);
+    let stack_ret = asm.stack_push(Type::CHash);
     asm.mov(stack_ret, hash);
 
     Some(KeepCompiling)
@@ -2303,12 +2321,12 @@ fn gen_newhash(
         asm.cpop_into(new_hash); // x86 alignment
 
         asm.stack_pop(num.try_into().unwrap());
-        let stack_ret = asm.stack_push(Type::THash);
+        let stack_ret = asm.stack_push(Type::CHash);
         asm.mov(stack_ret, new_hash);
     } else {
         // val = rb_hash_new();
         let new_hash = asm.ccall(rb_hash_new as *const u8, vec![]);
-        let stack_ret = asm.stack_push(Type::THash);
+        let stack_ret = asm.stack_push(Type::CHash);
         asm.mov(stack_ret, new_hash);
     }
 
@@ -2330,7 +2348,7 @@ fn gen_putstring(
         vec![EC, put_val.into(), 0.into()]
     );
 
-    let stack_top = asm.stack_push(Type::TString);
+    let stack_top = asm.stack_push(Type::CString);
     asm.mov(stack_top, str_opnd);
 
     Some(KeepCompiling)
@@ -2351,7 +2369,7 @@ fn gen_putchilledstring(
         vec![EC, put_val.into(), 1.into()]
     );
 
-    let stack_top = asm.stack_push(Type::TString);
+    let stack_top = asm.stack_push(Type::CString);
     asm.mov(stack_top, str_opnd);
 
     Some(KeepCompiling)
@@ -4493,8 +4511,18 @@ fn jit_guard_known_klass(
     let val_type = asm.ctx.get_opnd_type(insn_opnd);
 
     if val_type.known_class() == Some(known_klass) {
-        // We already know from type information that this is a match
-        return;
+        // Unless frozen, Array, Hash, and String objects may change their RBASIC_CLASS
+        // when they get a singleton class. Those types need invalidations.
+        if unsafe { [rb_cArray, rb_cHash, rb_cString].contains(&known_klass) } {
+            if jit.assume_no_singleton_class(asm, ocb, known_klass) {
+                // Speculate that this object will not have a singleton class,
+                // and invalidate the block in case it does.
+                return;
+            }
+        } else {
+            // We already know from type information that this is a match
+            return;
+        }
     }
 
     if unsafe { known_klass == rb_cNilClass } {
@@ -4613,14 +4641,11 @@ fn jit_guard_known_klass(
         jit_chain_guard(JCC_JNE, jit, asm, ocb, max_chain_depth, counter);
 
         if known_klass == unsafe { rb_cString } {
-            // Upgrading to Type::CString here is incorrect.
-            // The guard we put only checks RBASIC_CLASS(obj),
-            // which adding a singleton class can change. We
-            // additionally need to know the string is frozen
-            // to claim Type::CString.
-            asm.ctx.upgrade_opnd_type(insn_opnd, Type::TString);
+            asm.ctx.upgrade_opnd_type(insn_opnd, Type::CString);
         } else if known_klass == unsafe { rb_cArray } {
-            asm.ctx.upgrade_opnd_type(insn_opnd, Type::TArray);
+            asm.ctx.upgrade_opnd_type(insn_opnd, Type::CArray);
+        } else if known_klass == unsafe { rb_cHash } {
+            asm.ctx.upgrade_opnd_type(insn_opnd, Type::CHash);
         }
     }
 }

yjit/src/core.rs

@@ -52,20 +52,18 @@ pub enum Type {
     Flonum,
     ImmSymbol,
 
-    #[allow(unused)]
-    HeapSymbol,
-
     TString, // An object with the T_STRING flag set, possibly an rb_cString
     CString, // An un-subclassed string of type rb_cString (can have instance vars in some cases)
     TArray, // An object with the T_ARRAY flag set, possibly an rb_cArray
+    CArray, // An un-subclassed array of type rb_cArray (can have instance vars in some cases)
     THash, // An object with the T_HASH flag set, possibly an rb_cHash
+    CHash, // An un-subclassed hash of type rb_cHash (can have instance vars in some cases)
 
     BlockParamProxy, // A special sentinel value indicating the block parameter should be read from
                      // the current surrounding cfp
 
     // The context currently relies on types taking at most 4 bits (max value 15)
-    // to encode, so if we add two more, we will need to refactor the context,
-    // or we could remove HeapSymbol, which is currently unused.
+    // to encode, so if we add any more, we will need to refactor the context.
 }
 
 // Default initialization
@@ -98,8 +96,11 @@ impl Type {
             // Core.rs can't reference rb_cString because it's linked by Rust-only tests.
             // But CString vs TString is only an optimisation and shouldn't affect correctness.
             #[cfg(not(test))]
-            if val.class_of() == unsafe { rb_cString } && val.is_frozen() {
-                return Type::CString;
+            match val.class_of() {
+                class if class == unsafe { rb_cArray }  => return Type::CArray,
+                class if class == unsafe { rb_cHash }   => return Type::CHash,
+                class if class == unsafe { rb_cString } => return Type::CString,
+                _ => {}
             }
             // We likewise can't reference rb_block_param_proxy, but it's again an optimisation;
             // we can just treat it as a normal Object.
@@ -150,8 +151,9 @@ impl Type {
         match self {
             Type::UnknownHeap => true,
             Type::TArray => true,
+            Type::CArray => true,
             Type::THash => true,
-            Type::HeapSymbol => true,
+            Type::CHash => true,
             Type::TString => true,
             Type::CString => true,
             Type::BlockParamProxy => true,
@@ -161,21 +163,17 @@ impl Type {
 
     /// Check if it's a T_ARRAY object (both TArray and CArray are T_ARRAY)
     pub fn is_array(&self) -> bool {
-        matches!(self, Type::TArray)
+        matches!(self, Type::TArray | Type::CArray)
     }
 
-    /// Check if it's a T_HASH object
+    /// Check if it's a T_HASH object (both THash and CHash are T_HASH)
     pub fn is_hash(&self) -> bool {
-        matches!(self, Type::THash)
+        matches!(self, Type::THash | Type::CHash)
     }
 
     /// Check if it's a T_STRING object (both TString and CString are T_STRING)
     pub fn is_string(&self) -> bool {
-        match self {
-            Type::TString => true,
-            Type::CString => true,
-            _ => false,
-        }
+        matches!(self, Type::TString | Type::CString)
     }
 
     /// Returns an Option with the T_ value type if it is known, otherwise None
@@ -186,9 +184,9 @@ impl Type {
             Type::False => Some(RUBY_T_FALSE),
             Type::Fixnum => Some(RUBY_T_FIXNUM),
             Type::Flonum => Some(RUBY_T_FLOAT),
-            Type::TArray => Some(RUBY_T_ARRAY),
-            Type::THash => Some(RUBY_T_HASH),
-            Type::ImmSymbol | Type::HeapSymbol => Some(RUBY_T_SYMBOL),
+            Type::TArray | Type::CArray => Some(RUBY_T_ARRAY),
+            Type::THash | Type::CHash => Some(RUBY_T_HASH),
+            Type::ImmSymbol => Some(RUBY_T_SYMBOL),
             Type::TString | Type::CString => Some(RUBY_T_STRING),
             Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None,
             Type::BlockParamProxy => None,
@@ -204,7 +202,9 @@ impl Type {
                 Type::False => Some(rb_cFalseClass),
                 Type::Fixnum => Some(rb_cInteger),
                 Type::Flonum => Some(rb_cFloat),
-                Type::ImmSymbol | Type::HeapSymbol => Some(rb_cSymbol),
+                Type::ImmSymbol => Some(rb_cSymbol),
+                Type::CArray => Some(rb_cArray),
+                Type::CHash => Some(rb_cHash),
                 Type::CString => Some(rb_cString),
                 _ => None,
             }
@@ -255,6 +255,16 @@ impl Type {
             return TypeDiff::Compatible(1);
         }
 
+        // A CArray is also a TArray.
+        if self == Type::CArray && dst == Type::TArray {
+            return TypeDiff::Compatible(1);
+        }
+
+        // A CHash is also a THash.
+        if self == Type::CHash && dst == Type::THash {
+            return TypeDiff::Compatible(1);
+        }
+
         // A CString is also a TString.
         if self == Type::CString && dst == Type::TString {
             return TypeDiff::Compatible(1);
@@ -1644,6 +1654,9 @@ impl JITState {
         if let Some(idlist) = self.stable_constant_names_assumption {
             track_stable_constant_names_assumption(blockref, idlist);
         }
+        for klass in self.no_singleton_class_assumptions {
+            track_no_singleton_class_assumption(blockref, klass);
+        }
 
         blockref
     }