ZJIT: Optimize ObjToString with type guards (ruby#14469)

altxtech · web-flow · commit f7fe43610b75 · 2025-09-09T08:58:03.000-07:00
* failing test for ObjToString optimization with GuardType

* profile ObjToString receiver and rewrite with guard

* adjust integration tests for objtostring type guard optimization

* Implement new GuardTypeNot HIR; objtostring sends to_s directly on profiled nonstrings

* codegen for GuardTypeNot

* typo fixes

* better name for tests; fix side exit reason for GuardTypeNot

* revert accidental change

* make bindgen

* Fix is_string to identify subclasses of String; fix codegen for identifying if val is String
diff --git a/insns.def b/insns.def
@@ -939,6 +939,7 @@ objtostring
 (VALUE recv)
 (VALUE val)
 // attr bool leaf = false;
+// attr bool zjit_profile = true;
 {
     val = vm_objtostring(GET_ISEQ(), recv, cd);
 
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -2088,6 +2088,61 @@ def test(str)
     }
   end
 
+  def test_objtostring_profiled_string_fastpath
+    assert_compiles '"foo"', %q{
+      def test(str)
+        "#{str}"
+      end
+      test('foo'); test('foo') # profile as string
+    }, call_threshold: 2
+  end
+
+  def test_objtostring_profiled_string_subclass_fastpath
+    assert_compiles '"foo"', %q{
+      class MyString < String; end
+
+      def test(str)
+        "#{str}"
+      end
+
+      foo = MyString.new("foo")
+      test(foo); test(foo) # still profiles as string
+    }, call_threshold: 2
+  end
+
+  def test_objtostring_profiled_string_fastpath_exits_on_nonstring
+    assert_compiles '"1"', %q{
+      def test(str)
+        "#{str}"
+      end
+
+      test('foo') # profile as string
+      test(1)
+    }, call_threshold: 2
+  end
+
+  def test_objtostring_profiled_nonstring_calls_to_s
+    assert_compiles '"[1, 2, 3]"', %q{
+      def test(str)
+        "#{str}"
+      end
+
+      test([1,2,3]); # profile as nonstring
+      test([1,2,3]);
+    }, call_threshold: 2
+  end
+
+  def test_objtostring_profiled_nonstring_guard_exits_when_string
+    assert_compiles '"foo"', %q{
+      def test(str)
+        "#{str}"
+      end
+
+      test([1,2,3]); # profiles as nonstring
+      test('foo');
+    }, call_threshold: 2
+  end
+
   def test_string_bytesize_with_guard
     assert_compiles '5', %q{
       def test(str)
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -389,6 +389,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::IsMethodCfunc { val, cd, cfunc } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc),
         Insn::Test { val } => gen_test(asm, opnd!(val)),
         Insn::GuardType { val, guard_type, state } => gen_guard_type(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state)),
+        Insn::GuardTypeNot { val, guard_type, state } => gen_guard_type_not(jit, asm, opnd!(val), *guard_type, &function.frame_state(*state)),
         Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(jit, asm, opnd!(val), *expected, &function.frame_state(*state)),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(asm, *cfun, opnds!(args)),
@@ -1375,6 +1376,26 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
 
         asm.cmp(klass, Opnd::Value(expected_class));
         asm.jne(side_exit);
