Improve set instance variable

This commit improves the set ivar implementation.

bootstraptest/test_yjit.rb

@@ -239,6 +239,91 @@ assert_normal_exit %q{
   end
 }
 
+# test setinstancevariable on extended objects
+assert_equal '1', %q{
+  class Extended
+    attr_reader :one
+
+    def write_many
+      @a = 1
+      @b = 2
+      @c = 3
+      @d = 4
+      @one = 1
+    end
+  end
+
+  foo = Extended.new
+  foo.write_many
+  foo.write_many
+  foo.write_many
+}
+
+# test setinstancevariable on embedded objects
+assert_equal '1', %q{
+  class Embedded
+    attr_reader :one
+
+    def write_one
+      @one = 1
+    end
+  end
+
+  foo = Embedded.new
+  foo.write_one
+  foo.write_one
+  foo.write_one
+}
+
+# test setinstancevariable after extension
+assert_equal '[10, 11, 12, 13, 1]', %q{
+  class WillExtend
+    attr_reader :one
+
+    def make_extended
+      @foo1 = 10
+      @foo2 = 11
+      @foo3 = 12
+      @foo4 = 13
+    end
+
+    def write_one
+      @one = 1
+    end
+
+    def read_all
+      [@foo1, @foo2, @foo3, @foo4, @one]
+    end
+  end
+
+  foo = WillExtend.new
+  foo.write_one
+  foo.write_one
+  foo.make_extended
+  foo.write_one
+  foo.read_all
+}
+
+# test setinstancevariable on frozen object
+assert_equal 'object was not modified', %q{
+  class WillFreeze
+    def write
+      @ivar = 1
+    end
+  end
+
+  wf = WillFreeze.new
+  wf.write
+  wf.write
+  wf.freeze
+
+  begin
+    wf.write
+  rescue FrozenError
+    "object was not modified"
+  end
+}
+
 # Test getinstancevariable and inline caches
 assert_equal '6', %q{
   class Foo

yjit_codegen.c

@@ -745,6 +745,112 @@ enum {
     OSWB_MAX_DEPTH = 5,           // up to 5 different classes
 };
 
+// Codegen for setting an instance variable.
+// Preconditions:
+//   - receiver is in REG0
+//   - receiver has the same class as CLASS_OF(comptime_receiver)
+//   - no stack push or pops to ctx since the entry to the codegen of the instruction being compiled
+static codegen_status_t
+gen_set_ivar(jitstate_t *jit, ctx_t *ctx, const int max_chain_depth, VALUE comptime_receiver, ID ivar_name, insn_opnd_t reg0_opnd, uint8_t *side_exit)
+{
+    VALUE comptime_val_klass = CLASS_OF(comptime_receiver);
+    const ctx_t starting_context = *ctx; // make a copy for use with jit_chain_guard
+
+    // If the class uses the default allocator, instances should all be T_OBJECT
+    // NOTE: This assumes nobody changes the allocator of the class after allocation.
+    //       Eventually, we can encode whether an object is T_OBJECT or not
+    //       inside object shapes.
+    if (rb_get_alloc_func(comptime_val_klass) != rb_class_allocate_instance) {
+        GEN_COUNTER_INC(cb, setivar_not_object);
+        return YJIT_CANT_COMPILE;
+    }
+    RUBY_ASSERT(BUILTIN_TYPE(comptime_receiver) == T_OBJECT); // because we checked the allocator
+
+    // ID for the name of the ivar
+    ID id = ivar_name;
+    struct rb_iv_index_tbl_entry *ent;
+    struct st_table *iv_index_tbl = ROBJECT_IV_INDEX_TBL(comptime_receiver);
+
+    // Bail if this is a heap object, because this needs a write barrier
+    ADD_COMMENT(cb, "guard value is immediate");
+    test(cb, REG1, imm_opnd(RUBY_IMMEDIATE_MASK));
+    jz_ptr(cb, COUNTED_EXIT(side_exit, setivar_val_heapobject));
+
+    // Lookup index for the ivar the instruction loads
+    if (iv_index_tbl && rb_iv_index_tbl_lookup(iv_index_tbl, id, &ent)) {
+        uint32_t ivar_index = ent->index;
+
+        x86opnd_t val_to_write = ctx_stack_pop(ctx, 1);
+        mov(cb, REG1, val_to_write);
+
+        x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
+
+        // Bail if this object is frozen
+        ADD_COMMENT(cb, "guard self is not frozen");
+        test(cb, flags_opnd, imm_opnd(RUBY_FL_FREEZE));
+        jz_ptr(cb, COUNTED_EXIT(side_exit, setivar_frozen));
+
+        // Pop receiver if it's on the temp stack
+        if (!reg0_opnd.is_self) {
+            (void)ctx_stack_pop(ctx, 1);
+        }
+
+        // Compile time self is embedded and the ivar index lands within the object
+        if (RB_FL_TEST_RAW(comptime_receiver, ROBJECT_EMBED) && ivar_index < ROBJECT_EMBED_LEN_MAX) {
+            // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
+
+            // Guard that self is embedded
+            // TODO: BT and JC is shorter
+            ADD_COMMENT(cb, "guard embedded setivar");
+            test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
+            jit_chain_guard(JCC_JZ, jit, &starting_context, max_chain_depth, side_exit);
+
+            // Load the variable
+            x86opnd_t ivar_opnd = mem_opnd(64, REG0, offsetof(struct RObject, as.ary) + ivar_index * SIZEOF_VALUE);
+
+            mov(cb, ivar_opnd, REG1);
+
+            // Push the ivar on the stack
+            // For attr_writer we'll need to push the value on the stack
+            //x86opnd_t out_opnd = ctx_stack_push(ctx, TYPE_UNKNOWN);
+        }
+        else {
+            // Compile time value is *not* embeded.
+
+            // Guard that value is *not* embedded
+            // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
+            ADD_COMMENT(cb, "guard extended setivar");
+            x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
+            test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
+            jit_chain_guard(JCC_JNZ, jit, &starting_context, max_chain_depth, side_exit);
+
+            // check that the extended table is big enough
+            if (ivar_index >= ROBJECT_EMBED_LEN_MAX + 1) {
+                // Check that the slot is inside the extended table (num_slots > index)
+                x86opnd_t num_slots = mem_opnd(32, REG0, offsetof(struct RObject, as.heap.numiv));
+                cmp(cb, num_slots, imm_opnd(ivar_index));
+                jle_ptr(cb, COUNTED_EXIT(side_exit, setivar_idx_out_of_range));
+            }
+
+            // Get a pointer to the extended table
+            x86opnd_t tbl_opnd = mem_opnd(64, REG0, offsetof(struct RObject, as.heap.ivptr));
+            mov(cb, REG0, tbl_opnd);
+
+            // Write the ivar to the extended table
+            x86opnd_t ivar_opnd = mem_opnd(64, REG0, sizeof(VALUE) * ivar_index);
+            mov(cb, REG1, val_to_write);
+            mov(cb, ivar_opnd, REG1);
+        }
+
+        // Jump to next instruction. This allows guard chains to share the same successor.
+        jit_jump_to_next_insn(jit, ctx);
+        return YJIT_END_BLOCK;
+    }
+
+    GEN_COUNTER_INC(cb, setivar_name_not_mapped);
+    return YJIT_CANT_COMPILE;
+}
+
 // Codegen for getting an instance variable.
 // Preconditions:
 //   - receiver is in REG0
@@ -867,7 +973,7 @@ gen_getinstancevariable(jitstate_t *jit, ctx_t *ctx)
 
     // Guard that the receiver has the same class as the one from compile time.
     mov(cb, REG0, member_opnd(REG_CFP, rb_control_frame_t, self));
-    guard_self_is_heap(cb, REG0, side_exit, ctx);
+    guard_self_is_heap(cb, REG0, COUNTED_EXIT(side_exit, getivar_se_self_not_heap), ctx);
 
     jit_guard_known_klass(jit, ctx, comptime_val_klass, OPND_SELF, GETIVAR_MAX_DEPTH, side_exit);
 
@@ -877,69 +983,27 @@ gen_getinstancevariable(jitstate_t *jit, ctx_t *ctx)
 static codegen_status_t
 gen_setinstancevariable(jitstate_t* jit, ctx_t* ctx)
 {
-    IVC ic = (IVC)jit_get_arg(jit, 1);
+    // Defer compilation so we can specialize on a runtime `self`
-
+    if (!jit_at_current_insn(jit)) {
-    // Check that the inline cache has been set, slot index is known
+        defer_compilation(jit->block, jit->insn_idx, ctx);
-    if (!ic->entry) {
+        return YJIT_END_BLOCK;
-        return YJIT_CANT_COMPILE;
     }
 
-    // If the class uses the default allocator, instances should all be T_OBJECT
+    ID ivar_name = (ID)jit_get_arg(jit, 0);
-    // NOTE: This assumes nobody changes the allocator of the class after allocation.
-    //       Eventually, we can encode whether an object is T_OBJECT or not
-    //       inside object shapes.
-    if (rb_get_alloc_func(ic->entry->class_value) != rb_class_allocate_instance) {
-        return YJIT_CANT_COMPILE;
-    }
 
-    uint32_t ivar_index = ic->entry->index;
+    VALUE comptime_val = jit_peek_at_self(jit, ctx);
+    VALUE comptime_val_klass = CLASS_OF(comptime_val);
 
-    // Create a size-exit to fall back to the interpreter
+    // Generate a side exit
-    uint8_t* side_exit = yjit_side_exit(jit, ctx);
+    uint8_t *side_exit = yjit_side_exit(jit, ctx);
 
-    // Load self from CFP
+    // Guard that the receiver has the same class as the one from compile time.
     mov(cb, REG0, member_opnd(REG_CFP, rb_control_frame_t, self));
+    guard_self_is_heap(cb, REG0, COUNTED_EXIT(side_exit, setivar_se_self_not_heap), ctx);
 
-    guard_self_is_heap(cb, REG0, side_exit, ctx);
+    jit_guard_known_klass(jit, ctx, comptime_val_klass, OPND_SELF, GETIVAR_MAX_DEPTH, side_exit);
 
-    // Bail if receiver class is different from compiled time call cache class
+    return gen_set_ivar(jit, ctx, GETIVAR_MAX_DEPTH, comptime_val, ivar_name, OPND_SELF, side_exit);
-    x86opnd_t klass_opnd = mem_opnd(64, REG0, offsetof(struct RBasic, klass));
-    mov(cb, REG1, klass_opnd);
-    x86opnd_t serial_opnd = mem_opnd(64, REG1, offsetof(struct RClass, class_serial));
-    cmp(cb, serial_opnd, imm_opnd(ic->entry->class_serial));
-    jne_ptr(cb, side_exit);
-
-    // Bail if the ivars are not on the extended table
-    // See ROBJECT_IVPTR() from include/ruby/internal/core/robject.h
-    x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
-    test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
-    jnz_ptr(cb, side_exit);
-
-    // If we can't guarantee that the extended table is big enoughg
-    if (ivar_index >= ROBJECT_EMBED_LEN_MAX + 1) {
-        // Check that the slot is inside the extended table (num_slots > index)
-        x86opnd_t num_slots = mem_opnd(32, REG0, offsetof(struct RObject, as.heap.numiv));
-        cmp(cb, num_slots, imm_opnd(ivar_index));
-        jle_ptr(cb, side_exit);
-    }
-
-    // Get a pointer to the extended table
-    x86opnd_t tbl_opnd = mem_opnd(64, REG0, offsetof(struct RObject, as.heap.ivptr));
-    mov(cb, REG0, tbl_opnd);
-
-    // Pop the value to write from the stack
-    x86opnd_t stack_top = ctx_stack_pop(ctx, 1);
-    mov(cb, REG1, stack_top);
-
-    // Bail if this is a heap object, because this needs a write barrier
-    test(cb, REG1, imm_opnd(RUBY_IMMEDIATE_MASK));
-    jz_ptr(cb, side_exit);
-
-    // Write the ivar to the extended table
-    x86opnd_t ivar_opnd = mem_opnd(64, REG0, sizeof(VALUE) * ivar_index);
-    mov(cb, ivar_opnd, REG1);
-
-    return YJIT_KEEP_COMPILING;
 }
 
 static void

yjit_iface.h

@@ -57,6 +57,13 @@ YJIT_DECLARE_COUNTERS(
     getivar_name_not_mapped,
     getivar_not_object,
 
+    setivar_se_self_not_heap,
+    setivar_idx_out_of_range,
+    setivar_val_heapobject,
+    setivar_name_not_mapped,
+    setivar_not_object,
+    setivar_frozen,
+
     oaref_argc_not_one,
     oaref_arg_not_fixnum,