Merge pull request ruby#91 from jhawthorn/singleton_typecheck

maximecb · web-flow · commit 6698e433934d · 2021-06-17T16:29:03.000-04:00
Simplify known class check for singleton classes
diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
@@ -1153,3 +1153,34 @@ def foo(foo, bar)
   foo(5,2)
   foo(5,2)
 }
+
+# Call to object with singleton
+assert_equal '123', %q{
+  obj = Object.new
+  def obj.foo
+    123
+  end
+
+  def foo(obj)
+    obj.foo()
+  end
+
+  foo(obj)
+  foo(obj)
+}
+
+# Call to singleton class
+assert_equal '123', %q{
+  class Foo
+    def self.foo
+      123
+    end
+  end
+
+  def foo(obj)
+    obj.foo()
+  end
+
+  foo(Foo)
+  foo(Foo)
+}
diff --git a/yjit_codegen.c b/yjit_codegen.c
@@ -2170,6 +2170,7 @@ static bool
 jit_guard_known_klass(jitstate_t *jit, ctx_t* ctx, VALUE known_klass, insn_opnd_t insn_opnd, const int max_chain_depth, uint8_t *side_exit)
 {
     val_type_t val_type = ctx_get_opnd_type(ctx, insn_opnd);
+    bool singleton_klass = FL_TEST(known_klass, FL_SINGLETON);
 
     if (known_klass == rb_cNilClass) {
         if (val_type.type != ETYPE_NIL) {
@@ -2188,7 +2189,6 @@ jit_guard_known_klass(jitstate_t *jit, ctx_t* ctx, VALUE known_klass, insn_opnd_
 
             ctx_set_opnd_type(ctx, insn_opnd, TYPE_TRUE);
         }
-
     }
     else if (known_klass == rb_cFalseClass) {
         if (val_type.type != ETYPE_FALSE) {
@@ -2200,15 +2200,26 @@ jit_guard_known_klass(jitstate_t *jit, ctx_t* ctx, VALUE known_klass, insn_opnd_
             ctx_set_opnd_type(ctx, insn_opnd, TYPE_FALSE);
         }
     }
-    else {
-        // Can't guard for for these classes because some of they are sometimes immediate (special const).
-        // Can remove this by adding appropriate dynamic checks.
-        if (known_klass == rb_cInteger ||
+    else if (known_klass == rb_cInteger ||
             known_klass == rb_cSymbol ||
             known_klass == rb_cFloat) {
-            return false;
-        }
-
+        // Can't guard for for these classes because some of they are sometimes
+        // immediate (special const). Can remove this by adding appropriate
+        // dynamic checks.
+        return false;
+    }
+    else if (singleton_klass) {
+        // Singleton classes are attached to one specific object, so we can
+        // avoid one memory access (and potentially the is_heap check) by
+        // looking for the expected object directly.
+        ADD_COMMENT(cb, "guard known object with singleton class");
+        VALUE known_obj = rb_attr_get(known_klass, id__attached__);
+        // TODO: jit_mov_gc_ptr keeps a strong reference, which leaks the object.
+        jit_mov_gc_ptr(jit, cb, REG1, known_obj);
+        cmp(cb, REG0, REG1);
+        jit_chain_guard(JCC_JNE, jit, ctx, max_chain_depth, side_exit);
+    }
+    else {
         // Check that the receiver is a heap object
         // Note: if we get here, the class doesn't have immediate instances.
         if (!val_type.is_heap) {
@@ -2232,7 +2243,6 @@ jit_guard_known_klass(jitstate_t *jit, ctx_t* ctx, VALUE known_klass, insn_opnd_
         jit_chain_guard(JCC_JNE, jit, ctx, max_chain_depth, side_exit);
     }
 
-    // Pointer to the klass field of the receiver &(recv->klass)
     return true;
 }
 
