[WIP] NoSingletonOverride patchpoint. New functionality, same old NoSingleton name

Author: XrXr

class.c

@@ -24,6 +24,7 @@
 #include "internal/box.h"
 #include "internal/class.h"
 #include "internal/eval.h"
+#include "internal/error.h"
 #include "internal/hash.h"
 #include "internal/object.h"
 #include "internal/string.h"
@@ -55,6 +56,9 @@
  * 2:    RCLASS_PRIME_CLASSEXT_WRITABLE
  *           This module's prime classext is the only classext and writable from any boxes.
  *           If unset, the prime classext is writable only from the root box.
+ * 3:    RICLASS_IS_SINGLETON_ANCESTOR
+ *           This iclass is in the singleton class portion of an ancestry chain.
+ *           Set when a module is included/prepended into a singleton class.
  * 4:    RCLASS_BOXABLE
  *           Is a builtin class that may be boxed. It larger than a normal class.
  */
@@ -1422,7 +1426,6 @@ make_singleton_class(VALUE obj)
     RBASIC_SET_CLASS(obj, klass);
     rb_singleton_class_attached(klass, obj);
     rb_yjit_invalidate_no_singleton_class(orig_class);
-    rb_zjit_invalidate_no_singleton_class(orig_class);
 
     SET_METACLASS_OF(klass, METACLASS_OF(rb_class_real(orig_class)));
     return klass;
@@ -1855,6 +1858,18 @@ clear_module_cache_i(ID id, VALUE val, void *data)
     return ID_TABLE_CONTINUE;
 }
 
