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);