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.

3 files changed, 107 insertions, 1 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 0264751a6f..33f81366c2 100644
--- a/yjit/bindgen/src/main.rs
+++ b/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")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 8ca1644fe6..d4aee9528e 100644
--- a/yjit/src/codegen.rs
+++ b/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(
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 2be42f5c63..ac54ba4446 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -173,6 +173,9 @@ extern "C" {
     ) -> VALUE;
 }
 extern "C" {
+    pub fn rb_str_append(dst: VALUE, src: VALUE) -> VALUE;
+}
+extern "C" {
     pub fn rb_str_intern(str_: VALUE) -> VALUE;
 }
 extern "C" {
@@ -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);
 }
@@ -732,6 +740,9 @@ 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);
 }
 extern "C" {
@@ -744,6 +755,9 @@ 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;
 }
 extern "C" {