YJIT: Save PC and SP before calling leaf builtins (#7090)

Previously, we did not update `cfp->sp` before calling the C function of
ISEQs marked with `Primitive.attr! "inline"` (leaf builtins). This
caused the GC to miss temporary values on the stack in case the function
allocates and triggers a GC run. Right now, there is only a few leaf
builtins in numeric.rb on Integer methods such as `Integer#~`. Since
these methods only allocate when operating on big numbers, we missed
this issue.

Fix by saving PC and SP before calling the functions -- our usual
protocol for calling C functions that may allocate on the GC heap.

[Bug #19316]

2 files changed, 20 insertions, 0 deletions

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 1a552a1074..6b6ea6619e 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1042,6 +1042,22 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
+  def test_bug_19316
+    n = 2 ** 64
+    # foo's extra param and the splats are relevant
+    assert_compiles(<<~'RUBY', result: [[n, -n], [n, -n]])
+      def foo(_, a, b, c)
+        [a & b, ~c]
+      end
+
+      n = 2 ** 64
+      args = [0, -n, n, n-1]
+
+      GC.stress = true
+      [foo(*args), foo(*args)]
+    RUBY
+  end
+
   private
 
   def code_gc_helpers
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index f99c592e18..a7c2b37eeb 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5181,6 +5181,10 @@ fn gen_send_iseq(
         if builtin_argc + 1 < (C_ARG_OPNDS.len() as i32) {
             asm.comment("inlined leaf builtin");
 
+            // Save the PC and SP because the callee may allocate
+            // e.g. Integer#abs on a bignum
+            jit_prepare_routine_call(jit, ctx, asm);
+
             // Call the builtin func (ec, recv, arg1, arg2, ...)
             let mut args = vec![EC];