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, 26 insertions, 2 deletions

diff --git a/shape.h b/shape.h
index 8e1bf46ec9..a7450cdeea 100644
--- a/shape.h
+++ b/shape.h
@@ -40,13 +40,17 @@ typedef uint16_t shape_id_t;
 # define MAX_SHAPE_ID (SHAPE_MASK - 1)
 # define INVALID_SHAPE_ID SHAPE_MASK
 # define ROOT_SHAPE_ID 0x0
-# define FROZEN_ROOT_SHAPE_ID 0x1
+// We use SIZE_POOL_COUNT number of shape IDs for transitions out of different size pools
+// The next available shapd ID will be the SPECIAL_CONST_SHAPE_ID
+# define SPECIAL_CONST_SHAPE_ID SIZE_POOL_COUNT
 
 struct rb_shape {
     struct rb_id_table * edges; // id_table from ID (ivar) to next shape
     ID edge_name; // ID (ivar) for transition from parent to rb_shape
     attr_index_t next_iv_index;
+    uint32_t capacity; // Total capacity of the object with this shape
     uint8_t type;
+    uint8_t size_pool_index;
     shape_id_t parent_id;
 };
 
@@ -56,7 +60,9 @@ enum shape_type {
     SHAPE_ROOT,
     SHAPE_IVAR,
     SHAPE_FROZEN,
+    SHAPE_CAPACITY_CHANGE,
     SHAPE_IVAR_UNDEF,
+    SHAPE_INITIAL_CAPACITY,
 };
 
 #if SHAPE_IN_BASIC_FLAGS
@@ -124,6 +130,7 @@ static inline shape_id_t RCLASS_SHAPE_ID(VALUE obj) {
 #endif
 
 bool rb_shape_root_shape_p(rb_shape_t* shape);
+rb_shape_t * rb_shape_get_root_shape(void);
 
 rb_shape_t* rb_shape_get_shape_by_id_without_assertion(shape_id_t shape_id);
 
@@ -135,22 +142,38 @@ rb_shape_t* rb_shape_get_shape(VALUE obj);
 int rb_shape_frozen_shape_p(rb_shape_t* shape);
 void rb_shape_transition_shape_frozen(VALUE obj);
 void rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape);
+rb_shape_t * rb_shape_transition_shape_capa(rb_shape_t * shape, uint32_t new_capacity);
 void rb_shape_transition_shape(VALUE obj, ID id, rb_shape_t *shape);
+rb_shape_t * rb_shape_get_next_iv_shape(rb_shape_t * shape, ID id);
 rb_shape_t* rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id);
 bool rb_shape_get_iv_index(rb_shape_t * shape, ID id, attr_index_t * value);
 shape_id_t rb_shape_id(rb_shape_t * shape);
 MJIT_SYMBOL_EXPORT_END
 
+rb_shape_t * rb_shape_rebuild_shape(rb_shape_t * initial_shape, rb_shape_t * dest_shape);
+
+static inline uint32_t
+ROBJECT_IV_CAPACITY(VALUE obj)
+{
+    RBIMPL_ASSERT_TYPE(obj, RUBY_T_OBJECT);
+    return rb_shape_get_shape_by_id(ROBJECT_SHAPE_ID(obj))->capacity;
+}
+
 static inline uint32_t
 ROBJECT_IV_COUNT(VALUE obj)
 {
     RBIMPL_ASSERT_TYPE(obj, RUBY_T_OBJECT);
     uint32_t ivc = rb_shape_get_shape_by_id(ROBJECT_SHAPE_ID(obj))->next_iv_index;
-    RUBY_ASSERT(ivc <= ROBJECT_NUMIV(obj));
     return ivc;
 }
 
 static inline uint32_t
+RBASIC_IV_COUNT(VALUE obj)
+{
+    return rb_shape_get_shape_by_id(rb_shape_get_shape_id(obj))->next_iv_index;
+}
+
+static inline uint32_t
 RCLASS_IV_COUNT(VALUE obj)
 {
     RUBY_ASSERT(RB_TYPE_P(obj, RUBY_T_CLASS) || RB_TYPE_P(obj, RUBY_T_MODULE));
@@ -159,6 +182,7 @@ RCLASS_IV_COUNT(VALUE obj)
 }
 
 rb_shape_t * rb_shape_alloc(ID edge_name, rb_shape_t * parent);
+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 * rb_shape_alloc_with_parent_id(ID edge_name, shape_id_t parent_id);
 
 bool rb_shape_set_shape_id(VALUE obj, shape_id_t shape_id);