Fix rb_copy_generic_ivar to handle object_id in shapes

Author: byroot

gc.c

@@ -1835,14 +1835,17 @@ object_id(VALUE obj)
 
     if (rb_shape_too_complex_p(shape)) {
         st_table *table = ROBJECT_IV_HASH(obj);
-        if (st_lookup(table, (st_data_t)internal_object_id, (st_data_t *)&id)) {
+        if (rb_shape_has_object_id(shape)) {
+            st_lookup(table, (st_data_t)internal_object_id, (st_data_t *)&id);
             rb_gc_vm_unlock(lock_lev);
             return id;
         }
 
         id = ULL2NUM(next_object_id);
         next_object_id += OBJ_ID_INCREMENT;
+        rb_shape_t *object_id_shape = rb_shape_object_id_shape(obj);
         st_insert(table, (st_data_t)internal_object_id, (st_data_t)id);
+        rb_shape_set_shape(obj, object_id_shape);
         if (RB_UNLIKELY(id_to_obj_tbl)) {
             st_insert(id_to_obj_tbl, (st_data_t)id, (st_data_t)obj);
         }

object.c

@@ -357,6 +357,7 @@ rb_obj_copy_ivar(VALUE dest, VALUE obj)
 
         shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
         if (UNLIKELY(rb_shape_too_complex_p(shape_to_set_on_dest))) {
+            rb_bug("here");
             st_table *table = rb_st_init_numtable_with_size(src_num_ivs);
             rb_obj_copy_ivs_to_hash_table(obj, table);
             rb_obj_init_too_complex(dest, table);

shape.c

@@ -1072,7 +1072,7 @@ rb_shape_rebuild_shape(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
 
     rb_shape_t *midway_shape;
 
-    RUBY_ASSERT(initial_shape->type == SHAPE_T_OBJECT);
+    RUBY_ASSERT(initial_shape->type == SHAPE_T_OBJECT || initial_shape->type == SHAPE_ROOT);
 
     if (dest_shape->type != initial_shape->type) {
         midway_shape = rb_shape_rebuild_shape(initial_shape, rb_shape_get_parent(dest_shape));

test/ruby/test_object_id.rb

@@ -0,0 +1,135 @@
+require 'test/unit'
+
+class TestObjectId < Test::Unit::TestCase
+  def setup
+    @obj = Object.new
+  end
+
+  def test_dup_new_id
+    id = @obj.object_id
+    refute_equal id, @obj.dup.object_id
+  end
+
+  def test_dup_with_ivar_and_id
+    id = @obj.object_id
+    @obj.instance_variable_set(:@foo, 42)
+
+    copy = @obj.dup
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+  end
+
+  def test_dup_with_id_and_ivar
+    @obj.instance_variable_set(:@foo, 42)
+    id = @obj.object_id
+
+    copy = @obj.dup
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+  end
+
+  def test_dup_with_id_and_ivar_and_frozen
+    @obj.instance_variable_set(:@foo, 42)
+    @obj.freeze
+    id = @obj.object_id
+
+    copy = @obj.dup
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+    refute_predicate copy, :frozen?
+  end
+
+  def test_clone_new_id
+    id = @obj.object_id
+    refute_equal id, @obj.clone.object_id
+  end
+
+  def test_clone_with_ivar_and_id
+    id = @obj.object_id
+    @obj.instance_variable_set(:@foo, 42)
+
+    copy = @obj.clone
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+  end
+
+  def test_clone_with_id_and_ivar
+    @obj.instance_variable_set(:@foo, 42)
+    id = @obj.object_id
+
+    copy = @obj.clone
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+  end
+
+  def test_clone_with_id_and_ivar_and_frozen
+    @obj.instance_variable_set(:@foo, 42)
+    @obj.freeze
+    id = @obj.object_id
+
+    copy = @obj.clone
+    refute_equal id, copy.object_id
+    assert_equal 42, copy.instance_variable_get(:@foo)
+    assert_predicate copy, :frozen?
+  end
+end
+
+class TestObjectIdClass < TestObjectId
+  def setup
+    @obj = Class.new
+  end
+end
+
+class TestObjectIdGeneric < TestObjectId
+  def setup
+    @obj = Array.new
+  end
+end
+
+class TestObjectIdTooComplex < TestObjectId
+  class TooComplex
+  end
+
+  def setup
+    if defined?(RubyVM::Shape::SHAPE_MAX_VARIATIONS)
+      assert_equal 8, RubyVM::Shape::SHAPE_MAX_VARIATIONS
+    end
+    8.times do |i|
+      TooComplex.new.instance_variable_set("@a#{i}", 1)
+    end
+    @obj = TooComplex.new
+    @obj.instance_variable_set(:@test, 1)
+  end
+end
+
+class TestObjectIdTooComplexClass < TestObjectId
+  class TooComplex < Module
+  end
+
+  def setup
+    if defined?(RubyVM::Shape::SHAPE_MAX_VARIATIONS)
+      assert_equal 8, RubyVM::Shape::SHAPE_MAX_VARIATIONS
+    end
+    8.times do |i|
+      TooComplex.new.instance_variable_set("@a#{i}", 1)
+    end
+    @obj = TooComplex.new
+    @obj.instance_variable_set(:@test, 1)
+  end
+end
+
+class TestObjectIdTooComplexGeneric < TestObjectId
+  class TooComplex < Array
+  end
+
+  def setup
+    if defined?(RubyVM::Shape::SHAPE_MAX_VARIATIONS)
+      assert_equal 8, RubyVM::Shape::SHAPE_MAX_VARIATIONS
+    end
+    8.times do |i|
+      TooComplex.new.instance_variable_set("@a#{i}", 1)
+    end
+    @obj = TooComplex.new
+    @obj.instance_variable_set(:@test, 1)
+  end
+end

variable.c

@@ -2276,35 +2276,83 @@ class_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg,
 }
 
 void
-rb_copy_generic_ivar(VALUE clone, VALUE obj)
+rb_copy_generic_ivar(VALUE dest, VALUE obj)
 {
     struct gen_ivtbl *obj_ivtbl;
     struct gen_ivtbl *new_ivtbl;
 
-    rb_check_frozen(clone);
+    rb_check_frozen(dest);
 
     if (!FL_TEST(obj, FL_EXIVAR)) {
         goto clear;
     }
 
+    unsigned long src_num_ivs = rb_ivar_count(obj);
+    if (!src_num_ivs) {
+        goto clear;
+    }
+
+    rb_shape_t *src_shape = rb_shape_get_shape(obj);
+
     if (rb_gen_ivtbl_get(obj, 0, &obj_ivtbl)) {
         if (gen_ivtbl_count(obj, obj_ivtbl) == 0)
             goto clear;
 
-        FL_SET(clone, FL_EXIVAR);
+        FL_SET(dest, FL_EXIVAR);
 
-        if (rb_shape_obj_too_complex(obj)) {
-            new_ivtbl = xmalloc(sizeof(struct gen_ivtbl));
-#if !SHAPE_IN_BASIC_FLAGS
-            new_ivtbl->shape_id = SHAPE_OBJ_TOO_COMPLEX;
-#endif
-            new_ivtbl->as.complex.table = st_copy(obj_ivtbl->as.complex.table);
+        if (rb_shape_too_complex_p(src_shape)) {
+            // obj is TOO_COMPLEX so we can copy its iv_hash
+            st_table *table = st_copy(ROBJECT_IV_HASH(obj));
+            if (rb_shape_has_object_id(src_shape)) {
+                st_data_t id = (st_data_t)internal_object_id;
+                st_delete(table, &id, NULL);
+            }
+            rb_obj_init_too_complex(dest, table);
+
+            return;
+        }
+        
+        rb_shape_t *shape_to_set_on_dest = src_shape;
+        rb_shape_t *initial_shape = rb_shape_get_shape(dest);
+
+        if (initial_shape->heap_index != src_shape->heap_index || !rb_shape_canonical_p(src_shape)) {
+            RUBY_ASSERT(initial_shape->type == SHAPE_ROOT);
+
+            shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
+            if (UNLIKELY(rb_shape_too_complex_p(shape_to_set_on_dest))) {
+                st_table *table = rb_st_init_numtable_with_size(src_num_ivs);
+                rb_obj_copy_ivs_to_hash_table(obj, table);
+                rb_obj_init_too_complex(dest, table);
+
+                return;
+            }
+        }
+
+        new_ivtbl = gen_ivtbl_resize(0, shape_to_set_on_dest->next_field_index);
+
+        VALUE *src_buf = obj_ivtbl->as.shape.ivptr;
+        VALUE *dest_buf = new_ivtbl->as.shape.ivptr;
+
+        if (src_shape->next_field_index == shape_to_set_on_dest->next_field_index) {
+            // Happy path, we can just memcpy the ivptr content
+            MEMCPY(dest_buf, src_buf, VALUE, src_num_ivs);
+
+            // Fire write barriers
+            for (uint32_t i = 0; i < src_num_ivs; i++) {
+                RB_OBJ_WRITTEN(dest, Qundef, dest_buf[i]);
+            }
         }
         else {
-            new_ivtbl = gen_ivtbl_resize(0, obj_ivtbl->as.shape.numiv);
+            rb_shape_t *dest_shape = shape_to_set_on_dest;
+            while (src_shape->parent_id != INVALID_SHAPE_ID) {
+                if (src_shape->type == SHAPE_IVAR) {
+                    while (dest_shape->edge_name != src_shape->edge_name) {
+                        dest_shape = rb_shape_get_shape_by_id(dest_shape->parent_id);
+                    }
 
-            for (uint32_t i=0; i<obj_ivtbl->as.shape.numiv; i++) {
-                RB_OBJ_WRITE(clone, &new_ivtbl->as.shape.ivptr[i], obj_ivtbl->as.shape.ivptr[i]);
+                    RB_OBJ_WRITE(dest, &dest_buf[dest_shape->next_field_index - 1], src_buf[src_shape->next_field_index - 1]);
+                }
+                src_shape = rb_shape_get_shape_by_id(src_shape->parent_id);
             }
         }
 
@@ -2314,25 +2362,18 @@ rb_copy_generic_ivar(VALUE clone, VALUE obj)
          */
         RB_VM_LOCK_ENTER();
         {
-            generic_ivtbl_no_ractor_check(clone);
-            st_insert(generic_ivtbl_no_ractor_check(obj), (st_data_t)clone, (st_data_t)new_ivtbl);
+            generic_ivtbl_no_ractor_check(dest);
+            st_insert(generic_ivtbl_no_ractor_check(obj), (st_data_t)dest, (st_data_t)new_ivtbl);
         }
         RB_VM_LOCK_LEAVE();
-
-        rb_shape_t * obj_shape = rb_shape_get_shape(obj);
-        if (rb_shape_frozen_shape_p(obj_shape)) {
-            rb_shape_set_shape_id(clone, obj_shape->parent_id);
-        }
-        else {
-            rb_shape_set_shape(clone, obj_shape);
-        }
+        rb_shape_set_shape(dest, shape_to_set_on_dest);
     }
     return;
 
   clear:
-    if (FL_TEST(clone, FL_EXIVAR)) {
-        rb_free_generic_ivar(clone);
-        FL_UNSET(clone, FL_EXIVAR);
+    if (FL_TEST(dest, FL_EXIVAR)) {
+        rb_free_generic_ivar(dest);
+        FL_UNSET(dest, FL_EXIVAR);
     }
 }