YJIT: Support inlining putself (#10137)

1 files changed, 27 insertions, 11 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1e8f046403..eb85239998 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6638,8 +6638,14 @@ fn gen_send_bmethod(
     perf_call! { gen_send_iseq(jit, asm, ocb, iseq, ci, frame_type, Some(capture.ep), cme, block, flags, argc, None) }
 }
 
+/// The kind of a value an ISEQ returns
+enum IseqReturn {
+    Value(VALUE),
+    Receiver,
+}
+
 /// Return the ISEQ's return value if it consists of only putnil/putobject and leave.
-fn iseq_get_return_value(iseq: IseqPtr) -> Option<VALUE> {
+fn iseq_get_return_value(iseq: IseqPtr, captured_opnd: Option<Opnd>) -> 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) {
@@ -6655,10 +6661,12 @@ fn iseq_get_return_value(iseq: IseqPtr) -> Option<VALUE> {
         return None;
     }
     match first_insn {
-        YARVINSN_putnil => Some(Qnil),
-        YARVINSN_putobject => unsafe { Some(*rb_iseq_pc_at_idx(iseq, 1)) },
-        YARVINSN_putobject_INT2FIX_0_ => Some(VALUE::fixnum_from_usize(0)),
-        YARVINSN_putobject_INT2FIX_1_ => Some(VALUE::fixnum_from_usize(1)),
+        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))),
+        YARVINSN_putobject_INT2FIX_1_ => Some(IseqReturn::Value(VALUE::fixnum_from_usize(1))),
+        // We don't support invokeblock for now. Such ISEQs are likely not used by blocks anyway.
+        YARVINSN_putself if captured_opnd.is_none() => Some(IseqReturn::Receiver),
         _ => None,
     }
 }
@@ -6934,16 +6942,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), block_arg_type, opt_send_call) {
+    if let (Some(value), None, false) = (iseq_get_return_value(iseq, captured_opnd), block_arg_type, opt_send_call) {
         asm_comment!(asm, "inlined simple ISEQ");
         gen_counter_incr(asm, Counter::num_send_iseq_inline);
 
-        // Pop receiver and arguments
-        asm.stack_pop(argc as usize + if captured_opnd.is_some() { 0 } else { 1 });
+        match value {
+            IseqReturn::Value(value) => {
+                // Pop receiver and arguments
+                asm.stack_pop(argc as usize + if captured_opnd.is_some() { 0 } else { 1 });
 
-        // Push the return value
-        let stack_ret = asm.stack_push(Type::from(value));
-        asm.mov(stack_ret, value.into());
+                // Push the return value
+                let stack_ret = asm.stack_push(Type::from(value));
+                asm.mov(stack_ret, value.into());
+            },
+            IseqReturn::Receiver => {
+                // Just pop arguments and leave the receiver on stack
+                asm.stack_pop(argc as usize);
+            }
+        }
 
         // Let guard chains share the same successor
         jump_to_next_insn(jit, asm, ocb);