YJIT: Fix autosplat miscomp for blocks with optionals (#8006)

* YJIT: Fix autosplat miscomp for blocks with optionals

When passing an array as the sole argument to `yield`, and the yieldee
takes more than 1 optional parameter, the array is expanded similar
to `*array` splat calls. This is called "autosplat" in
`setup_parameters_complex()`.

Previously, YJIT did not detect this autosplat condition. It passed the
array without expanding it, deviating from interpreter behavior.
Detect this conditon and refuse to compile it.

Fixes: Shopify/yjit#313

* RJIT: Fix autosplat miscomp for blocks with optionals

This is mirrors the same issue as YJIT. See previous commit.

bootstraptest/test_yjit.rb

@@ -1,3 +1,13 @@
+# Regression test for yielding with autosplat to block with
+# optional parameters. https://github.com/Shopify/yjit/issues/313
+assert_equal '[:a, :b, :a, :b]', %q{
+  def yielder(arg) = yield(arg) + yield(arg)
+
+  yielder([:a, :b]) do |c = :c, d = :d|
+    [c, d]
+  end
+}
+
 # Regression test for GC mishap while doing shape transition
 assert_equal '[:ok]', %q{
   # [Bug #19601]

lib/ruby_vm/rjit/insn_compiler.rb

@@ -4553,7 +4553,8 @@ module RubyVM::RJIT
       # Check if we need the arg0 splat handling of vm_callee_setup_block_arg
       arg_setup_block = (calling.block_handler == :captured) # arg_setup_type: arg_setup_block (invokeblock)
       block_arg0_splat = arg_setup_block && argc == 1 &&
-        iseq.body.param.flags.has_lead && !iseq.body.param.flags.ambiguous_param0
+        (iseq.body.param.flags.has_lead || opt_num > 1) &&
+        !iseq.body.param.flags.ambiguous_param0
       if block_arg0_splat
         # If block_arg0_splat, we still need side exits after splat, but
         # doing push_splat_args here disallows it. So bail out.
@@ -4567,6 +4568,13 @@ module RubyVM::RJIT
           asm.incr_counter(:invokeblock_iseq_arg0_has_kw)
           return CantCompile
         end
+        # The block_arg0_splat implementation cannot deal with optional parameters.
+        # This is a setup_parameters_complex() situation and interacts with the
+        # starting position of the callee.
+        if opt_num > 1
+          asm.incr_counter(:invokeblock_iseq_arg0_optional)
+          return CantCompile
+        end
       end
       if flags & C::VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest
         array = jit.peek_at_stack(block_arg ? 1 : 0)

yjit/src/codegen.rs

@@ -5590,10 +5590,12 @@ fn gen_send_iseq(
         }
     }
 
-    // Check if we need the arg0 splat handling of vm_callee_setup_block_arg
+    // Check if we need the arg0 splat handling of vm_callee_setup_block_arg()
+    // Also known as "autosplat" inside setup_parameters_complex()
     let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock)
     let block_arg0_splat = arg_setup_block && argc == 1 && unsafe {
-         get_iseq_flags_has_lead(iseq) && !get_iseq_flags_ambiguous_param0(iseq)
+         (get_iseq_flags_has_lead(iseq) || opt_num > 1)
+             && !get_iseq_flags_ambiguous_param0(iseq)
     };
     if block_arg0_splat {
         // If block_arg0_splat, we still need side exits after splat, but
@@ -5608,6 +5610,13 @@ fn gen_send_iseq(
             gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_has_kw);
             return None;
         }
+        // The block_arg0_splat implementation cannot deal with optional parameters.
+        // This is a setup_parameters_complex() situation and interacts with the
+        // starting position of the callee.
+        if opt_num > 1 {
+            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_optional);
+            return None;
+        }
     }
 
     let splat_array_length = if flags & VM_CALL_ARGS_SPLAT != 0 {

yjit/src/stats.rs

@@ -296,6 +296,7 @@ make_counters! {
     invokesuper_block,
 
     invokeblock_none,
+    invokeblock_iseq_arg0_optional,
     invokeblock_iseq_arg0_has_kw,
     invokeblock_iseq_arg0_args_splat,
     invokeblock_iseq_arg0_not_array,