YJIT: Fix staying in invalidated code after proc calls

Previously, there is no instruction boundary patch point after
the call to a non-leaf C function we generate for
OPTIMIZED_METHOD_TYPE_CALL. This meant that if code GC is triggered
while inside the C function, we would keep running invalidated code when
we return from the C function. This had the effect of running
stale branch stubs, jumping to bad code, etc.

Use jit_prepare_routine_call() to make sure we exit from the invalidated
region as soon as possible after the C call in case of invalidation.

1 files changed, 2 insertions, 5 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index b1b854ad7f..421e14c553 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5555,11 +5555,8 @@ fn gen_send_general(
 
                         let sp = asm.lea(ctx.sp_opnd(0));
 
-                        // Write interpreter SP into CFP.
-                        // Needed in case the callee yields to the block.
-                        jit_save_pc(jit, asm);
-                        // Store incremented PC into current control frame in case callee raises.
-                        gen_save_sp(jit, asm, ctx);
+                        // Save the PC and SP because the callee can make Ruby calls
+                        jit_prepare_routine_call(jit, ctx, asm);
 
                         let kw_splat = flags & VM_CALL_KW_SPLAT;
                         let stack_argument_pointer = asm.lea(Opnd::mem(64, sp, -(argc) * SIZEOF_VALUE_I32));