Ruby shovel operator (<<) speedup. (#5896)

For string concat, see if compile-time encoding of strings matches.
If so, use simple buffer string concat at runtime. Otherwise, use
encoding-checking string concat.

test/ruby/test_yjit.rb

@@ -316,6 +316,32 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_string_concat_utf8
+    assert_compiles(<<~RUBY, frozen_string_literal: true, result: true)
+      def str_cat_utf8
+        s = String.new
+        10.times { s << "✅" }
+        s
+      end
+
+      str_cat_utf8 == "✅" * 10
+    RUBY
+  end
+
+  def test_string_concat_ascii
+    # Constant-get for classes (e.g. String, Encoding) can cause a side-exit in getinlinecache. For now, ignore exits.
+    assert_compiles(<<~RUBY, exits: :any)
+      str_arg = "b".encode(Encoding::ASCII)
+      def str_cat_ascii(arg)
+        s = String.new(encoding: Encoding::ASCII)
+        10.times { s << arg }
+        s
+      end
+
+      str_cat_ascii(str_arg) == str_arg * 10
+    RUBY
+  end
+
   def test_opt_length_in_method
     assert_compiles(<<~RUBY, insns: %i[opt_length], result: 5)
       def foo(str)
@@ -646,7 +672,7 @@ class TestYJIT < Test::Unit::TestCase
     disasm = stats[:disasm]
 
     # Check that exit counts are as expected
-    # Full stats are only available when RUBY_DEBUG enabled
+    # Full stats are only available when --enable-yjit=dev
     if runtime_stats[:all_stats]
       recorded_exits = runtime_stats.select { |k, v| k.to_s.start_with?("exit_") }
       recorded_exits = recorded_exits.reject { |k, v| v == 0 }
@@ -658,7 +684,7 @@ class TestYJIT < Test::Unit::TestCase
       end
     end
 
-    # Only available when RUBY_DEBUG enabled
+    # Only available when --enable-yjit=dev
     if runtime_stats[:all_stats]
       missed_insns = insns.dup
 
@@ -675,13 +701,18 @@ class TestYJIT < Test::Unit::TestCase
     end
   end
 
+  def script_shell_encode(s)
+    # We can't pass utf-8-encoded characters directly in a shell arg. But we can use Ruby \u constants.
+    s.chars.map { |c| c.ascii_only? ? c : "\\u%x" % c.codepoints[0] }.join
+  end
+
   def eval_with_jit(script, call_threshold: 1, timeout: 1000)
     args = [
       "--disable-gems",
       "--yjit-call-threshold=#{call_threshold}",
       "--yjit-stats"
     ]
-    args << "-e" << script
+    args << "-e" << script_shell_encode(script)
     stats_r, stats_w = IO.pipe
     out, err, status = EnvUtil.invoke_ruby(args,
       '', true, true, timeout: timeout, ios: {3 => stats_w}

yjit.c

@@ -557,6 +557,12 @@ rb_leaf_builtin_function(const rb_iseq_t *iseq)
     return (const struct rb_builtin_function *)iseq->body->iseq_encoded[1];
 }
 
+VALUE
+rb_yjit_str_simple_append(VALUE str1, VALUE str2)
+{
+    return rb_str_cat(str1, RSTRING_PTR(str2), RSTRING_LEN(str2));
+}
+
 struct rb_control_frame_struct *
 rb_get_ec_cfp(rb_execution_context_t *ec)
 {
@@ -692,6 +698,13 @@ rb_RCLASS_ORIGIN(VALUE c)
     return RCLASS_ORIGIN(c);
 }
 
+// Return the string encoding index
+int
+rb_ENCODING_GET(VALUE obj)
+{
+    return RB_ENCODING_GET(obj);
+}
+
 bool
 rb_yjit_multi_ractor_p(void)
 {

yjit/bindgen/src/main.rs

@@ -33,6 +33,7 @@ fn main() {
 
     let bindings = bindgen::builder()
         .clang_args(filtered_clang_args)
+        .header("encindex.h")
         .header("internal.h")
         .header("internal/re.h")
         .header("include/ruby/ruby.h")
@@ -57,6 +58,7 @@ fn main() {
 
         // From include/ruby/internal/intern/string.h
         .allowlist_function("rb_utf8_str_new")
+        .allowlist_function("rb_str_append")
 
         // This struct is public to Ruby C extensions
         // From include/ruby/internal/core/rbasic.h
@@ -69,6 +71,9 @@ fn main() {
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")
 
+        // From ruby/internal/encoding/encoding.h
+        .allowlist_type("ruby_encoding_consts")
+
         // From include/hash.h
         .allowlist_function("rb_hash_new")
 
@@ -228,6 +233,8 @@ fn main() {
         .allowlist_function("rb_yjit_dump_iseq_loc")
         .allowlist_function("rb_yjit_for_each_iseq")
         .allowlist_function("rb_yjit_obj_written")
+        .allowlist_function("rb_yjit_str_simple_append")
+        .allowlist_function("rb_ENCODING_GET")
 
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")

yjit/src/codegen.rs

@@ -3626,6 +3626,91 @@ fn jit_rb_str_to_s(
     false
 }
 
+// Codegen for rb_str_concat()
+// Frequently strings are concatenated using "out_str << next_str".
+// This is common in Erb and similar templating languages.
+fn jit_rb_str_concat(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    cb: &mut CodeBlock,
+    ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    _cme: *const rb_callable_method_entry_t,
+    _block: Option<IseqPtr>,
+    _argc: i32,
+    _known_recv_class: *const VALUE,
+) -> bool {
+    let comptime_arg = jit_peek_at_stack(jit, ctx, 0);
+    let comptime_arg_type = ctx.get_opnd_type(StackOpnd(0));
+
+    // String#<< can take an integer codepoint as an argument, but we don't optimise that.
+    // Also, a non-string argument would have to call .to_str on itself before being treated
+    // as a string, and that would require saving pc/sp, which we don't do here.
+    if comptime_arg_type != Type::String {
+        return false;
+    }
+
+    // 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_opnd = ctx.stack_opnd(0);
+    mov(cb, REG0, arg_opnd);
+    if !jit_guard_known_klass(
+        jit,
+        ctx,
+        cb,
+        ocb,
+        unsafe { rb_cString },
+        StackOpnd(0),
+        comptime_arg,
+        SEND_MAX_DEPTH,
+        side_exit,
+    ) {
+        return false;
+    }
+
+    let concat_arg = ctx.stack_pop(1);
+    let recv = ctx.stack_pop(1);
+
+    // 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.
+    add_comment(cb, "<< on strings");
+
+    // Both rb_str_append and rb_yjit_str_simple_append take identical args
+    mov(cb, C_ARG_REGS[0], recv);
+    mov(cb, C_ARG_REGS[1], concat_arg);
+
+    // Take receiver's object flags XOR arg's flags. If any
+    // string-encoding flags are different between the two,
+    // the encodings don't match.
+    mov(cb, REG0, recv);
+    mov(cb, REG1, concat_arg);
+    mov(cb, REG0, mem_opnd(64, REG0, RUBY_OFFSET_RBASIC_FLAGS));
+    xor(cb, REG0, mem_opnd(64, REG1, RUBY_OFFSET_RBASIC_FLAGS));
+    test(cb, REG0, uimm_opnd(RUBY_ENCODING_MASK as u64));
+
+    let enc_mismatch = cb.new_label("enc_mismatch".to_string());
+    jne_label(cb, enc_mismatch);
+
+    // If encodings match, call the simple append function and jump to return
+    call_ptr(cb, REG0, rb_yjit_str_simple_append as *const u8);
+    let ret_label: usize = cb.new_label("stack_return".to_string());
+    jmp_label(cb, ret_label);
+
+    // If encodings are different, use a slower encoding-aware concatenate
+    cb.write_label(enc_mismatch);
+    call_ptr(cb, REG0, rb_str_append as *const u8);
+    // Drop through to return
+
+    cb.write_label(ret_label);
+    let stack_ret = ctx.stack_push(Type::String);
+    mov(cb, stack_ret, RAX);
+
+    cb.link_labels();
+    true
+}
+
 fn jit_thread_s_current(
     _jit: &mut JITState,
     ctx: &mut Context,
@@ -3887,7 +3972,6 @@ fn gen_send_cfunc(
         // Copy the arguments from the stack to the C argument registers
         // self is the 0th argument and is at index argc from the stack top
         for i in 0..=passed_argc as usize {
-            // "as usize?" Yeah, you can't index an array by an i32.
             let stack_opnd = mem_opnd(64, RAX, -(argc + 1 - (i as i32)) * SIZEOF_VALUE_I32);
             let c_arg_reg = C_ARG_REGS[i];
             mov(cb, c_arg_reg, stack_opnd);
@@ -5839,6 +5923,7 @@ impl CodegenGlobals {
             self.yjit_reg_method(rb_cString, "to_s", jit_rb_str_to_s);
             self.yjit_reg_method(rb_cString, "to_str", jit_rb_str_to_s);
             self.yjit_reg_method(rb_cString, "bytesize", jit_rb_str_bytesize);
+            self.yjit_reg_method(rb_cString, "<<", jit_rb_str_concat);
 
             // Thread.current
             self.yjit_reg_method(

yjit/src/cruby_bindings.inc.rs

@@ -172,6 +172,9 @@ extern "C" {
         len: ::std::os::raw::c_long,
     ) -> VALUE;
 }
+extern "C" {
+    pub fn rb_str_append(dst: VALUE, src: VALUE) -> VALUE;
+}
 extern "C" {
     pub fn rb_str_intern(str_: VALUE) -> VALUE;
 }
@@ -181,6 +184,11 @@ extern "C" {
 extern "C" {
     pub fn rb_attr_get(obj: VALUE, name: ID) -> VALUE;
 }
+pub const RUBY_ENCODING_INLINE_MAX: ruby_encoding_consts = 127;
+pub const RUBY_ENCODING_SHIFT: ruby_encoding_consts = 22;
+pub const RUBY_ENCODING_MASK: ruby_encoding_consts = 532676608;
+pub const RUBY_ENCODING_MAXNAMELEN: ruby_encoding_consts = 42;
+pub type ruby_encoding_consts = u32;
 extern "C" {
     pub fn rb_obj_info_dump(obj: VALUE);
 }
@@ -731,6 +739,9 @@ extern "C" {
 extern "C" {
     pub fn rb_leaf_builtin_function(iseq: *const rb_iseq_t) -> *const rb_builtin_function;
 }
+extern "C" {
+    pub fn rb_yjit_str_simple_append(str1: VALUE, str2: VALUE) -> VALUE;
+}
 extern "C" {
     pub fn rb_set_cfp_pc(cfp: *mut rb_control_frame_struct, pc: *const VALUE);
 }
@@ -743,6 +754,9 @@ extern "C" {
 extern "C" {
     pub fn rb_yjit_dump_iseq_loc(iseq: *const rb_iseq_t, insn_idx: u32);
 }
+extern "C" {
+    pub fn rb_ENCODING_GET(obj: VALUE) -> ::std::os::raw::c_int;
+}
 extern "C" {
     pub fn rb_yjit_multi_ractor_p() -> bool;
 }