+static enum rb_id_table_iterator_result
+zjit_invalid_no_singleton_override_i(ID method_id, VALUE val, void *data)
+{
+    VALUE insertion_point = (VALUE)data;
+    // If this method is an override, bust invariant and stop iterating
+    if (rb_method_entry(insertion_point, method_id)) {
+        rb_zjit_invalidate_no_singleton_class(rb_class_real(insertion_point));
+        return ID_TABLE_STOP;
+    }
+    return ID_TABLE_CONTINUE;
+}
+
 static bool
 module_in_super_chain(const VALUE klass, VALUE module)
 {
@@ -1944,6 +1959,15 @@ do_include_modules_at(const VALUE klass, VALUE c, VALUE module, int search_super
 
         // setup T_ICLASS for the include/prepend module
         iclass = rb_include_class_new(module, super_class);
+        // Mark the iclass as being in the singleton portion of the chain.
+        // c is the insertion point; if it's a singleton class or already in
+        // the singleton portion, the new iclass is too.
+        if (rb_zjit_enabled_p && (RCLASS_SINGLETON_P(c) || (RB_TYPE_P(c, T_ICLASS) && RICLASS_SINGLETON_ANCESTOR_P(c)))) {
+            FL_SET(iclass, RICLASS_IS_SINGLETON_ANCESTOR);
+            if (method_changed && tbl && rb_id_table_size(tbl)) {
+                rb_id_table_foreach(tbl, zjit_invalid_no_singleton_override_i, (void *)c);
+            }
+        }
         c = rb_class_set_super(c, iclass);
         RCLASS_SET_INCLUDER(iclass, klass);
         if (module != RCLASS_ORIGIN(module)) {
@@ -2082,6 +2106,10 @@ rb_prepend_module(VALUE klass, VALUE module)
                     rb_id_table_foreach(RCLASS_M_TBL(subclass), clear_module_cache_i, (void *)subclass);
                     RCLASS_WRITE_M_TBL(subclass, klass_m_tbl);
                     VALUE origin = rb_include_class_new(klass_origin, RCLASS_SUPER(subclass));
+                    if (rb_zjit_enabled_p && RICLASS_SINGLETON_ANCESTOR_P(subclass)) {
+                        RUBY_ASSERT_BUILTIN_TYPE(subclass, T_ICLASS);
+                        FL_SET(origin, RICLASS_IS_SINGLETON_ANCESTOR);
+                    }
                     rb_class_set_super(subclass, origin);
                     RCLASS_SET_INCLUDER(origin, RCLASS_INCLUDER(subclass));
                     RCLASS_WRITE_ORIGIN(subclass, origin);

internal/class.h

@@ -293,6 +293,8 @@ static inline void RCLASS_WRITE_CLASSPATH(VALUE klass, VALUE classpath, bool per
 #define RCLASS_PRIME_CLASSEXT_WRITABLE FL_USER2
 #define RCLASS_IS_INITIALIZED FL_USER3
 // 3 is RMODULE_IS_REFINEMENT for RMODULE
+// 3 is also RICLASS_IS_SINGLETON_ANCESTOR for T_ICLASS (not T_CLASS/T_MODULE)
+#define RICLASS_IS_SINGLETON_ANCESTOR FL_USER3
 #define RCLASS_BOXABLE FL_USER4
 
 static inline st_table *
@@ -708,6 +710,12 @@ RICLASS_OWNS_M_TBL_P(VALUE iclass)
     return RCLASSEXT_ICLASS_IS_ORIGIN(ext) && !RCLASSEXT_ICLASS_ORIGIN_SHARED_MTBL(ext);
 }
 
+static inline bool
+RICLASS_SINGLETON_ANCESTOR_P(VALUE iclass)
+{
+    return FL_TEST_RAW(iclass, RICLASS_IS_SINGLETON_ANCESTOR);
+}
+
 static inline void
 RCLASS_SET_INCLUDER(VALUE iclass, VALUE klass)
 {

vm_method.c

@@ -1355,6 +1355,17 @@ check_override_opt_method(VALUE klass, VALUE mid)
     }
 }
 
+static void
+check_singleton_override_by_method_i(VALUE subclass, VALUE arg)
+{
+    ID mid = arg;
+    VM_ASSERT_TYPE(subclass, T_ICLASS);
+    if (RICLASS_SINGLETON_ANCESTOR_P(subclass) && rb_method_entry(subclass, mid)) {
+        // TODO fuse real and search above
+        rb_zjit_invalidate_no_singleton_class(rb_class_real(subclass));
+    }
+}
+
 static inline rb_method_entry_t* search_method0(VALUE klass, ID id, VALUE *defined_class_ptr, bool skip_refined);
 /*
  * klass->method_table[mid] = method_entry(defined_class, visi, def)
@@ -1485,6 +1496,16 @@ rb_method_entry_make(VALUE klass, ID mid, VALUE defined_class, rb_method_visibil
             rb_warn("redefining '%s' may cause serious problems", rb_id2name(mid));
         }
     }
+    // Check ZJIT NoSingletonClassOverride invariant
+    if (rb_zjit_enabled_p) {
+        // Invalidate ZJIT code that assumes no singleton class override
+        if (RCLASS_SINGLETON_P(klass) && rb_method_entry(klass, mid)) {
+            rb_zjit_invalidate_no_singleton_class(rb_class_real(klass));
+        }
+        else if (RB_TYPE_P(klass, T_MODULE)) {
+            rb_class_foreach_subclass(klass, check_singleton_override_by_method_i, (VALUE)mid);
+        }
+    }
 
     if (make_refined) {
         make_method_entry_refined(klass, me);

zjit/src/codegen.rs

@@ -11,7 +11,7 @@ use crate::backend::current::ALLOC_REGS;
 use crate::invariants::{
     track_bop_assumption, track_cme_assumption, track_no_ep_escape_assumption, track_no_trace_point_assumption,
     track_single_ractor_assumption, track_stable_constant_names_assumption, track_no_singleton_class_assumption,
-    track_root_box_assumption
+    track_no_singleton_override_assumption, track_root_box_assumption
 };
 use crate::gc::append_gc_offsets;
 use crate::payload::{get_or_create_iseq_payload, IseqCodePtrs, IseqVersion, IseqVersionRef, IseqStatus};
@@ -877,6 +877,9 @@ pub fn split_patch_point(asm: &mut Assembler, target: &Target, invariant: Invari
             Invariant::NoSingletonClass { klass } => {
                 track_no_singleton_class_assumption(klass, code_ptr, side_exit_ptr, version);
             }
+            Invariant::NoSingletonClassOverride { klass } => {
+                track_no_singleton_override_assumption(klass, code_ptr, side_exit_ptr, version);
+            }
             Invariant::RootBoxOnly => {
                 track_root_box_assumption(code_ptr, side_exit_ptr, version);
             }

zjit/src/hir.rs

@@ -153,6 +153,12 @@ pub enum Invariant {
     NoSingletonClass {
         klass: VALUE,
     },
+    /// No method override has occurred in the singleton portion of this class's ancestry.
+    /// Invalidated when a method is added to a singleton class or a module is
+    /// included/prepended into the singleton chain that introduces methods.
+    NoSingletonClassOverride {
+        klass: VALUE,
+    },
     /// Only the root box is active, so we can safely read from the prime classext.
     /// Invalidated if a non-root box duplicates any classext.
     RootBoxOnly,
@@ -295,6 +301,12 @@ impl<'a> std::fmt::Display for InvariantPrinter<'a> {
                     class_name,
                     self.ptr_map.map_ptr(klass.as_ptr::<VALUE>()))
             }
+            Invariant::NoSingletonClassOverride { klass } => {
+                let class_name = get_class_name(klass);
+                write!(f, "NoSingletonClassOverride({}@{:p})",
+                    class_name,
+                    self.ptr_map.map_ptr(klass.as_ptr::<VALUE>()))
+            }
             Invariant::RootBoxOnly => write!(f, "RootBoxOnly"),
         }
     }