Mark and update object references in generated code

bootstraptest/test_ujit.rb

@@ -168,7 +168,8 @@ def foo
     end
 
     new.foo
-    UJIT.install_entry(RubyVM::InstructionSequence.of(instance_method(:foo)))
+    new.foo
+    new.foo
     new.foo
   end
 
@@ -199,18 +200,36 @@ def read
 }
 
 # Test that opt_aref checks the class of the receiver
-assert_equal ":special\n", %q{
+assert_equal 'special', %q{
   def foo(array)
     array[30]
   end
 
-  UJIT.install_entry(RubyVM::InstructionSequence.of(method(:foo)))
+  foo([])
+  foo([])
 
   special = []
   def special.[](idx)
-    :special
+    'special'
+  end
+
+  foo(special)
+}
+
+# Test that object references in generated code get marked and moved
+assert_equal "good", %q{
+  def bar
+    "good"
   end
 
-  p foo(special)
-  nil
+  def foo
+    bar
+  end
+
+  foo
+  foo
+
+  GC.verify_compaction_references(double_heap: true, toward: :empty)
+
+  foo
 }

ujit_codegen.c

@@ -59,6 +59,22 @@ jit_get_arg(jitstate_t* jit, size_t arg_idx)
     return *(jit->pc + arg_idx + 1);
 }
 
+// Load a pointer to a GC'd object into a register and keep track of the reference
+static void
+jit_mov_gc_ptr(jitstate_t* jit, codeblock_t* cb, x86opnd_t reg, VALUE ptr)
+{
+    RUBY_ASSERT(reg.type == OPND_REG && reg.num_bits == 64);
+    RUBY_ASSERT(!SPECIAL_CONST_P(ptr));
+
+    mov(cb, reg, const_ptr_opnd((void*)ptr));
+    // The pointer immediate is encoded as the last part of the mov written out.
+    uint32_t ptr_offset = cb->write_pos - sizeof(VALUE);
+
+    if (!rb_darray_append(&jit->block->gc_object_offsets, ptr_offset)) {
+        rb_bug("allocation failed");
+    }
+}
+
 /**
 Generate an inline exit to return to the interpreter
 */
@@ -1083,7 +1099,7 @@ gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_c
     assume_method_lookup_stable(cd->cc, cme, jit->block);
 
     // Bail if receiver class is different from compile-time call cache class
-    mov(cb, REG1, imm_opnd(cd->cc->klass));
+    jit_mov_gc_ptr(jit, cb, REG1, (VALUE)cd->cc->klass);
     cmp(cb, klass_opnd, REG1);
     jne_ptr(cb, side_exit);
 
@@ -1107,7 +1123,7 @@ gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_c
 
         // Put compile time cme into REG1. It's assumed to be valid because we are notified when
         // any cme we depend on become outdated. See rb_ujit_method_lookup_change().
-        mov(cb, REG1, const_ptr_opnd(cme));
+        jit_mov_gc_ptr(jit, cb, REG1, (VALUE)cme);
         // Write method entry at sp[-3]
         // sp[-3] = me;
         mov(cb, mem_opnd(64, REG0, 8 * -3), REG1);
@@ -1161,9 +1177,9 @@ gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_c
 
         // Call check_cfunc_dispatch
         mov(cb, RDI, recv);
-        mov(cb, RSI, const_ptr_opnd(cd));
+        jit_mov_gc_ptr(jit, cb, RSI, (VALUE)cd);
         mov(cb, RDX, const_ptr_opnd((void *)cfunc->func));
-        mov(cb, RCX, const_ptr_opnd(cme));
+        jit_mov_gc_ptr(jit, cb, RCX, (VALUE)cme);
         call_ptr(cb, REG0, (void *)&check_cfunc_dispatch);
 
         // Restore registers
@@ -1242,7 +1258,7 @@ gen_oswb_cfunc(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_c
 
 bool rb_simple_iseq_p(const rb_iseq_t *iseq);
 
-void
+static void
 gen_return_branch(codeblock_t* cb, uint8_t* target0, uint8_t* target1, uint8_t shape)
 {
     switch (shape)
@@ -1315,7 +1331,7 @@ gen_oswb_iseq(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_ca
     assume_method_lookup_stable(cd->cc, cme, jit->block);
 
     // Bail if receiver class is different from compile-time call cache class
-    mov(cb, REG1, imm_opnd(cd->cc->klass));
+    jit_mov_gc_ptr(jit, cb, REG1, (VALUE)cd->cc->klass);
     cmp(cb, klass_opnd, REG1);
     jne_ptr(cb, side_exit);
 
@@ -1343,7 +1359,7 @@ gen_oswb_iseq(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_ca
 
     // Put compile time cme into REG1. It's assumed to be valid because we are notified when
     // any cme we depend on become outdated. See rb_ujit_method_lookup_change().
-    mov(cb, REG1, const_ptr_opnd(cme));
+    jit_mov_gc_ptr(jit, cb, REG1, (VALUE)cme);
     // Write method entry at sp[-3]
     // sp[-3] = me;
     mov(cb, mem_opnd(64, REG0, 8 * -3), REG1);
@@ -1378,7 +1394,7 @@ gen_oswb_iseq(jitstate_t* jit, ctx_t* ctx, struct rb_call_data * cd, const rb_ca
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, ep), REG0);
     mov(cb, REG0, recv);
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, self), REG0);
-    mov(cb, REG0, const_ptr_opnd(iseq));
+    jit_mov_gc_ptr(jit, cb, REG0, (VALUE)iseq);
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, iseq), REG0);
     mov(cb, REG0, const_ptr_opnd(start_pc));
     mov(cb, member_opnd(REG_CFP, rb_control_frame_t, pc), REG0);

