Stop transitioning to UNDEF when undefining an instance variable

Cases like this:

```ruby
obj = Object.new
loop do
  obj.instance_variable_set(:@foo, 1)
  obj.remove_instance_variable(:@foo)
end
```

can cause us to use many more shapes than we want (and even run out).
This commit changes the code such that when an instance variable is
removed, we'll walk up the shape tree, find the shape, then rebuild any
child nodes that happened to be below the "targetted for removal" IV.

This also requires moving any instance variables so that indexes derived
from the shape tree will work correctly.

Co-Authored-By: Jemma Issroff <jemmaissroff@gmail.com>
Co-authored-by: John Hawthorn <jhawthorn@github.com>

1 files changed, 94 insertions, 10 deletions

diff --git a/shape.c b/shape.c
index 097458b5eb..973a8a6328 100644
--- a/shape.c
+++ b/shape.c
@@ -5,6 +5,7 @@
 #include "internal/class.h"
 #include "internal/symbol.h"
 #include "internal/variable.h"
+#include "variable.h"
 #include <stdbool.h>
 
 #ifndef SHAPE_DEBUG
@@ -96,6 +97,19 @@ rb_shape_get_shape_id(VALUE obj)
 #endif
 }
 
+unsigned int
+rb_shape_depth(rb_shape_t * shape)
+{
+    unsigned int depth = 1;
+
+    while (shape->parent_id != INVALID_SHAPE_ID) {
+        depth++;
+        shape = rb_shape_get_parent(shape);
+    }
+
+    return depth;
+}
+
 rb_shape_t*
 rb_shape_get_shape(VALUE obj)
 {
@@ -131,7 +145,6 @@ get_next_shape_internal(rb_shape_t * shape, ID id, enum shape_type shape_type)
                 new_shape->next_iv_index = shape->next_iv_index + 1;
                 break;
               case SHAPE_CAPACITY_CHANGE:
-              case SHAPE_IVAR_UNDEF:
               case SHAPE_FROZEN:
               case SHAPE_T_OBJECT:
                 new_shape->next_iv_index = shape->next_iv_index;
@@ -157,12 +170,88 @@ rb_shape_frozen_shape_p(rb_shape_t* shape)
     return SHAPE_FROZEN == (enum shape_type)shape->type;
 }
 
-void
-rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape)
+static void
+move_iv(VALUE obj, ID id, attr_index_t from, attr_index_t to)
 {
-    rb_shape_t * next_shape = get_next_shape_internal(shape, id, SHAPE_IVAR_UNDEF);
+    switch(BUILTIN_TYPE(obj)) {
+      case T_CLASS:
+      case T_MODULE:
+        RCLASS_IVPTR(obj)[to] = RCLASS_IVPTR(obj)[from];
+        break;
+      case T_OBJECT:
+        ROBJECT_IVPTR(obj)[to] = ROBJECT_IVPTR(obj)[from];
+        break;
+      default: {
+        struct gen_ivtbl *ivtbl;
+        rb_gen_ivtbl_get(obj, id, &ivtbl);
+        ivtbl->ivptr[to] = ivtbl->ivptr[from];
+        break;
+      }
+    }
+}
 
-    rb_shape_set_shape(obj, next_shape);
+static rb_shape_t *
+remove_shape_recursive(VALUE obj, ID id, rb_shape_t * shape, VALUE * removed)
+{
+    if (shape->parent_id == INVALID_SHAPE_ID) {
+        // We've hit the top of the shape tree and couldn't find the
+        // IV we wanted to remove, so return NULL
+        return NULL;
+    }
+    else {
+        if (shape->type == SHAPE_IVAR && shape->edge_name == id) {
+            // We've hit the edge we wanted to remove, return it's _parent_
+            // as the new parent while we go back down the stack.
+            attr_index_t index = shape->next_iv_index - 1;
+
+            switch(BUILTIN_TYPE(obj)) {
+              case T_CLASS:
+              case T_MODULE:
+                *removed = RCLASS_IVPTR(obj)[index];
+                break;
+              case T_OBJECT:
+                *removed = ROBJECT_IVPTR(obj)[index];
+                break;
+              default: {
+                struct gen_ivtbl *ivtbl;
+                rb_gen_ivtbl_get(obj, id, &ivtbl);
+                *removed = ivtbl->ivptr[index];
+                break;
+              }
+            }
+            return rb_shape_get_parent(shape);
+        }
+        else {
+            // This isn't the IV we want to remove, keep walking up.
+            rb_shape_t * new_parent = remove_shape_recursive(obj, id, rb_shape_get_parent(shape), removed);
+
+            // We found a new parent.  Create a child of the new parent that
+            // has the same attributes as this shape.
+            if (new_parent) {
+                rb_shape_t * new_child = get_next_shape_internal(new_parent, shape->edge_name, shape->type);
+                if (new_child->type == SHAPE_IVAR) {
+                    move_iv(obj, id, shape->next_iv_index - 1, new_child->next_iv_index - 1);
+
+                }
+
+                return new_child;
+            }
+            else {
+                // We went all the way to the top of the shape tree and couldn't
+                // find an IV to remove, so return NULL
+                return NULL;
+            }
+        }
+    }
+}
+
+void
+rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape, VALUE * removed)
+{
+    rb_shape_t * new_shape = remove_shape_recursive(obj, id, shape, removed);
+    if (new_shape) {
+        rb_shape_set_shape(obj, new_shape);
+    }
 }
 
 void
@@ -238,7 +327,6 @@ rb_shape_get_iv_index(rb_shape_t * shape, ID id, attr_index_t *value)
                 *value = shape->next_iv_index - 1;
                 return true;
               case SHAPE_CAPACITY_CHANGE:
-              case SHAPE_IVAR_UNDEF:
               case SHAPE_ROOT:
               case SHAPE_INITIAL_CAPACITY:
               case SHAPE_T_OBJECT:
@@ -329,9 +417,6 @@ rb_shape_rebuild_shape(rb_shape_t * initial_shape, rb_shape_t * dest_shape)
 
         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:
@@ -638,7 +723,6 @@ Init_shape(void)
     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_T_OBJECT", INT2NUM(SHAPE_T_OBJECT));
-    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_ID_NUM_BITS", INT2NUM(SHAPE_ID_NUM_BITS));
     rb_define_const(rb_cShape, "SHAPE_FLAG_SHIFT", INT2NUM(SHAPE_FLAG_SHIFT));