YJIT: Bail out if proc would be stored above stack top

Fixes [Bug #21266].

Backport of 9168cad4d63a5d281d443bde4edea6be213b0b25 to 3.3

4 files changed, 28 insertions, 0 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 02767a53bd..0e9461f8e2 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -4585,3 +4585,13 @@ assert_normal_exit %{
     new.foo
   end
 }
+
+# regression test for splat with &proc{} when the target has rest (Bug #21266)
+assert_equal '[]', %q{
+  def foo(args) = bar(*args, &proc { _1 })
+  def bar(_, _, _, _, *rest) = yield rest
+
+  GC.stress = true
+  foo([1,2,3,4])
+  foo([1,2,3,4])
+}
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 206193d86d..0d40d38872 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1526,6 +1526,18 @@ class TestYJIT < Test::Unit::TestCase
     assert_in_out_err(%w[--yjit-stats --yjit-disable])
   end
 
+  def test_proc_block_with_kwrest
+    # When the bug was present this required --yjit-stats to trigger.
+    assert_compiles(<<~RUBY, result: {extra: 5})
+      def foo = bar(w: 1, x: 2, y: 3, z: 4, extra: 5, &proc { _1 })
+      def bar(w:, x:, y:, z:, **kwrest) = yield kwrest
+
+      GC.stress = true
+      foo
+      foo
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1d1e59161e..b881d24a87 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6338,6 +6338,11 @@ fn gen_send_iseq(
                 gen_counter_incr(asm, Counter::send_iseq_clobbering_block_arg);
                 return None;
             }
+            if iseq_has_rest {
+                // The proc would be stored above the current stack top, where GC can't see it
+                gen_counter_incr(asm, Counter::send_iseq_block_arg_gc_unsafe);
+                return None;
+            }
             let proc = asm.stack_pop(1); // Pop first, as argc doesn't account for the block arg
             let callee_specval = asm.ctx.sp_opnd(callee_specval as isize * SIZEOF_VALUE as isize);
             asm.store(callee_specval, proc);
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 3e5b6db176..9325b09e94 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -328,6 +328,7 @@ make_counters! {
     send_attrset_kwargs,
     send_iseq_tailcall,
     send_iseq_arity_error,
+    send_iseq_block_arg_gc_unsafe,
     send_iseq_clobbering_block_arg,
     send_iseq_leaf_builtin_block_arg_block_param,
     send_iseq_only_keywords,