Implement setivar method calls

yjit_codegen.c

@@ -1430,7 +1430,6 @@ enum {
     SEND_MAX_DEPTH = 5,           // up to 5 different classes
 };
 
-/*
 // Codegen for setting an instance variable.
 // Preconditions:
 //   - receiver is in REG0
@@ -1442,6 +1441,122 @@ gen_set_ivar(jitstate_t *jit, ctx_t *ctx, const int max_chain_depth, VALUE compt
     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_TYPE_P(comptime_receiver, T_OBJECT) ||
+            rb_get_alloc_func(comptime_val_klass) != rb_class_allocate_instance) {
+        // General case. Call rb_ivar_get(). No need to reconstruct interpreter
+        // state since the routine never raises exceptions or allocate objects
+        // visibile to Ruby.
+        // VALUE rb_ivar_set(VALUE obj, ID id, VALUE val)
+        ADD_COMMENT(cb, "call rb_ivar_set()");
+        mov(cb, C_ARG_REGS[0], REG0);
+        mov(cb, C_ARG_REGS[1], imm_opnd((int64_t)ivar_name));
+        mov(cb, C_ARG_REGS[2], ctx_stack_pop(ctx, 1));
+        call_ptr(cb, REG1, (void *)rb_ivar_set);
+
+        if (!reg0_opnd.is_self) {
+            (void)ctx_stack_pop(ctx, 1);
+        }
+        // FIXME: setting an ivar pushes the same value back on the stack, so we shouldn't
+        // pop and push.
+        // Push the ivar on the stack
+        x86opnd_t out_opnd = ctx_stack_push(ctx, TYPE_UNKNOWN);
+        mov(cb, out_opnd, RAX);
+
+        // Jump to next instruction. This allows guard chains to share the same successor.
+        jit_jump_to_next_insn(jit, ctx);
+        return YJIT_END_BLOCK;
+    }
+
+    // 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);
+
+    // 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;
+
+        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");
+            x86opnd_t flags_opnd = member_opnd(REG0, struct RBasic, flags);
+            test(cb, flags_opnd, imm_opnd(ROBJECT_EMBED));
+            jit_chain_guard(JCC_JZ, jit, &starting_context, max_chain_depth, side_exit);
+
+            // Write the variable
+            x86opnd_t ivar_opnd = mem_opnd(64, REG0, offsetof(struct RObject, as.ary) + ivar_index * SIZEOF_VALUE);
+            mov(cb, REG1, ctx_stack_pop(ctx, 1));
+            mov(cb, ivar_opnd, REG1);
+
+            // Pop receiver if it's on the temp stack
+            // ie. this is an attribute method
+            if (!reg0_opnd.is_self) {
+              ctx_stack_pop(ctx, 1);
+            }
+
+            // Push the ivar on the stack
+            x86opnd_t out_opnd = ctx_stack_push(ctx, TYPE_UNKNOWN);
+            mov(cb, out_opnd, REG1);
+        }
+        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, getivar_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);
+
+            // Read the ivar from the extended table
+            x86opnd_t ivar_opnd = mem_opnd(64, REG0, sizeof(VALUE) * ivar_index);
+            mov(cb, REG1, ctx_stack_pop(ctx, 1));
+            mov(cb, ivar_opnd, REG1);
+
+            // Pop receiver if it's on the temp stack
+            // ie. this is an attribute method
+            if (!reg0_opnd.is_self) {
+              ctx_stack_pop(ctx, 1);
+            }
+
+            // Push the ivar on the stack
+            x86opnd_t out_opnd = ctx_stack_push(ctx, TYPE_UNKNOWN);
+            mov(cb, out_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;
+}
+
+/*
+{
+    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
@@ -3459,7 +3574,15 @@ gen_send_general(jitstate_t *jit, ctx_t *ctx, struct rb_call_data *cd, rb_iseq_t
             }
         case VM_METHOD_TYPE_ATTRSET:
             GEN_COUNTER_INC(cb, send_ivar_set_method);
-            return YJIT_CANT_COMPILE;
+
+            if (argc != 1) {
+                return YJIT_CANT_COMPILE;
+            } else {
+                mov(cb, REG0, recv);
+
+                ID ivar_name = cme->def->body.attr.id;
+                return gen_set_ivar(jit, ctx, SEND_MAX_DEPTH, comptime_recv, ivar_name, recv_opnd, side_exit);
+            }
         case VM_METHOD_TYPE_BMETHOD:
             GEN_COUNTER_INC(cb, send_bmethod);
             return YJIT_CANT_COMPILE;