From df7d9812cc504f3361792f3dd4843d1ffa3c5ead Mon Sep 17 00:00:00 2001
From: Stan Lo <stan.lo@shopify.com>
Date: Wed, 13 Aug 2025 01:39:46 +0100
Subject: [PATCH] ZJIT: Prepare non-leaf calls for SetGlobal (#14197)

When trace_var is used, setting a global variable can cause exceptions
to be raised. We need to prepare for that.
---
 test/ruby/test_zjit.rb | 24 ++++++++++++++++++++++++
 zjit/src/codegen.rs    |  8 ++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 4dc0919b6b..98146b80db 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -61,6 +61,30 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_setglobal
+    assert_compiles '1', %q{
+      def test
+        $a = 1
+        $a
+      end
+
+      test
+    }, insns: [:setglobal]
+  end
+
+  def test_setglobal_with_trace_var_exception
+    assert_compiles '"rescued"', %q{
+      def test
+        $a = 1
+      rescue
+        "rescued"
+      end
+
+      trace_var(:$a) { raise }
+      test
+    }, insns: [:setglobal]
+  end
+
   def test_setlocal
     assert_compiles '3', %q{
       def test(n)
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 01ed0f0590..e7bd3285dd 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -367,7 +367,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::PatchPoint { invariant, state } => return gen_patch_point(jit, asm, invariant, &function.frame_state(*state)),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(asm, *cfun, opnds!(args))?,
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
-        Insn::SetGlobal { id, val, state: _ } => return Some(gen_setglobal(asm, *id, opnd!(val))),
+        Insn::SetGlobal { id, val, state } => return gen_setglobal(jit, asm, *id, opnd!(val), &function.frame_state(*state)),
         Insn::GetGlobal { id, state: _ } => gen_getglobal(asm, *id),
         &Insn::GetLocal { ep_offset, level } => gen_getlocal_with_ep(asm, ep_offset, level)?,
         Insn::SetLocal { val, ep_offset, level } => return gen_setlocal_with_ep(asm, opnd!(val), *ep_offset, *level),
@@ -592,8 +592,12 @@ fn gen_getglobal(asm: &mut Assembler, id: ID) -> Opnd {
 }
 
 /// Set global variables
-fn gen_setglobal(asm: &mut Assembler, id: ID, val: Opnd) {
+fn gen_setglobal(jit: &mut JITState, asm: &mut Assembler, id: ID, val: Opnd, state: &FrameState) -> Option<()> {
+    // When trace_var is used, setting a global variable can cause exceptions
+    gen_prepare_non_leaf_call(jit, asm, state)?;
+
     asm_ccall!(asm, rb_gvar_set, id.0.into(), val);
+    Some(())
 }
 
 /// Side-exit into the interpreter