ZJIT: Compile IsA into load + compare for String/Array/Hash (#15878)

Resolves https://github.com/Shopify/ruby/issues/880

Implemented this by using the code generation for `GuardType` as a reference.
Not sure if this is the best way to go about it, but it seems to work.

2 files changed, 94 insertions, 1 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index ad2df806d5..2cfe9dd7e3 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -4417,6 +4417,60 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 14, num_profiles: 5
   end
 
+  def test_is_a_string_special_case
+    assert_compiles '[true, false, false, false, false, false]', %q{
+      def test(x)
+        x.is_a?(String)
+      end
+      test("foo")
+      [test("bar"), test(1), test(false), test(:foo), test([]), test({})]
+    }
+  end
+
+  def test_is_a_array_special_case
+    assert_compiles '[true, true, false, false, false, false, false]', %q{
+      def test(x)
+        x.is_a?(Array)
+      end
+      test([])
+      [test([1,2,3]), test([]), test(1), test(false), test(:foo), test("foo"), test({})]
+    }
+  end
+
+  def test_is_a_hash_special_case
+    assert_compiles '[true, true, false, false, false, false, false]', %q{
+      def test(x)
+        x.is_a?(Hash)
+      end
+      test({})
+      [test({:a => "b"}), test({}), test(1), test(false), test(:foo), test([]), test("foo")]
+    }
+  end
+
+  def test_is_a_hash_subclass
+    assert_compiles 'true', %q{
+      class MyHash < Hash
+      end
+      def test(x)
+        x.is_a?(Hash)
+      end
+      test({})
+      test(MyHash.new)
+    }
+  end
+
+  def test_is_a_normal_case
+    assert_compiles '[true, false]', %q{
+      class MyClass
+      end
+      def test(x)
+        x.is_a?(MyClass)
+      end
+      test("a")
+      [test(MyClass.new), test("a")]
+    }
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 7afcff5863..0ae85c24a2 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -1743,7 +1743,46 @@ fn gen_dup_array_include(
 }
 
 fn gen_is_a(asm: &mut Assembler, obj: Opnd, class: Opnd) -> lir::Opnd {
-    asm_ccall!(asm, rb_obj_is_kind_of, obj, class)
+    let builtin_type = match class {
+        Opnd::Value(value) if value == unsafe { rb_cString } => Some(RUBY_T_STRING),
+        Opnd::Value(value) if value == unsafe { rb_cArray } => Some(RUBY_T_ARRAY),
+        Opnd::Value(value) if value == unsafe { rb_cHash } => Some(RUBY_T_HASH),
+        _ => None
+    };
+
+    if let Some(builtin_type) = builtin_type {
+        asm_comment!(asm, "IsA by matching builtin type");
+        let ret_label = asm.new_label("is_a_ret");
+        let false_label = asm.new_label("is_a_false");
+
+        let val = match obj {
+            Opnd::Reg(_) | Opnd::VReg { .. } => obj,
+            _ => asm.load(obj),
+        };
+
+        // Check special constant
+        asm.test(val, Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
+        asm.jnz(ret_label.clone());
+
+        // Check false
+        asm.cmp(val, Qfalse.into());
+        asm.je(false_label.clone());
+
+        let flags = asm.load(Opnd::mem(VALUE_BITS, val, RUBY_OFFSET_RBASIC_FLAGS));
+        let obj_builtin_type = asm.and(flags, Opnd::UImm(RUBY_T_MASK as u64));
+        asm.cmp(obj_builtin_type, Opnd::UImm(builtin_type as u64));
+        asm.jmp(ret_label.clone());
+
+        // If we get here then the value was false, unset the Z flag
+        // so that csel_e will select false instead of true
+        asm.write_label(false_label);
+        asm.test(Opnd::UImm(1), Opnd::UImm(1));
+
+        asm.write_label(ret_label);
+        asm.csel_e(Qtrue.into(), Qfalse.into())
+    } else {
+        asm_ccall!(asm, rb_obj_is_kind_of, obj, class)
+    }
 }
 
 /// Compile a new hash instruction