YJIT: Inline simple getlocal+leave iseqs

This mainly targets things like `T.unsafe()` from Sorbet, which is just an
identity function at runtime and only a hint for the static checker.
Only deal with simple caller and callees (no keywords and splat etc.).

Co-authored-by: Takashi Kokubun (k0kubun) <takashikkbn@gmail.com>

3 files changed, 49 insertions, 3 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 24253a10c9..e5a160fa9f 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -4753,3 +4753,19 @@ assert_equal 'foo', %q{
   entry(false)
   entry(true)
 }
+
+assert_equal '[:ok, :ok, :ok]', %q{
+  def identity(x) = x
+  def foo(x, _) = x
+  def bar(_, _, _, _, x) = x
+
+  def tests
+    [
+      identity(:ok),
+      foo(:ok, 2),
+      bar(1, 2, 3, 4, :ok),
+    ]
+  end
+
+  tests
+}
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 7d9085d4b3..0d536907c1 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6754,11 +6754,16 @@ fn gen_send_bmethod(
 /// The kind of a value an ISEQ returns
 enum IseqReturn {
     Value(VALUE),
+    LocalVariable(u32),
     Receiver,
 }
 
-/// Return the ISEQ's return value if it consists of only putnil/putobject and leave.
-fn iseq_get_return_value(iseq: IseqPtr, captured_opnd: Option<Opnd>) -> Option<IseqReturn> {
+extern {
+    fn rb_simple_iseq_p(iseq: IseqPtr) -> bool;
+}
+
+/// Return the ISEQ's return value if it consists of one simple instruction and leave.
+fn iseq_get_return_value(iseq: IseqPtr, captured_opnd: Option<Opnd>, ci_flags: u32) -> Option<IseqReturn> {
     // Expect only two instructions and one possible operand
     let iseq_size = unsafe { get_iseq_encoded_size(iseq) };
     if !(2..=3).contains(&iseq_size) {
@@ -6774,6 +6779,17 @@ fn iseq_get_return_value(iseq: IseqPtr, captured_opnd: Option<Opnd>) -> Option<I
         return None;
     }
     match first_insn {
+        YARVINSN_getlocal_WC_0  => {
+            // Only accept simple positional only cases for both the caller and the callee.
+            // Reject block ISEQs to avoid autosplat and other block parameter complications.
+            if captured_opnd.is_none() && unsafe { rb_simple_iseq_p(iseq) } && ci_flags & VM_CALL_ARGS_SIMPLE != 0 {
+                let ep_offset = unsafe { *rb_iseq_pc_at_idx(iseq, 1) }.as_u32();
+                let local_idx = ep_offset_to_local_idx(iseq, ep_offset);
+                Some(IseqReturn::LocalVariable(local_idx))
+            } else {
+                None
+            }
+        }
         YARVINSN_putnil => Some(IseqReturn::Value(Qnil)),
         YARVINSN_putobject => Some(IseqReturn::Value(unsafe { *rb_iseq_pc_at_idx(iseq, 1) })),
         YARVINSN_putobject_INT2FIX_0_ => Some(IseqReturn::Value(VALUE::fixnum_from_usize(0))),
@@ -7055,11 +7071,24 @@ fn gen_send_iseq(
     }
 
     // Inline simple ISEQs whose return value is known at compile time
-    if let (Some(value), None, false) = (iseq_get_return_value(iseq, captured_opnd), block_arg_type, opt_send_call) {
+    if let (Some(value), None, false) = (iseq_get_return_value(iseq, captured_opnd, flags), block_arg_type, opt_send_call) {
         asm_comment!(asm, "inlined simple ISEQ");
         gen_counter_incr(asm, Counter::num_send_iseq_inline);
 
         match value {
+            IseqReturn::LocalVariable(local_idx) => {
+                // Put the local variable at the return slot
+                let stack_local = asm.stack_opnd(argc - 1 - local_idx as i32);
+                let stack_return = asm.stack_opnd(argc);
+                asm.mov(stack_return, stack_local);
+
+                // Update the mapping for the return value
+                let mapping = asm.ctx.get_opnd_mapping(stack_local.into());
+                asm.ctx.set_opnd_mapping(stack_return.into(), mapping);
+
+                // Pop everything but the return value
+                asm.stack_pop(argc as usize);
+            }
             IseqReturn::Value(value) => {
                 // Pop receiver and arguments
                 asm.stack_pop(argc as usize + if captured_opnd.is_some() { 0 } else { 1 });
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 31d842531f..9547e3fa2c 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -711,6 +711,7 @@ mod manual_defs {
     pub const RUBY_FIXNUM_MAX: isize = RUBY_LONG_MAX / 2;
 
     // From vm_callinfo.h - uses calculation that seems to confuse bindgen
+    pub const VM_CALL_ARGS_SIMPLE: u32 = 1 << VM_CALL_ARGS_SIMPLE_bit;
     pub const VM_CALL_ARGS_SPLAT: u32 = 1 << VM_CALL_ARGS_SPLAT_bit;
     pub const VM_CALL_ARGS_BLOCKARG: u32 = 1 << VM_CALL_ARGS_BLOCKARG_bit;
     pub const VM_CALL_FCALL: u32 = 1 << VM_CALL_FCALL_bit;