RJIT: Implement attr_writer

5 files changed, 75 insertions, 9 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 2dfe52713b..64a1d1a148 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -4028,8 +4028,7 @@ module RubyVM::RJIT
       in C::VM_METHOD_TYPE_CFUNC
         jit_call_cfunc(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
       in C::VM_METHOD_TYPE_ATTRSET
-        asm.incr_counter(:send_attrset)
-        return CantCompile
+        jit_call_attrset(jit, ctx, asm, cme, flags, argc, comptime_recv, recv_opnd, send_shift:)
       in C::VM_METHOD_TYPE_IVAR
         jit_call_ivar(jit, ctx, asm, cme, flags, argc, comptime_recv, recv_opnd, send_shift:)
       in C::VM_METHOD_TYPE_MISSING
@@ -4163,7 +4162,7 @@ module RubyVM::RJIT
       end
 
       # EXEC_EVENT_HOOK: RUBY_EVENT_C_CALL and RUBY_EVENT_C_RETURN
-      if C.rb_rjit_global_events & (C::RUBY_EVENT_C_CALL | C::RUBY_EVENT_C_RETURN) != 0
+      if c_method_tracing_currently_enabled?
         asm.incr_counter(:send_c_tracing)
         return CantCompile
       end
@@ -4261,6 +4260,58 @@ module RubyVM::RJIT
       EndBlock
     end
 
+    # vm_call_attrset
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
+    def jit_call_attrset(jit, ctx, asm, cme, flags, argc, comptime_recv, recv_opnd, send_shift:)
+      if flags & C::VM_CALL_ARGS_SPLAT != 0
+        asm.incr_counter(:send_attrset_splat)
+        return CantCompile
+      end
+      if flags & C::VM_CALL_KWARG != 0
+        asm.incr_counter(:send_attrset_kwarg)
+        return CantCompile
+      elsif argc != 1 || !C.RB_TYPE_P(comptime_recv, C::RUBY_T_OBJECT)
+        asm.incr_counter(:send_attrset_method)
+        return CantCompile
+      elsif c_method_tracing_currently_enabled?
+        # Can't generate code for firing c_call and c_return events
+        # See :attr-tracing:
+        asm.incr_counter(:send_c_tracingg)
+        return CantCompile
+      elsif flags & C::VM_CALL_ARGS_BLOCKARG != 0
+        asm.incr_counter(:send_attrset_blockarg)
+        return CantCompile
+      end
+
+      ivar_name = cme.def.body.attr.id
+
+      # This is a .send call and we need to adjust the stack
+      if flags & C::VM_CALL_OPT_SEND != 0
+        jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
+      end
+
+      # Save the PC and SP because the callee may allocate
+      # Note that this modifies REG_SP, which is why we do it first
+      jit_prepare_routine_call(jit, ctx, asm)
+
+      # Get the operands from the stack
+      val_opnd = ctx.stack_pop(1)
+      recv_opnd = ctx.stack_pop(1)
+
+      # Call rb_vm_set_ivar_id with the receiver, the ivar name, and the value
+      asm.mov(C_ARGS[0], recv_opnd)
+      asm.mov(C_ARGS[1], ivar_name)
+      asm.mov(C_ARGS[2], val_opnd)
+      asm.call(C.rb_vm_set_ivar_id)
+
+      out_opnd = ctx.stack_push
+      asm.mov(out_opnd, C_RET)
+
+      KeepCompiling
+    end
+
     # vm_call_ivar (+ part of vm_call_method_each_type)
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
@@ -5111,5 +5162,9 @@ module RubyVM::RJIT
         @ocb.write(asm)
       end
     end
+
+    def c_method_tracing_currently_enabled?
+      C.rb_rjit_global_events & (C::RUBY_EVENT_C_CALL | C::RUBY_EVENT_C_RETURN) != 0
+    end
   end
 end
diff --git a/rjit_c.c b/rjit_c.c
index f4b34691ec..9431afd896 100644
--- a/rjit_c.c
+++ b/rjit_c.c
@@ -506,6 +506,7 @@ extern void rb_vm_setclassvariable(const rb_iseq_t *iseq, const rb_control_frame
 extern VALUE rb_str_bytesize(VALUE str);
 extern const rb_callable_method_entry_t *rb_callable_method_entry_or_negative(VALUE klass, ID mid);
 extern VALUE rb_vm_yield_with_cfunc(rb_execution_context_t *ec, const struct rb_captured_block *captured, int argc, const VALUE *argv);
+extern VALUE rb_vm_set_ivar_id(VALUE obj, ID id, VALUE val);
 
 #include "rjit_c.rbinc"
 
diff --git a/rjit_c.h b/rjit_c.h
index d6e7cc01d6..3ae4070266 100644
--- a/rjit_c.h
+++ b/rjit_c.h
@@ -35,8 +35,6 @@ RJIT_RUNTIME_COUNTERS(
     send_protected_check_failed,
     send_tailcall,
     send_notimplemented,
-    send_cfunc,
-    send_attrset,
     send_missing,
     send_bmethod,
     send_alias,
@@ -74,7 +72,11 @@ RJIT_RUNTIME_COUNTERS(
     send_cfunc_too_many_args,
     send_cfunc_ruby_array_varg,
 
-    send_ivar,
+    send_attrset_splat,
+    send_attrset_kwarg,
+    send_attrset_method,
+    send_attrset_blockarg,
+
     send_ivar_splat,
     send_ivar_opt_send,
     send_ivar_blockarg,
diff --git a/rjit_c.rb b/rjit_c.rb
index aa2ca3fa2b..fbd18c5d5b 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -408,6 +408,7 @@ module RubyVM::RJIT # :nodoc: all
   C::RUBY_T_ICLASS = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_ICLASS) }
   C::RUBY_T_MASK = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_MASK) }
   C::RUBY_T_MODULE = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_MODULE) }
