Add setglobal to yjit

Adds yjit support for setting global variables.

Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>
Co-authored-by: John Hawthorn <john@hawthorn.email>

3 files changed, 36 insertions, 1 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 91b63cf076..1a683c1c41 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,5 +1,13 @@
-# Check that global variables work
+# Check that global variable set works
+assert_equal 'string', %q{
+  def foo
+    $foo = "string"
+  end
 
+  foo
+}
+
+# Check that global variables work
 assert_equal 'string', %q{
   $foo = "string"
 
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 8b61a1dbfc..22b78a4e2d 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -55,6 +55,10 @@ class TestYJIT < Test::Unit::TestCase
     assert_compiles('-"foo" == -"bar"', insns: %i[opt_eq], result: false)
   end
 
+  def test_compile_set_and_get_global
+    assert_compiles('$foo = 123; $foo', insns: %i[setglobal], result: 123)
+  end
+
   def test_getlocal_with_level
     assert_compiles(<<~RUBY, insns: %i[getlocal opt_plus], result: [[7]])
       def foo(foo, bar)
diff --git a/yjit_codegen.c b/yjit_codegen.c
index 42beb7dd94..3658d208fa 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -3448,6 +3448,28 @@ gen_getglobal(jitstate_t* jit, ctx_t* ctx)
 }
 
 static codegen_status_t
+gen_setglobal(jitstate_t* jit, ctx_t* ctx)
+{
+    ID gid = jit_get_arg(jit, 0);
+
+    // Save YJIT registers
+    yjit_save_regs(cb);
+
+    mov(cb, C_ARG_REGS[0], imm_opnd(gid));
+
+    x86opnd_t val = ctx_stack_pop(ctx, 1);
+
+    mov(cb, C_ARG_REGS[1], val);
+
+    call_ptr(cb, REG0, (void *)&rb_gvar_set);
+
+    // Load YJIT registers
+    yjit_load_regs(cb);
+
+    return YJIT_KEEP_COMPILING;
+}
+
+static codegen_status_t
 gen_opt_getinlinecache(jitstate_t *jit, ctx_t *ctx)
 {
     VALUE jump_offset = jit_get_arg(jit, 0);
@@ -3690,6 +3712,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(send), gen_send);
     yjit_reg_op(BIN(leave), gen_leave);
     yjit_reg_op(BIN(getglobal), gen_getglobal);
+    yjit_reg_op(BIN(setglobal), gen_setglobal);
 
     yjit_method_codegen_table = st_init_numtable();