ZJIT: Optimize class guards by directly reading klass field (#14136)

st0012 · web-flow · commit d25eb1eb5c05 · 2025-08-07T15:38:02.000-07:00
Replace `rb_yarv_class_of` call with:
- a constant check for special constants (nil, fixnums, symbols, etc)
- a check for false
- direct memory read at offset 8 for regular heap objects for the class check
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -1432,6 +1432,30 @@ def test(val) = val.nil?
     }, call_threshold: 2, insns: [:opt_nil_p]
   end
 
+  def test_basic_object_guard_works_with_immediate
+    assert_compiles 'NilClass', %q{
+      class Foo; end
+
+      def test(val) = val.class
+
+      test(Foo.new)
+      test(Foo.new)
+      test(nil)
+    }, call_threshold: 2
+  end
+
+  def test_basic_object_guard_works_with_false
+    assert_compiles 'FalseClass', %q{
+      class Foo; end
+
+      def test(val) = val.class
+
+      test(Foo.new)
+      test(Foo.new)
+      test(false)
+    }, call_threshold: 2
+  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
@@ -1069,17 +1069,29 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
         asm.cmp(val, Qtrue.into());
         asm.jne(side_exit(jit, state, GuardType(guard_type))?);
     } else if guard_type.is_subtype(types::FalseClass) {
-        assert!(Qfalse.as_i64() == 0);
-        asm.test(val, val);
+        asm.cmp(val, Qfalse.into());
         asm.jne(side_exit(jit, state, GuardType(guard_type))?);
+    } else if guard_type.is_immediate() {
+        // All immediate types' guard should have been handled above
+        panic!("unexpected immediate guard type: {guard_type}");
     } else if let Some(expected_class) = guard_type.runtime_exact_ruby_class() {
-        asm_comment!(asm, "guard exact class");
+        asm_comment!(asm, "guard exact class for non-immediate types");
 
-        // Get the class of the value
-        let klass = asm.ccall(rb_yarv_class_of as *const u8, vec![val]);
+        let side_exit = side_exit(jit, state, GuardType(guard_type))?;
+
+        // Check if it's a special constant
+        asm.test(val, (RUBY_IMMEDIATE_MASK as u64).into());
+        asm.jnz(side_exit.clone());
+
+        // Check if it's false
+        asm.cmp(val, Qfalse.into());
+        asm.je(side_exit.clone());
+
+        // Load the class from the object's klass field
+        let klass = asm.load(Opnd::mem(64, val, RUBY_OFFSET_RBASIC_KLASS));
 
         asm.cmp(klass, Opnd::Value(expected_class));
-        asm.jne(side_exit(jit, state, GuardType(guard_type))?);
+        asm.jne(side_exit);
     } else {
         unimplemented!("unsupported type: {guard_type}");
     }
diff --git a/zjit/src/hir_type/mod.rs b/zjit/src/hir_type/mod.rs
@@ -497,7 +497,7 @@ impl Type {
         }
     }
 
-    fn is_immediate(&self) -> bool {
+    pub fn is_immediate(&self) -> bool {
         self.is_subtype(types::Immediate)
     }
 

Original file line number	Diff line number	Diff line change
`@@ -497,7 +497,7 @@ impl Type {`
`497`	`497`	`}`
`498`	`498`	`}`
`499`	`499`
`500`		`- fn is_immediate(&self) -> bool {`
	`500`	`+ pub fn is_immediate(&self) -> bool {`
`501`	`501`	`self.is_subtype(types::Immediate)`
`502`	`502`	`}`
`503`	`503`