Transition shape when object's capacity changes

This commit adds a `capacity` field to shapes, and adds shape
transitions whenever an object's capacity changes. Objects which are
allocated out of a bigger size pool will also make a transition from the
root shape to the shape with the correct capacity for their size pool
when they are allocated.

This commit will allow us to remove numiv from objects completely, and
will also mean we can guarantee that if two objects share shapes, their
IVs are in the same positions (an embedded and extended object cannot
share shapes). This will enable us to implement ivar sets in YJIT using
object shapes.

Co-Authored-By: Aaron Patterson <[email protected]>

1 files changed, 54 insertions, 23 deletions

diff --git a/object.c b/object.c
index aa337ea2ce..9a06500b6b 100644
--- a/object.c
+++ b/object.c
@@ -33,6 +33,7 @@
 #include "internal/string.h"
 #include "internal/symbol.h"
 #include "internal/variable.h"
+#include "variable.h"
 #include "probes.h"
 #include "ruby/encoding.h"
 #include "ruby/st.h"
@@ -268,21 +269,64 @@ rb_obj_singleton_class(VALUE obj)
 MJIT_FUNC_EXPORTED void
 rb_obj_copy_ivar(VALUE dest, VALUE obj)
 {
-    uint32_t dest_len = ROBJECT_NUMIV(dest);
-    uint32_t src_len = ROBJECT_NUMIV(obj);
+    RUBY_ASSERT(!RB_TYPE_P(obj, T_CLASS) && !RB_TYPE_P(obj, T_MODULE));
 
-    if (dest_len < src_len) {
-        rb_ensure_iv_list_size(dest, dest_len, src_len);
-        RUBY_ASSERT(!(RBASIC(dest)->flags & ROBJECT_EMBED));
+    RUBY_ASSERT(BUILTIN_TYPE(dest) == BUILTIN_TYPE(obj));
+    uint32_t src_num_ivs = RBASIC_IV_COUNT(obj);
+    rb_shape_t * src_shape = rb_shape_get_shape(obj);
+    rb_shape_t * shape_to_set_on_dest = src_shape;
+    VALUE * src_buf;
+    VALUE * dest_buf;
+
+    if (!src_num_ivs) {
+        return;
     }
-    else {
-        RUBY_ASSERT((RBASIC(dest)->flags & ROBJECT_EMBED));
+
+    // The copy should be mutable, so we don't want the frozen shape
+    if (rb_shape_frozen_shape_p(src_shape)) {
+        shape_to_set_on_dest = rb_shape_get_shape_by_id(src_shape->parent_id);
+    }
+
+    src_buf = ROBJECT_IVPTR(obj);
+    dest_buf = ROBJECT_IVPTR(dest);
+
+    rb_shape_t * initial_shape = rb_shape_get_shape(dest);
+
+    if (initial_shape->size_pool_index != src_shape->size_pool_index) {
+        RUBY_ASSERT(initial_shape->parent_id == ROOT_SHAPE_ID || initial_shape->type == SHAPE_ROOT);
+
+        shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
     }
 
-    VALUE * dest_buf = ROBJECT_IVPTR(dest);
-    VALUE * src_buf = ROBJECT_IVPTR(obj);
+    RUBY_ASSERT(src_num_ivs <= shape_to_set_on_dest->capacity);
+    if (initial_shape->capacity < shape_to_set_on_dest->capacity) {
+        rb_ensure_iv_list_size(dest, initial_shape->capacity, shape_to_set_on_dest->capacity);
+        dest_buf = ROBJECT_IVPTR(dest);
+
+        rb_shape_t * initial_shape = rb_shape_get_shape(dest);
+
+        if (initial_shape->size_pool_index != src_shape->size_pool_index) {
+            RUBY_ASSERT(initial_shape->parent_id == ROOT_SHAPE_ID || initial_shape->type == SHAPE_ROOT);
+
+            shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
+        }
 
-    MEMCPY(dest_buf, src_buf, VALUE, ROBJECT_IV_COUNT(obj));
+        RUBY_ASSERT(src_num_ivs <= shape_to_set_on_dest->capacity);
+        if (initial_shape->capacity < shape_to_set_on_dest->capacity) {
+            rb_ensure_iv_list_size(dest, initial_shape->capacity, shape_to_set_on_dest->capacity);
+            dest_buf = ROBJECT_IVPTR(dest);
+        }
+    }
+
+    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]);
+    }
+
+    rb_shape_set_shape(dest, shape_to_set_on_dest);
+    RUBY_ASSERT(!RB_TYPE_P(obj, T_OBJECT) || ROBJECT_IV_CAPACITY(dest) == ROBJECT_NUMIV(dest));
 }
 
 static void
@@ -301,19 +345,6 @@ init_copy(VALUE dest, VALUE obj)
     if (RB_TYPE_P(obj, T_OBJECT)) {
         rb_obj_copy_ivar(dest, obj);
     }
-
-    if (!RB_TYPE_P(obj, T_CLASS) && !RB_TYPE_P(obj, T_MODULE)) {
-        rb_shape_t *shape_to_set = rb_shape_get_shape(obj);
-
-        // If the object is frozen, the "dup"'d object will *not* be frozen,
-        // so we need to copy the frozen shape's parent to the new object.
-        if (rb_shape_frozen_shape_p(shape_to_set)) {
-            shape_to_set = rb_shape_get_shape_by_id(shape_to_set->parent_id);
-        }
-
-        // shape ids are different
-        rb_shape_set_shape(dest, shape_to_set);
-    }
 }
 
 static VALUE immutable_obj_clone(VALUE obj, VALUE kwfreeze);