YJIT: Propagate Array, Hash, and String classes (#10323)

bootstraptest/test_yjit.rb

@@ -4725,3 +4725,31 @@ assert_equal '["", "1/2", [0, [:ok, 1]]]', %q{
 
   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

@@ -3106,6 +3106,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
     }

yjit/src/invariants.rs

@@ -53,6 +53,12 @@ pub struct Invariants {
     /// A map from a block to a set of IDs that it is assuming have not been
     /// redefined.
     block_constant_states: HashMap<BlockRef, HashSet<ID>>,
+
+    /// A map from a class to a set of blocks that assume objects of the class
+    /// will have no singleton class. When the set is empty, it means that
+    /// there has been a singleton class for the class after boot, so you cannot
+    /// assume no singleton class going forward.
+    no_singleton_classes: HashMap<VALUE, HashSet<BlockRef>>,
 }
 
 /// Private singleton instance of the invariants global struct.
@@ -69,6 +75,7 @@ impl Invariants {
                 single_ractor: HashSet::new(),
                 constant_state_blocks: HashMap::new(),
                 block_constant_states: HashMap::new(),
+                no_singleton_classes: HashMap::new(),
             });
         }
     }
@@ -130,6 +137,23 @@ pub fn track_method_lookup_stability_assumption(
         .insert(uninit_block);
 }
 
+/// Track that a block will assume that `klass` objects will have no singleton class.
+pub fn track_no_singleton_class_assumption(uninit_block: BlockRef, klass: VALUE) {
+    Invariants::get_instance()
+        .no_singleton_classes
+        .entry(klass)
+        .or_default()
+        .insert(uninit_block);
+}
+
+/// Returns true if we've seen a singleton class of a given class since boot.
+pub fn has_singleton_class_of(klass: VALUE) -> bool {
+    Invariants::get_instance()
+        .no_singleton_classes
+        .get(&klass)
+        .map_or(false, |blocks| blocks.is_empty())
+}
+
 // Checks rb_method_basic_definition_p and registers the current block for invalidation if method
 // lookup changes.
 // A "basic method" is one defined during VM boot, so we can use this to check assumptions based on
@@ -391,6 +415,11 @@ pub fn block_assumptions_free(blockref: BlockRef) {
     if invariants.constant_state_blocks.is_empty() {
         invariants.constant_state_blocks.shrink_to_fit();
     }
+
+    // Remove tracking for blocks assumping no singleton class
+    for (_, blocks) in invariants.no_singleton_classes.iter_mut() {
+        blocks.remove(&blockref);
+    }
 }
 
 /// Callback from the opt_setinlinecache instruction in the interpreter.
@@ -457,6 +486,35 @@ pub extern "C" fn rb_yjit_constant_ic_update(iseq: *const rb_iseq_t, ic: IC, ins
     });
 }
 
+/// Invalidate blocks that assume objects of a given class will have no singleton class.
+#[no_mangle]
+pub extern "C" fn rb_yjit_invalidate_no_singleton_class(klass: VALUE) {
+    // Skip tracking singleton classes during boot. Such objects already have a singleton class
+    // before entering JIT code, so they get rejected when they're checked for the first time.
+    if unsafe { INVARIANTS.is_none() } {
+        return;
+    }
+
+    // We apply this optimization only to Array, Hash, and String for now.
+    if unsafe { [rb_cArray, rb_cHash, rb_cString].contains(&klass) } {
+        let no_singleton_classes = &mut Invariants::get_instance().no_singleton_classes;
+        match no_singleton_classes.get_mut(&klass) {
+            Some(blocks) => {
+                // Invalidate existing blocks and let has_singleton_class_of()
+                // return true when they are compiled again
+                for block in mem::take(blocks) {
+                    invalidate_block_version(&block);
+                    incr_counter!(invalidate_no_singleton_class);
+                }
+            }
+            None => {
+                // Let has_singleton_class_of() return true for this class
+                no_singleton_classes.insert(klass, HashSet::new());
+            }
+        }
+    }
+}
+
 // Invalidate all generated code and patch C method return code to contain
 // logic for firing the c_return TracePoint event. Once rb_vm_barrier()
 // returns, all other ractors are pausing inside RB_VM_LOCK_ENTER(), which

yjit/src/stats.rs

@@ -261,7 +261,7 @@ macro_rules! make_counters {
 
 /// The list of counters that are available without --yjit-stats.
 /// They are incremented only by `incr_counter!` and don't use `gen_counter_incr`.
-pub const DEFAULT_COUNTERS: [Counter; 15] = [
+pub const DEFAULT_COUNTERS: [Counter; 16] = [
     Counter::code_gc_count,
     Counter::compiled_iseq_entry,
     Counter::cold_iseq_entry,
@@ -278,6 +278,7 @@ pub const DEFAULT_COUNTERS: [Counter; 15] = [
     Counter::invalidate_ractor_spawn,
     Counter::invalidate_constant_state_bump,
     Counter::invalidate_constant_ic_fill,
+    Counter::invalidate_no_singleton_class,
 ];
 
 /// Macro to increase a counter by name and count
@@ -559,6 +560,7 @@ make_counters! {
     invalidate_ractor_spawn,
     invalidate_constant_state_bump,
     invalidate_constant_ic_fill,
+    invalidate_no_singleton_class,
 
     // Currently, it's out of the ordinary (might be impossible) for YJIT to leave gaps in
     // executable memory, so this should be 0.