+  C::RUBY_T_OBJECT = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_OBJECT) }
   C::RUBY_T_STRING = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_STRING) }
   C::RUBY_T_SYMBOL = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_SYMBOL) }
   C::SHAPE_CAPACITY_CHANGE = Primitive.cexpr! %q{ SIZET2NUM(SHAPE_CAPACITY_CHANGE) }
@@ -693,6 +694,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_opt_newarray_min) }
   end
 
+  def C.rb_vm_set_ivar_id
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_set_ivar_id) }
+  end
+
   def C.rb_vm_setclassvariable
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_setclassvariable) }
   end
@@ -1289,8 +1294,6 @@ module RubyVM::RJIT # :nodoc: all
       send_protected_check_failed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_protected_check_failed)")],
       send_tailcall: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_tailcall)")],
       send_notimplemented: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_notimplemented)")],
-      send_cfunc: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc)")],
-      send_attrset: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_attrset)")],
       send_missing: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_missing)")],
       send_bmethod: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_bmethod)")],
       send_alias: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_alias)")],
@@ -1324,7 +1327,10 @@ module RubyVM::RJIT # :nodoc: all
       send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_variadic)")],
       send_cfunc_too_many_args: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_too_many_args)")],
       send_cfunc_ruby_array_varg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_ruby_array_varg)")],
-      send_ivar: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_ivar)")],
+      send_attrset_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_attrset_splat)")],
+      send_attrset_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_attrset_kwarg)")],
+      send_attrset_method: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_attrset_method)")],
+      send_attrset_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_attrset_blockarg)")],
       send_ivar_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_ivar_splat)")],
       send_ivar_opt_send: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_ivar_opt_send)")],
       send_ivar_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_ivar_blockarg)")],
diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
index e9994feaa2..410ea016dd 100755
--- a/tool/rjit/bindgen.rb
+++ b/tool/rjit/bindgen.rb
@@ -433,6 +433,7 @@ generator = BindingGenerator.new(
       RUBY_T_MODULE
       RUBY_T_STRING
       RUBY_T_SYMBOL
+      RUBY_T_OBJECT
       SHAPE_CAPACITY_CHANGE
       SHAPE_FLAG_SHIFT
       SHAPE_FROZEN
@@ -558,6 +559,7 @@ generator = BindingGenerator.new(
     rb_str_buf_append
     rb_str_dup
     rb_vm_yield_with_cfunc
+    rb_vm_set_ivar_id
   ],
   types: %w[
     CALL_DATA