From 296782ab60e63fe825461a199369c5dd8505c2b7 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Tue, 4 Jul 2023 10:45:29 -0400
Subject: [PATCH] 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        | 10 ++++++++++
 lib/ruby_vm/rjit/insn_compiler.rb | 10 +++++++++-
 yjit/src/codegen.rs               | 13 +++++++++++--
 yjit/src/stats.rs                 |  1 +
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index a6525e7af9c239..8be24cfdb4b963 100644
--- a/bootstraptest/test_yjit.rb
+++ b/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]
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index d9a2ab52ae2341..1f0376dd04d871 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -4553,7 +4553,8 @@ def jit_call_iseq(jit, ctx, asm, cme, calling, iseq, frame_type: nil, prev_ep: n
       # 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 @@ def jit_call_iseq(jit, ctx, asm, cme, calling, iseq, frame_type: nil, prev_ep: n
           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)
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 18d53fe65ebbae..fcf12be8ef8070 100644
--- a/yjit/src/codegen.rs
+++ b/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 {
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 91ecc8209b3328..7c5b3344214a6a 100644
--- a/yjit/src/stats.rs
+++ b/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,