ujit_core.c

@@ -155,7 +155,7 @@ get_first_version(const rb_iseq_t *iseq, unsigned idx)
     return rb_darray_get(body->ujit_blocks, idx);
 }
 
-// Add a block version to the map. Block should be fully constructed
+// Keep track of a block version. Block should be fully constructed.
 static void
 add_block_version(blockid_t blockid, block_t* block)
 {
@@ -195,6 +195,17 @@ add_block_version(blockid_t blockid, block_t* block)
         RB_OBJ_WRITTEN(iseq, Qundef, block->dependencies.iseq);
         RB_OBJ_WRITTEN(iseq, Qundef, block->dependencies.cc);
         RB_OBJ_WRITTEN(iseq, Qundef, block->dependencies.cme);
+
+        // Run write barrier for all objects in generated code.
+        uint32_t *offset_element;
+        rb_darray_foreach(block->gc_object_offsets, offset_idx, offset_element) {
+            uint32_t offset_to_value = *offset_element;
+            uint8_t *value_address = cb_get_ptr(cb, offset_to_value);
+
+            VALUE object;
+            memcpy(&object, value_address, SIZEOF_VALUE);
+            RB_OBJ_WRITTEN(iseq, Qundef, object);
+        }
     }
 }
 
@@ -589,6 +600,7 @@ ujit_free_block(block_t *block)
     ujit_unlink_method_lookup_dependency(block);
     rb_darray_free(block->incoming);
     free(block);
+    rb_darray_free(block->gc_object_offsets);
 }
 
 // Invalidate one specific block version

ujit_core.h

@@ -123,12 +123,12 @@ typedef struct ujit_block_version
     // List of incoming branches indices
     int32_array_t incoming;
 
+    // Offsets for GC managed objects in the mainline code block
+    int32_array_t gc_object_offsets;
+
     // Next block version for this blockid (singly-linked list)
     struct ujit_block_version *next;
 
-    // List node for all block versions in an iseq
-    struct list_node iseq_block_node;
-
     // GC managed objects that this block depend on
     struct {
         VALUE cc;

ujit_iface.c

@@ -566,6 +566,17 @@ rb_ujit_iseq_mark(const struct rb_iseq_constant_body *body)
             rb_gc_mark_movable(block->dependencies.cc);
             rb_gc_mark_movable(block->dependencies.cme);
             rb_gc_mark_movable(block->dependencies.iseq);
+
+            // Walk over references to objects in generated code.
+            uint32_t *offset_element;
+            rb_darray_foreach(block->gc_object_offsets, offset_idx, offset_element) {
+                uint32_t offset_to_value = *offset_element;
+                uint8_t *value_address = cb_get_ptr(cb, offset_to_value);
+
+                VALUE object;
+                memcpy(&object, value_address, SIZEOF_VALUE);
+                rb_gc_mark_movable(object);
+            }
         }
     }
 }
@@ -581,6 +592,21 @@ rb_ujit_iseq_update_references(const struct rb_iseq_constant_body *body)
             block->dependencies.cc = rb_gc_location(block->dependencies.cc);
             block->dependencies.cme = rb_gc_location(block->dependencies.cme);
             block->dependencies.iseq = rb_gc_location(block->dependencies.iseq);
+
+            // Walk over references to objects in generated code.
+            uint32_t *offset_element;
+            rb_darray_foreach(block->gc_object_offsets, offset_idx, offset_element) {
+                uint32_t offset_to_value = *offset_element;
+                uint8_t *value_address = cb_get_ptr(cb, offset_to_value);
+
+                VALUE object;
+                memcpy(&object, value_address, SIZEOF_VALUE);
+                VALUE possibly_moved = rb_gc_location(object);
+                // Only write when the VALUE moves, to be CoW friendly.
+                if (possibly_moved != object) {
+                    memcpy(value_address, &possibly_moved, SIZEOF_VALUE);
+                }
+            }
         }
     }
 }