+    } else if guard_type.is_subtype(types::String) {
+        let side = side_exit(jit, state, GuardType(guard_type));
+
+        // Check special constant
+        asm.test(val, Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
+        asm.jnz(side.clone());
+
+        // Check false
+        asm.cmp(val, Qfalse.into());
+        asm.je(side.clone());
+
+        let val = match val {
+            Opnd::Reg(_) | Opnd::VReg { .. } => val,
+            _ => asm.load(val),
+        };
+
+        let flags = asm.load(Opnd::mem(VALUE_BITS, val, RUBY_OFFSET_RBASIC_FLAGS));
+        let tag   = asm.and(flags, Opnd::UImm(RUBY_T_MASK as u64));
+        asm.cmp(tag, Opnd::UImm(RUBY_T_STRING as u64));
+        asm.jne(side);
     } else if guard_type.bit_equal(types::HeapObject) {
         let side_exit = side_exit(jit, state, GuardType(guard_type));
         asm.cmp(val, Opnd::Value(Qfalse));
@@ -1387,6 +1408,38 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
     val
 }
 
+fn gen_guard_type_not(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard_type: Type, state: &FrameState) -> lir::Opnd {
+    if guard_type.is_subtype(types::String) {
+        // We only exit if val *is* a String. Otherwise we fall through.
+        let cont = asm.new_label("guard_type_not_string_cont");
+        let side = side_exit(jit, state, GuardTypeNot(guard_type));
+
+        // Continue if special constant (not string)
+        asm.test(val, Opnd::UImm(RUBY_IMMEDIATE_MASK as u64));
+        asm.jnz(cont.clone());
+
+        // Continue if false (not string)
+        asm.cmp(val, Qfalse.into());
+        asm.je(cont.clone());
+
+        let val = match val {
+            Opnd::Reg(_) | Opnd::VReg { .. } => val,
+            _ => asm.load(val),
+        };
+
+        let flags = asm.load(Opnd::mem(VALUE_BITS, val, RUBY_OFFSET_RBASIC_FLAGS));
+        let tag   = asm.and(flags, Opnd::UImm(RUBY_T_MASK as u64));
+        asm.cmp(tag, Opnd::UImm(RUBY_T_STRING as u64));
+        asm.je(side);
+
+        // Otherwise (non-string heap object), continue.
+        asm.write_label(cont);
+    } else {
+        unimplemented!("unsupported type: {guard_type}");
+    }
+    val
+}
+
 /// Compile an identity check with a side exit
 fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: VALUE, state: &FrameState) -> lir::Opnd {
     asm.cmp(val, Opnd::Value(expected));
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -455,6 +455,7 @@ pub enum SideExitReason {
     FixnumSubOverflow,
     FixnumMultOverflow,
     GuardType(Type),
+    GuardTypeNot(Type),
     GuardShape(ShapeId),
     GuardBitEquals(VALUE),
     PatchPoint(Invariant),
@@ -474,6 +475,7 @@ impl std::fmt::Display for SideExitReason {
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_PACK_BUFFER) => write!(f, "UnknownNewarraySend(PACK_BUFFER)"),
             SideExitReason::UnknownNewarraySend(VM_OPT_NEWARRAY_SEND_INCLUDE_P) => write!(f, "UnknownNewarraySend(INCLUDE_P)"),
             SideExitReason::GuardType(guard_type) => write!(f, "GuardType({guard_type})"),
+            SideExitReason::GuardTypeNot(guard_type) => write!(f, "GuardTypeNot({guard_type})"),
             SideExitReason::GuardBitEquals(value) => write!(f, "GuardBitEquals({})", value.print(&PtrPrintMap::identity())),
             SideExitReason::PatchPoint(invariant) => write!(f, "PatchPoint({invariant})"),
             _ => write!(f, "{self:?}"),
@@ -623,6 +625,7 @@ pub enum Insn {
 
     /// Side-exit if val doesn't have the expected type.
     GuardType { val: InsnId, guard_type: Type, state: InsnId },
+    GuardTypeNot { val: InsnId, guard_type: Type, state: InsnId },
     /// Side-exit if val is not the expected VALUE.
     GuardBitEquals { val: InsnId, expected: VALUE, state: InsnId },
     /// Side-exit if val doesn't have the expected shape.
@@ -859,6 +862,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::FixnumAnd  { left, right, .. } => { write!(f, "FixnumAnd {left}, {right}") },
             Insn::FixnumOr   { left, right, .. } => { write!(f, "FixnumOr {left}, {right}") },
             Insn::GuardType { val, guard_type, .. } => { write!(f, "GuardType {val}, {}", guard_type.print(self.ptr_map)) },
+            Insn::GuardTypeNot { val, guard_type, .. } => { write!(f, "GuardTypeNot {val}, {}", guard_type.print(self.ptr_map)) },
             Insn::GuardBitEquals { val, expected, .. } => { write!(f, "GuardBitEquals {val}, {}", expected.print(self.ptr_map)) },
             &Insn::GuardShape { val, shape, .. } => { write!(f, "GuardShape {val}, {:p}", self.ptr_map.map_shape(shape)) },
             Insn::PatchPoint { invariant, .. } => { write!(f, "PatchPoint {}", invariant.print(self.ptr_map)) },
@@ -1285,6 +1289,7 @@ impl Function {
             &IfTrue { val, ref target } => IfTrue { val: find!(val), target: find_branch_edge!(target) },
             &IfFalse { val, ref target } => IfFalse { val: find!(val), target: find_branch_edge!(target) },
             &GuardType { val, guard_type, state } => GuardType { val: find!(val), guard_type, state },
+            &GuardTypeNot { val, guard_type, state } => GuardTypeNot { val: find!(val), guard_type, state },
             &GuardBitEquals { val, expected, state } => GuardBitEquals { val: find!(val), expected, state },
             &GuardShape { val, shape, state } => GuardShape { val: find!(val), shape, state },
             &FixnumAdd { left, right, state } => FixnumAdd { left: find!(left), right: find!(right), state },
@@ -1430,6 +1435,7 @@ impl Function {
             Insn::ObjectAlloc { .. } => types::HeapObject,
             Insn::CCall { return_type, .. } => *return_type,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
+            Insn::GuardTypeNot { .. } => types::BasicObject,
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
             Insn::GuardShape { val, .. } => self.type_of(*val),
             Insn::FixnumAdd  { .. } => types::Fixnum,
@@ -1546,9 +1552,10 @@ impl Function {
     fn chase_insn(&self, insn: InsnId) -> InsnId {
         let id = self.union_find.borrow().find_const(insn);
         match self.insns[id.0] {
-            Insn::GuardType { val, .. } => self.chase_insn(val),
-            Insn::GuardShape { val, .. } => self.chase_insn(val),
-            Insn::GuardBitEquals { val, .. } => self.chase_insn(val),
+            Insn::GuardType { val, .. }
+            | Insn::GuardTypeNot { val, .. }
+            | Insn::GuardShape { val, .. }
+            | Insn::GuardBitEquals { val, .. } => self.chase_insn(val),
             _ => id,
         }
     }
@@ -1791,12 +1798,26 @@ impl Function {
                         self.insn_types[replacement.0] = self.infer_type(replacement);
                         self.make_equal_to(insn_id, replacement);
                     }
-                    Insn::ObjToString { val, .. } => {
+                    Insn::ObjToString { val, cd, state, .. } => {
                         if self.is_a(val, types::String) {
                             // behaves differently from `SendWithoutBlock` with `mid:to_s` because ObjToString should not have a patch point for String to_s being redefined
-                            self.make_equal_to(insn_id, val);
+                            self.make_equal_to(insn_id, val); continue;
+                        }
+
+                        let frame_state = self.frame_state(state);
+                        let Some(recv_type) = self.profiled_type_of_at(val, frame_state.insn_idx) else {
+                            self.push_insn_id(block, insn_id); continue
+                        };
+
+                        if recv_type.is_string() {
+                            let guard = self.push_insn(block, Insn::GuardType { val: val, guard_type: types::String, state: state });
+                            // Infer type so AnyToString can fold off this
+                            self.insn_types[guard.0] = self.infer_type(guard);
+                            self.make_equal_to(insn_id, guard);
                         } else {
-                            self.push_insn_id(block, insn_id);
+                            self.push_insn(block, Insn::GuardTypeNot { val: val, guard_type: types::String, state: state});
+                            let send_to_s = self.push_insn(block, Insn::SendWithoutBlock { self_val: val, cd: cd, args: vec![], state: state});
+                            self.make_equal_to(insn_id, send_to_s);
                         }
                     }
                     Insn::AnyToString { str, .. } => {
@@ -2193,6 +2214,7 @@ impl Function {
             | &Insn::StringCopy { val, state, .. }
             | &Insn::ObjectAlloc { val, state }
             | &Insn::GuardType { val, state, .. }
+            | &Insn::GuardTypeNot { val, state, .. }
             | &Insn::GuardBitEquals { val, state, .. }
             | &Insn::GuardShape { val, state, .. }
             | &Insn::ToArray { val, state }
@@ -8281,6 +8303,71 @@ mod opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_objtostring_anytostring_recv_profiled() {
+        eval("
+            def test(a)
+              \"#{a}\"
+            end
+            test('foo'); test('foo')
+        ");
+
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v5:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v17:String = GuardType v1, String
+          v11:StringExact = StringConcat v5, v17
+          CheckInterrupts
+          Return v11
+        ");
+    }
+
+    #[test]
+    fn test_optimize_objtostring_anytostring_recv_profiled_string_subclass() {
+        eval("
+            class MyString < String; end
+
+            def test(a)
+              \"#{a}\"
+            end
+            foo = MyString.new('foo')
+            test(MyString.new(foo)); test(MyString.new(foo))
+        ");
+
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:5:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v5:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v17:String = GuardType v1, String
+          v11:StringExact = StringConcat v5, v17
+          CheckInterrupts
+          Return v11
+        ");
+    }
+
+    #[test]
+    fn test_optimize_objtostring_profiled_nonstring_falls_back_to_send() {
+        eval("
+            def test(a)
+              \"#{a}\"
+            end
+            test([1,2,3]); test([1,2,3]) # No fast path for array
+        ");
+
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject):
+          v5:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v17:BasicObject = GuardTypeNot v1, String
+          v18:BasicObject = SendWithoutBlock v1, :to_s
+          v9:String = AnyToString v1, str: v18
+          v11:StringExact = StringConcat v5, v9
+          CheckInterrupts
+          Return v11
+        ");
+    }
+
     #[test]
     fn test_branchnil_nil() {
         eval("
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
@@ -75,6 +75,7 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
         YARVINSN_opt_empty_p => profile_operands(profiler, profile, 1),
         YARVINSN_opt_not   => profile_operands(profiler, profile, 1),
         YARVINSN_getinstancevariable => profile_self(profiler, profile),
+        YARVINSN_objtostring   => profile_operands(profiler, profile, 1),
         YARVINSN_opt_send_without_block => {
             let cd: *const rb_call_data = profiler.insn_opnd(0).as_ptr();
             let argc = unsafe { vm_ci_argc((*cd).ci) };
@@ -235,6 +236,20 @@ impl ProfiledType {
         self.class == unsafe { rb_cInteger } && self.flags.is_immediate()
     }
 
+    pub fn is_string(&self) -> bool {
+        // Fast paths for immediates and exact-class
+        if self.flags.is_immediate() {
+            return false;
+        }
+
+        let string = unsafe { rb_cString };
+        if self.class == string{
+            return true;
+        }
+
+        self.class.is_subclass_of(string) == ClassRelationship::Subclass
+    }
+
     pub fn is_flonum(&self) -> bool {
         self.class == unsafe { rb_cFloat } && self.flags.is_immediate()
     }
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs

Original file line number	Diff line number	Diff line change
`@@ -939,6 +939,7 @@ objtostring`
`939`	`939`	`(VALUE recv)`
`940`	`940`	`(VALUE val)`
`941`	`941`	`// attr bool leaf = false;`
	`942`	`+// attr bool zjit_profile = true;`
`942`	`943`	`{`
`943`	`944`	`val = vm_objtostring(GET_ISEQ(), recv, cd);`
`944`	`945`