YJIT: Add codegen for Array#<< (#7645)

3 files changed, 28 insertions, 0 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 627a88fd31..c990e3e629 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -136,6 +136,7 @@ fn main() {
         .allowlist_function("rb_ary_resurrect")
         .allowlist_function("rb_ary_clear")
         .allowlist_function("rb_ary_dup")
+        .allowlist_function("rb_ary_push")
         .allowlist_function("rb_ary_unshift_m")
         .allowlist_function("rb_yjit_rb_ary_subseq_length")
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 750ac9e5f5..f3c671d23b 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4542,6 +4542,31 @@ fn jit_rb_ary_empty_p(
     return true;
 }
 
+fn jit_rb_ary_push(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    _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 {
+    asm.comment("Array#<<");
+
+    // rb_ary_push allocates memory for buffer extension
+    jit_prepare_routine_call(jit, ctx, asm);
+
+    let item_opnd = ctx.stack_pop(1);
+    let ary_opnd = ctx.stack_pop(1);
+    let ret = asm.ccall(rb_ary_push as *const u8, vec![ary_opnd, item_opnd]);
+
+    let ret_opnd = ctx.stack_push(Type::TArray);
+    asm.mov(ret_opnd, ret);
+    true
+}
+
 fn jit_obj_respond_to(
     jit: &mut JITState,
     ctx: &mut Context,
@@ -8071,6 +8096,7 @@ impl CodegenGlobals {
 
             // rb_ary_empty_p() method in array.c
             self.yjit_reg_method(rb_cArray, "empty?", jit_rb_ary_empty_p);
+            self.yjit_reg_method(rb_cArray, "<<", jit_rb_ary_push);
 
             self.yjit_reg_method(rb_mKernel, "respond_to?", jit_obj_respond_to);
             self.yjit_reg_method(rb_mKernel, "block_given?", jit_rb_f_block_given_p);
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 3e05388308..34f0f73e07 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1094,6 +1094,7 @@ extern "C" {
     pub fn rb_ary_store(ary: VALUE, key: ::std::os::raw::c_long, val: VALUE);
     pub fn rb_ary_dup(ary: VALUE) -> VALUE;
     pub fn rb_ary_resurrect(ary: VALUE) -> VALUE;
+    pub fn rb_ary_push(ary: VALUE, elem: VALUE) -> VALUE;
     pub fn rb_ary_clear(ary: VALUE) -> VALUE;
     pub fn rb_hash_new() -> VALUE;
     pub fn rb_hash_aref(hash: VALUE, key: VALUE) -> VALUE;