YJIT: upgrade type in guard_object_is_string (#7489)

* YJIT: upgrade type in guard_object_is_string

Also make logic more in line with other guard_xxx methods

* Update yjit/src/core.rs

* Revert changes to Type::upgrade()

yjit/src/codegen.rs

@@ -1376,6 +1376,39 @@ fn guard_object_is_array(
     }
 }
 
+fn guard_object_is_string(
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    object: Opnd,
+    object_opnd: YARVOpnd,
+    side_exit: Target,
+) {
+    let object_type = ctx.get_opnd_type(object_opnd);
+    if object_type.is_string() {
+        return;
+    }
+
+    let object_reg = match object {
+        Opnd::Reg(_) => object,
+        _ => asm.load(object),
+    };
+    guard_object_is_heap(ctx, asm, object_reg, object_opnd, side_exit);
+
+    asm.comment("guard object is string");
+
+    // Pull out the type mask
+    let flags_reg = asm.load(Opnd::mem(VALUE_BITS, object_reg, RUBY_OFFSET_RBASIC_FLAGS));
+    let flags_reg = asm.and(flags_reg, Opnd::UImm(RUBY_T_MASK as u64));
+
+    // Compare the result with T_STRING
+    asm.cmp(flags_reg, Opnd::UImm(RUBY_T_STRING as u64));
+    asm.jne(side_exit);
+
+    if object_type.diff(Type::TString) != TypeDiff::Incompatible {
+        ctx.upgrade_opnd_type(object_opnd, Type::TString);
+    }
+}
+
 /// This guards that a special flag is not set on a hash.
 /// By passing a hash with this flag set as the last argument
 /// in a splat call, you can change the way keywords are handled
@@ -1410,22 +1443,6 @@ fn guard_object_is_not_ruby2_keyword_hash(
     asm.write_label(not_ruby2_keyword);
 }
 
-fn guard_object_is_string(
-    asm: &mut Assembler,
-    object_reg: Opnd,
-    side_exit: Target,
-) {
-    asm.comment("guard object is string");
-
-    // Pull out the type mask
-    let flags_reg = asm.load(Opnd::mem(VALUE_BITS, object_reg, RUBY_OFFSET_RBASIC_FLAGS));
-    let flags_reg = asm.and(flags_reg, Opnd::UImm(RUBY_T_MASK as u64));
-
-    // Compare the result with T_STRING
-    asm.cmp(flags_reg, Opnd::UImm(RUBY_T_STRING as u64));
-    asm.jne(side_exit);
-}
-
 // push enough nils onto the stack to fill out an array
 fn gen_expandarray(
     jit: &mut JITState,
@@ -4322,28 +4339,15 @@ fn jit_rb_str_concat(
     // Generate a side exit
     let side_exit = get_side_exit(jit, ocb, ctx);
 
-    // Guard that the argument is of class String at runtime.
-    let arg_type = ctx.get_opnd_type(StackOpnd(0));
+    // Guard that the concat argument is a string
+    guard_object_is_string(ctx, asm, ctx.stack_opnd(0), StackOpnd(0), side_exit);
 
     // Guard buffers from GC since rb_str_buf_append may allocate.
     gen_save_sp(asm, ctx);
 
     let concat_arg = ctx.stack_pop(1);
     let recv = ctx.stack_pop(1);
 
-    // If we're not compile-time certain that this will always be a string, guard at runtime
-    if arg_type != Type::CString && arg_type != Type::TString {
-        let arg_opnd = asm.load(concat_arg);
-        if !arg_type.is_heap() {
-            asm.comment("guard arg not immediate");
-            asm.test(arg_opnd, (RUBY_IMMEDIATE_MASK as u64).into());
-            asm.jnz(side_exit);
-            asm.cmp(arg_opnd, Qfalse.into());
-            asm.je(side_exit);
-        }
-        guard_object_is_string(asm, arg_opnd, side_exit);
-    }
-
     // Test if string encodings differ. If different, use rb_str_append. If the same,
     // use rb_yjit_str_simple_append, which calls rb_str_cat.
     asm.comment("<< on strings");
@@ -7048,6 +7052,7 @@ fn gen_objtostring(
             SEND_MAX_DEPTH,
             side_exit,
         );
+
         // No work needed. The string value is already on the top of the stack.
         KeepCompiling
     } else {

yjit/src/core.rs

@@ -154,6 +154,15 @@ impl Type {
         }
     }
 
+    /// Check if it's a T_STRING object (both TString and CString are T_STRING)
+    pub fn is_string(&self) -> bool {
+        match self {
+            Type::TString => true,
+            Type::CString => true,
+            _ => false,
+        }
+    }
+
     /// Returns an Option with the T_ value type if it is known, otherwise None
     pub fn known_value_type(&self) -> Option<ruby_value_type> {
         match self {
@@ -258,10 +267,10 @@ impl Type {
 
     /// Upgrade this type into a more specific compatible type
     /// The new type must be compatible and at least as specific as the previously known type.
-    fn upgrade(&mut self, src: Self) {
-        // Here we're checking that src is more specific than self
-        assert!(src.diff(*self) != TypeDiff::Incompatible);
-        *self = src;
+    fn upgrade(&mut self, new_type: Self) {
+        // We can only upgrade to a type that is more specific
+        assert!(new_type.diff(*self) != TypeDiff::Incompatible);
+        *self = new_type;
     }
 }