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 <tenderlove@ruby-lang.org>

1 files changed, 158 insertions, 45 deletions

diff --git a/shape.c b/shape.c
index 1de89d3f8f..e19667ae2c 100644
--- a/shape.c
+++ b/shape.c
@@ -1,15 +1,19 @@
 #include "vm_core.h"
 #include "vm_sync.h"
 #include "shape.h"
+#include "gc.h"
 #include "internal/class.h"
 #include "internal/symbol.h"
 #include "internal/variable.h"
 #include <stdbool.h>
 
+static ID id_frozen;
+static ID size_pool_edge_names[SIZE_POOL_COUNT];
+
 /*
  * Shape getters
  */
-static rb_shape_t*
+rb_shape_t *
 rb_shape_get_root_shape(void)
 {
     return GET_VM()->root_shape;
@@ -21,12 +25,6 @@ rb_shape_id(rb_shape_t * shape)
     return (shape_id_t)(shape - GET_VM()->shape_list);
 }
 
-static rb_shape_t*
-rb_shape_get_frozen_root_shape(void)
-{
-    return GET_VM()->frozen_root_shape;
-}
-
 bool
 rb_shape_root_shape_p(rb_shape_t* shape)
 {
@@ -68,7 +66,7 @@ shape_id_t
 rb_shape_get_shape_id(VALUE obj)
 {
     if (RB_SPECIAL_CONST_P(obj)) {
-        return FROZEN_ROOT_SHAPE_ID;
+        return SPECIAL_CONST_SHAPE_ID;
     }
 
 #if SHAPE_IN_BASIC_FLAGS
@@ -113,12 +111,9 @@ rb_shape_lookup_id(rb_shape_t* shape, ID id, enum shape_type shape_type)
 }
 
 static rb_shape_t*
-get_next_shape_internal(rb_shape_t* shape, ID id, VALUE obj, enum shape_type shape_type)
+get_next_shape_internal(rb_shape_t * shape, ID id, enum shape_type shape_type)
 {
     rb_shape_t *res = NULL;
-
-    RUBY_ASSERT(SHAPE_FROZEN != (enum shape_type)shape->type || RB_TYPE_P(obj, T_MODULE) || RB_TYPE_P(obj, T_CLASS));
-
     RB_VM_LOCK_ENTER();
     {
         if (rb_shape_lookup_id(shape, id, shape_type)) {
@@ -142,23 +137,18 @@ get_next_shape_internal(rb_shape_t* shape, ID id, VALUE obj, enum shape_type sha
                 rb_shape_t * new_shape = rb_shape_alloc(id, shape);
 
                 new_shape->type = (uint8_t)shape_type;
+                new_shape->capacity = shape->capacity;
 
                 switch (shape_type) {
                   case SHAPE_IVAR:
-                    new_shape->next_iv_index = rb_shape_get_shape_by_id(new_shape->parent_id)->next_iv_index + 1;
-
-                    // Check if we should update next_iv_index on the object's class
-                    if (BUILTIN_TYPE(obj) == T_OBJECT) {
-                        VALUE klass = rb_obj_class(obj);
-                        if (new_shape->next_iv_index > RCLASS_EXT(klass)->max_iv_count) {
-                            RCLASS_EXT(klass)->max_iv_count = new_shape->next_iv_index;
-                        }
-                    }
+                    new_shape->next_iv_index = shape->next_iv_index + 1;
                     break;
+                  case SHAPE_CAPACITY_CHANGE:
                   case SHAPE_IVAR_UNDEF:
                   case SHAPE_FROZEN:
-                    new_shape->next_iv_index = rb_shape_get_shape_by_id(new_shape->parent_id)->next_iv_index;
+                    new_shape->next_iv_index = shape->next_iv_index;
                     break;
+                  case SHAPE_INITIAL_CAPACITY:
                   case SHAPE_ROOT:
                     rb_bug("Unreachable");
                     break;
@@ -183,7 +173,7 @@ rb_shape_frozen_shape_p(rb_shape_t* shape)
 void
 rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape)
 {
-    rb_shape_t* next_shape = get_next_shape_internal(shape, id, obj, SHAPE_IVAR_UNDEF);
+    rb_shape_t * next_shape = get_next_shape_internal(shape, id, SHAPE_IVAR_UNDEF);
 
     if (shape == next_shape) {
         return;
@@ -206,16 +196,11 @@ rb_shape_transition_shape_frozen(VALUE obj)
     rb_shape_t* next_shape;
 
     if (shape == rb_shape_get_root_shape()) {
-        next_shape = rb_shape_get_frozen_root_shape();
+        rb_shape_set_shape_id(obj, SPECIAL_CONST_SHAPE_ID);
+        return;
     }
-    else {
-        static ID id_frozen;
-        if (!id_frozen) {
-            id_frozen = rb_make_internal_id();
-        }
 
-        next_shape = get_next_shape_internal(shape, (ID)id_frozen, obj, SHAPE_FROZEN);
-    }
+    next_shape = get_next_shape_internal(shape, (ID)id_frozen, SHAPE_FROZEN);
 
     RUBY_ASSERT(next_shape);
     rb_shape_set_shape(obj, next_shape);
@@ -231,10 +216,39 @@ rb_shape_transition_shape(VALUE obj, ID id, rb_shape_t *shape)
     rb_shape_set_shape(obj, next_shape);
 }
 
-rb_shape_t*
+/*
+ * This function is used for assertions where we don't want to increment
+ * max_iv_count
+ */
+rb_shape_t *
+rb_shape_get_next_iv_shape(rb_shape_t* shape, ID id)
+{
+    return get_next_shape_internal(shape, id, SHAPE_IVAR);
+}
+
+rb_shape_t *
 rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
 {
-    return get_next_shape_internal(shape, id, obj, SHAPE_IVAR);
+    rb_shape_t * new_shape = rb_shape_get_next_iv_shape(shape, id);
+
+    // Check if we should update max_iv_count on the object's class
+    if (BUILTIN_TYPE(obj) == T_OBJECT) {
+        VALUE klass = rb_obj_class(obj);
+        if (new_shape->next_iv_index > RCLASS_EXT(klass)->max_iv_count) {
+            RCLASS_EXT(klass)->max_iv_count = new_shape->next_iv_index;
+        }
+    }
+
+    return new_shape;
+}
+
+rb_shape_t *
+rb_shape_transition_shape_capa(rb_shape_t* shape, uint32_t new_capacity)
+{
+    ID edge_name = rb_make_temporary_id(new_capacity);
+    rb_shape_t * new_shape = get_next_shape_internal(shape, edge_name, SHAPE_CAPACITY_CHANGE);
+    new_shape->capacity = new_capacity;
+    return new_shape;
 }
 
 bool
@@ -250,11 +264,13 @@ rb_shape_get_iv_index(rb_shape_t * shape, ID id, attr_index_t *value)
                 RUBY_ASSERT(shape->next_iv_index > 0);
                 *value = shape->next_iv_index - 1;
                 return true;
+              case SHAPE_CAPACITY_CHANGE:
               case SHAPE_IVAR_UNDEF:
               case SHAPE_ROOT:
+              case SHAPE_INITIAL_CAPACITY:
                 return false;
               case SHAPE_FROZEN:
-                rb_bug("Ivar should not exist on frozen transition\n");
+                rb_bug("Ivar should not exist on transition\n");
             }
         }
         shape = rb_shape_get_shape_by_id(shape->parent_id);
@@ -290,9 +306,18 @@ rb_shape_alloc_with_parent_id(ID edge_name, shape_id_t parent_id)
 }
 
 rb_shape_t *
+rb_shape_alloc_with_size_pool_index(ID edge_name, rb_shape_t * parent, uint8_t size_pool_index)
+{
+    rb_shape_t * shape = rb_shape_alloc_with_parent_id(edge_name, rb_shape_id(parent));
+    shape->size_pool_index = size_pool_index;
+    return shape;
+}
+
+
+rb_shape_t *
 rb_shape_alloc(ID edge_name, rb_shape_t * parent)
 {
-    return rb_shape_alloc_with_parent_id(edge_name, rb_shape_id(parent));
+    return rb_shape_alloc_with_size_pool_index(edge_name, parent, parent->size_pool_index);
 }
 
 MJIT_FUNC_EXPORTED void
@@ -307,6 +332,39 @@ rb_shape_flags_mask(void)
     return SHAPE_FLAG_MASK;
 }
 
+rb_shape_t *
+rb_shape_rebuild_shape(rb_shape_t * initial_shape, rb_shape_t * dest_shape)
+{
+    rb_shape_t * midway_shape;
+
+    if (dest_shape->type != SHAPE_ROOT) {
+        midway_shape = rb_shape_rebuild_shape(initial_shape, rb_shape_get_shape_by_id(dest_shape->parent_id));
+    }
+    else {
+        midway_shape = initial_shape;
+    }
+
+    switch (dest_shape->type) {
+        case SHAPE_IVAR:
+            if (midway_shape->capacity < midway_shape->next_iv_index) {
+                // There isn't enough room to write this IV, so we need to increase the capacity
+                midway_shape = rb_shape_transition_shape_capa(midway_shape, midway_shape->capacity * 2);
+            }
+
+            midway_shape = rb_shape_get_next_iv_shape(midway_shape, dest_shape->edge_name);
+            break;
+        case SHAPE_IVAR_UNDEF:
+            midway_shape = get_next_shape_internal(midway_shape, dest_shape->edge_name, SHAPE_IVAR_UNDEF);
+            break;
+        case SHAPE_ROOT:
+        case SHAPE_FROZEN:
+        case SHAPE_CAPACITY_CHANGE:
+            break;
+    }
+
+    return midway_shape;
+}
+
 #if VM_CHECK_MODE > 0
 VALUE rb_cShape;
 
@@ -336,6 +394,14 @@ rb_shape_type(VALUE self)
 }
 
 static VALUE
+rb_shape_capacity(VALUE self)
+{
+    rb_shape_t * shape;
+    TypedData_Get_Struct(self, rb_shape_t, &shape_data_type, shape);
+    return INT2NUM(shape->capacity);
+}
+
+static VALUE
 rb_shape_parent_id(VALUE self)
 {
     rb_shape_t * shape;
@@ -398,11 +464,16 @@ rb_shape_edge_name(VALUE self)
     rb_shape_t* shape;
     TypedData_Get_Struct(self, rb_shape_t, &shape_data_type, shape);
 
-    if (shape->edge_name) {
-        return ID2SYM(shape->edge_name);
+    if ((shape->edge_name & (ID_INTERNAL)) == ID_INTERNAL) {
+        return INT2NUM(shape->capacity);
     }
     else {
-        return Qnil;
+        if (shape->edge_name) {
+            return ID2SYM(shape->edge_name);
+        }
+        else {
+            return Qnil;
+        }
     }
 }
 
@@ -416,6 +487,15 @@ rb_shape_next_iv_index(VALUE self)
 }
 
 static VALUE
+rb_shape_size_pool_index(VALUE self)
+{
+    rb_shape_t * shape;
+    TypedData_Get_Struct(self, rb_shape_t, &shape_data_type, shape);
+
+    return INT2NUM(shape->size_pool_index);
+}
+
+static VALUE
 rb_shape_export_depth(VALUE self)
 {
     rb_shape_t* shape;
@@ -454,12 +534,6 @@ rb_shape_root_shape(VALUE self)
     return rb_shape_t_to_rb_cShape(rb_shape_get_root_shape());
 }
 
-static VALUE
-rb_shape_frozen_root_shape(VALUE self)
-{
-    return rb_shape_t_to_rb_cShape(rb_shape_get_frozen_root_shape());
-}
-
 VALUE rb_obj_shape(rb_shape_t* shape);
 
 static enum rb_id_table_iterator_result collect_keys_and_values(ID key, VALUE value, void *ref)
@@ -519,6 +593,43 @@ rb_shape_find_by_id(VALUE mod, VALUE id)
 #endif
 
 void
+Init_default_shapes(void)
+{
+    id_frozen = rb_make_internal_id();
+
+    // Shapes by size pool
+    for (int i = 0; i < SIZE_POOL_COUNT; i++) {
+        size_pool_edge_names[i] = rb_make_internal_id();
+    }
+
+    // Root shape
+    rb_shape_t * root = rb_shape_alloc_with_parent_id(0, INVALID_SHAPE_ID);
+    root->capacity = (uint32_t)((rb_size_pool_slot_size(0) - offsetof(struct RObject, as.ary)) / sizeof(VALUE));
+    root->type = SHAPE_ROOT;
+    root->size_pool_index = 0;
+    GET_VM()->root_shape = root;
+    RUBY_ASSERT(rb_shape_id(GET_VM()->root_shape) == ROOT_SHAPE_ID);
+
+    // Shapes by size pool
+    for (int i = 1; i < SIZE_POOL_COUNT; i++) {
+        uint32_t capa = (uint32_t)((rb_size_pool_slot_size(i) - offsetof(struct RObject, as.ary)) / sizeof(VALUE));
+        rb_shape_t * new_shape = rb_shape_transition_shape_capa(root, capa);
+        new_shape->type = SHAPE_INITIAL_CAPACITY;
+        new_shape->size_pool_index = i;
+        RUBY_ASSERT(rb_shape_id(new_shape) == (shape_id_t)i);
+    }
+
+    // Special const shape
+#if RUBY_DEBUG
+    rb_shape_t * special_const_shape =
+#endif
+        get_next_shape_internal(root, (ID)id_frozen, SHAPE_FROZEN);
+    RUBY_ASSERT(rb_shape_id(special_const_shape) == SPECIAL_CONST_SHAPE_ID);
+    RUBY_ASSERT(SPECIAL_CONST_SHAPE_ID == (GET_VM()->next_shape_id - 1));
+    RUBY_ASSERT(rb_shape_frozen_shape_p(special_const_shape));
+}
+
+void
 Init_shape(void)
 {
 #if VM_CHECK_MODE > 0
@@ -530,21 +641,23 @@ Init_shape(void)
     rb_define_method(rb_cShape, "edges", rb_shape_edges, 0);
     rb_define_method(rb_cShape, "edge_name", rb_shape_edge_name, 0);
     rb_define_method(rb_cShape, "next_iv_index", rb_shape_next_iv_index, 0);
+    rb_define_method(rb_cShape, "size_pool_index", rb_shape_size_pool_index, 0);
     rb_define_method(rb_cShape, "depth", rb_shape_export_depth, 0);
     rb_define_method(rb_cShape, "id", rb_wrapped_shape_id, 0);
     rb_define_method(rb_cShape, "type", rb_shape_type, 0);
+    rb_define_method(rb_cShape, "capacity", rb_shape_capacity, 0);
     rb_define_const(rb_cShape, "SHAPE_ROOT", INT2NUM(SHAPE_ROOT));
     rb_define_const(rb_cShape, "SHAPE_IVAR", INT2NUM(SHAPE_IVAR));
     rb_define_const(rb_cShape, "SHAPE_IVAR_UNDEF", INT2NUM(SHAPE_IVAR_UNDEF));
     rb_define_const(rb_cShape, "SHAPE_FROZEN", INT2NUM(SHAPE_FROZEN));
     rb_define_const(rb_cShape, "SHAPE_BITS", INT2NUM(SHAPE_BITS));
     rb_define_const(rb_cShape, "SHAPE_FLAG_SHIFT", INT2NUM(SHAPE_FLAG_SHIFT));
+    rb_define_const(rb_cShape, "SPECIAL_CONST_SHAPE_ID", INT2NUM(SPECIAL_CONST_SHAPE_ID));
 
     rb_define_singleton_method(rb_cShape, "transition_tree", shape_transition_tree, 0);
     rb_define_singleton_method(rb_cShape, "find_by_id", rb_shape_find_by_id, 1);
     rb_define_singleton_method(rb_cShape, "next_shape_id", next_shape_id, 0);
     rb_define_singleton_method(rb_cShape, "of", rb_shape_debug_shape, 1);
     rb_define_singleton_method(rb_cShape, "root_shape", rb_shape_root_shape, 0);
-    rb_define_singleton_method(rb_cShape, "frozen_root_shape", rb_shape_frozen_root_shape, 0);
 #endif
 }