Special-case jit_guard_known_class for strings. This can remove (#5920)

runtime guard-checks for String#to_s, making some blocks too
short to invalidate later. Add NOPs in those cases to reserve space.

2 files changed, 44 insertions, 0 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1d62d74de0..e64f269cfe 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -30,6 +30,13 @@ pub const REG0_8: X86Opnd = AL;
 pub const REG1: X86Opnd = RCX;
 // pub const REG1_32: X86Opnd = ECX;
 
+// A block that can be invalidated needs space to write a jump.
+// We'll reserve a minimum size for any block that could
+// be invalidated. In this case the JMP takes 5 bytes, but
+// gen_send_general will always MOV the receiving object
+// into place, so 2 bytes are always written automatically.
+pub const JUMP_SIZE_IN_BYTES:usize = 3;
+
 /// Status returned by code generation functions
 #[derive(PartialEq, Debug)]
 enum CodegenStatus {
@@ -3411,6 +3418,32 @@ fn jit_guard_known_klass(
             jit_chain_guard(JCC_JNE, jit, ctx, cb, ocb, max_chain_depth, side_exit);
             ctx.upgrade_opnd_type(insn_opnd, Type::Flonum);
         }
+    } else if unsafe { known_klass == rb_cString } && sample_instance.string_p() {
+        assert!(!val_type.is_imm());
+        if val_type != Type::String {
+            assert!(val_type.is_unknown());
+
+            // Need the check for immediate, because trying to look up the klass field of an immediate will segfault
+            if !val_type.is_heap() {
+                add_comment(cb, "guard not immediate (for string)");
+                assert!(Qfalse.as_i32() < Qnil.as_i32());
+                test(cb, REG0, imm_opnd(RUBY_IMMEDIATE_MASK as i64));
+                jit_chain_guard(JCC_JNZ, jit, ctx, cb, ocb, max_chain_depth, side_exit);
+                cmp(cb, REG0, imm_opnd(Qnil.into()));
+                jit_chain_guard(JCC_JBE, jit, ctx, cb, ocb, max_chain_depth, side_exit);
+            }
+
+            add_comment(cb, "guard object is string");
+            let klass_opnd = mem_opnd(64, REG0, RUBY_OFFSET_RBASIC_KLASS);
+            mov(cb, REG1, uimm_opnd(unsafe { rb_cString }.into()));
+            cmp(cb, klass_opnd, REG1);
+            jit_chain_guard(JCC_JNE, jit, ctx, cb, ocb, max_chain_depth, side_exit);
+
+            // Upgrading here causes an error with invalidation writing past end of block
+            ctx.upgrade_opnd_type(insn_opnd, Type::String);
+        } else {
+            add_comment(cb, "skip guard - known to be a string");
+        }
     } else if unsafe {
         FL_TEST(known_klass, VALUE(RUBY_FL_SINGLETON)) != VALUE(0)
             && sample_instance == rb_attr_get(known_klass, id__attached__ as ID)
@@ -3837,7 +3870,13 @@ fn gen_send_cfunc(
     if kw_arg.is_null() {
         let codegen_p = lookup_cfunc_codegen(unsafe { (*cme).def });
         if let Some(known_cfunc_codegen) = codegen_p {
+            let start_pos = cb.get_write_ptr().raw_ptr() as usize;
             if known_cfunc_codegen(jit, ctx, cb, ocb, ci, cme, block, argc, recv_known_klass) {
+                let written_bytes = cb.get_write_ptr().raw_ptr() as usize - start_pos;
+                if written_bytes < JUMP_SIZE_IN_BYTES {
+                    add_comment(cb, "Writing NOPs to leave room for later invalidation code");
+                    nop(cb, (JUMP_SIZE_IN_BYTES - written_bytes) as u32);
+                }
                 // cfunc codegen generated code. Terminate the block so
                 // there isn't multiple calls in the same block.
                 jump_to_next_insn(jit, ctx, cb, ocb);
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 0afbd6451b..20091bf77a 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -461,6 +461,11 @@ impl VALUE {
         self == Qnil
     }
 
+    /// Returns true or false depending whether the value is a string
+    pub fn string_p(self) -> bool {
+        unsafe { CLASS_OF(self) == rb_cString }
+    }
+
     /// Read the flags bits from the RBasic object, then return a Ruby type enum (e.g. RUBY_T_ARRAY)
     pub fn builtin_type(self) -> ruby_value_type {
         assert!(!self.special_const_p());