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 <[email protected]>
Co-authored-by: John Hawthorn <[email protected]>

10 files changed, 152 insertions, 88 deletions

diff --git a/common.mk b/common.mk
index 0d7f39e4f0..3c89be7f3f 100644
--- a/common.mk
+++ b/common.mk
@@ -14219,6 +14219,7 @@ shape.$(OBJEXT): {$(VPATH)}st.h
 shape.$(OBJEXT): {$(VPATH)}subst.h
 shape.$(OBJEXT): {$(VPATH)}thread_$(THREAD_MODEL).h
 shape.$(OBJEXT): {$(VPATH)}thread_native.h
+shape.$(OBJEXT): {$(VPATH)}variable.h
 shape.$(OBJEXT): {$(VPATH)}vm_core.h
 shape.$(OBJEXT): {$(VPATH)}vm_debug.h
 shape.$(OBJEXT): {$(VPATH)}vm_opts.h
diff --git a/marshal.c b/marshal.c
index e1612303fe..7152be2924 100644
--- a/marshal.c
+++ b/marshal.c
@@ -721,13 +721,23 @@ w_ivar_each(VALUE obj, st_index_t num, struct dump_call_arg *arg)
     struct w_ivar_arg ivarg = {arg, num};
     if (!num) return;
     rb_ivar_foreach(obj, w_obj_each, (st_data_t)&ivarg);
-    if (ivarg.num_ivar) {
-        rb_raise(rb_eRuntimeError, "instance variable removed from %"PRIsVALUE" instance",
-                 CLASS_OF(arg->obj));
-    }
+
     if (shape_id != rb_shape_get_shape_id(arg->obj)) {
-        rb_raise(rb_eRuntimeError, "instance variable added to %"PRIsVALUE" instance",
-                 CLASS_OF(arg->obj));
+        rb_shape_t * expected_shape = rb_shape_get_shape_by_id(shape_id);
+        rb_shape_t * actual_shape = rb_shape_get_shape(arg->obj);
+
+        // If the shape tree got _shorter_ then we probably removed an IV
+        // If the shape tree got longer, then we probably added an IV.
+        // The exception message might not be accurate when someone adds and
+        // removes the same number of IVs, but they will still get an exception
+        if (rb_shape_depth(expected_shape) > rb_shape_depth(actual_shape)) {
+            rb_raise(rb_eRuntimeError, "instance variable removed from %"PRIsVALUE" instance",
+                    CLASS_OF(arg->obj));
+        }
+        else {
+            rb_raise(rb_eRuntimeError, "instance variable added to %"PRIsVALUE" instance",
+                    CLASS_OF(arg->obj));
+        }
     }
 }
 
diff --git a/mjit_c.rb b/mjit_c.rb
index 39ca89a0ea..9b85faec30 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -194,10 +194,6 @@ module RubyVM::MJIT
     Primitive.cexpr! %q{ UINT2NUM(SHAPE_IVAR) }
   end
 
-  def C.SHAPE_IVAR_UNDEF
-    Primitive.cexpr! %q{ UINT2NUM(SHAPE_IVAR_UNDEF) }
-  end
-
   def C.SHAPE_ROOT
     Primitive.cexpr! %q{ UINT2NUM(SHAPE_ROOT) }
   end
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));
diff --git a/shape.h b/shape.h
index 8ed51c6211..96feae99fd 100644
--- a/shape.h
+++ b/shape.h
@@ -51,7 +51,6 @@ enum shape_type {
     SHAPE_IVAR,
     SHAPE_FROZEN,
     SHAPE_CAPACITY_CHANGE,
-    SHAPE_IVAR_UNDEF,
     SHAPE_INITIAL_CAPACITY,
     SHAPE_T_OBJECT,
 };
@@ -125,6 +124,7 @@ bool rb_shape_root_shape_p(rb_shape_t* shape);
 rb_shape_t * rb_shape_get_root_shape(void);
 uint8_t rb_shape_id_num_bits(void);
 int32_t rb_shape_id_offset(void);
+unsigned int rb_shape_depth(rb_shape_t * shape);
 
 rb_shape_t* rb_shape_get_shape_by_id_without_assertion(shape_id_t shape_id);
 rb_shape_t * rb_shape_get_parent(rb_shape_t * shape);
@@ -136,7 +136,7 @@ shape_id_t rb_shape_get_shape_id(VALUE obj);
 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);
+void rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape, VALUE * removed);
 rb_shape_t * rb_shape_transition_shape_capa(rb_shape_t * shape, uint32_t new_capacity);
 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);
diff --git a/spec/bundler/bundler/fetcher/compact_index_spec.rb b/spec/bundler/bundler/fetcher/compact_index_spec.rb
index 00eb27edea..842088a20a 100644
--- a/spec/bundler/bundler/fetcher/compact_index_spec.rb
+++ b/spec/bundler/bundler/fetcher/compact_index_spec.rb
@@ -52,7 +52,9 @@ RSpec.describe Bundler::Fetcher::CompactIndex do
 
       context "when OpenSSL is FIPS-enabled" do
         def remove_cached_md5_availability
+          require "objspace"
           return unless Bundler::SharedHelpers.instance_variable_defined?(:@md5_available)
+          puts ObjectSpace.dump(Bundler::SharedHelpers)
           Bundler::SharedHelpers.remove_instance_variable(:@md5_available)
         end
 
diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
index 848bb4971a..b19238950b 100644
--- a/test/ruby/test_shapes.rb
+++ b/test/ruby/test_shapes.rb
@@ -28,7 +28,7 @@ class TestShapes < Test::Unit::TestCase
       @foo = 1
     end
 
-    def remove
+    def remove_foo
       remove_instance_variable(:@foo)
     end
 
@@ -64,26 +64,36 @@ class TestShapes < Test::Unit::TestCase
 
   def test_iv_index
     example = RemoveAndAdd.new
-    shape = RubyVM::Shape.of(example)
-    assert_equal 0, shape.next_iv_index
+    initial_shape = RubyVM::Shape.of(example)
+    assert_equal 0, initial_shape.next_iv_index
 
     example.add_foo # makes a transition
-    new_shape = RubyVM::Shape.of(example)
+    add_foo_shape = RubyVM::Shape.of(example)
     assert_equal([:@foo], example.instance_variables)
-    assert_equal(shape.id, new_shape.parent.id)
-    assert_equal(1, new_shape.next_iv_index)
+    assert_equal(initial_shape.id, add_foo_shape.parent.id)
+    assert_equal(1, add_foo_shape.next_iv_index)
 
-    example.remove # makes a transition
-    remove_shape = RubyVM::Shape.of(example)
+    example.remove_foo # makes a transition
+    remove_foo_shape = RubyVM::Shape.of(example)
     assert_equal([], example.instance_variables)
-    assert_equal(new_shape.id, remove_shape.parent.id)
-    assert_equal(1, remove_shape.next_iv_index)
+    assert_shape_equal(initial_shape, remove_foo_shape)
 
     example.add_bar # makes a transition
     bar_shape = RubyVM::Shape.of(example)
     assert_equal([:@bar], example.instance_variables)
-    assert_equal(remove_shape.id, bar_shape.parent.id)
-    assert_equal(2, bar_shape.next_iv_index)
+    assert_equal(initial_shape.id, bar_shape.parent_id)
+    assert_equal(1, bar_shape.next_iv_index)
+  end
+
+  def test_remove_then_add_again
+    example = RemoveAndAdd.new
+    initial_shape = RubyVM::Shape.of(example)
+
+    example.add_foo # makes a transition
+    add_foo_shape = RubyVM::Shape.of(example)
+    example.remove_foo # makes a transition
+    example.add_foo # makes a transition
+    assert_shape_equal(add_foo_shape, RubyVM::Shape.of(example))
   end
 
   class TestObject; end
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index b846a91664..c093cd423c 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -359,7 +359,6 @@ generator = BindingGenerator.new(
       SHAPE_FROZEN
       SHAPE_INITIAL_CAPACITY
       SHAPE_IVAR
-      SHAPE_IVAR_UNDEF
       SHAPE_ROOT
     ],
     ULONG: %w[
diff --git a/variable.c b/variable.c
index 12b0793756..42a73b5953 100644
--- a/variable.c
+++ b/variable.c
@@ -1219,7 +1219,7 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
     rb_check_frozen(obj);
 
     VALUE val = undef;
-    attr_index_t index;
+    rb_shape_t * shape = rb_shape_get_shape(obj);
 
     switch (BUILTIN_TYPE(obj)) {
       case T_CLASS:
@@ -1228,36 +1228,18 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
 
         RB_VM_LOCK_ENTER();
         {
-            rb_shape_t * shape = rb_shape_get_shape(obj);
-            if (rb_shape_get_iv_index(shape, id, &index)) {
-                rb_shape_transition_shape_remove_ivar(obj, id, shape);
-                val = RCLASS_IVPTR(obj)[index];
-                RCLASS_IVPTR(obj)[index] = Qundef;
-            }
+            rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
         }
         RB_VM_LOCK_LEAVE();
 
         break;
       case T_OBJECT: {
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-        if (rb_shape_get_iv_index(shape, id, &index)) {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape);
-            val = ROBJECT_IVPTR(obj)[index];
-            ROBJECT_IVPTR(obj)[index] = Qundef;
-        }
+        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
 
         break;
       }
       default: {
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-
-        if (rb_shape_get_iv_index(shape, id, &index)) {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape);
-            struct gen_ivtbl *ivtbl;
-            rb_gen_ivtbl_get(obj, id, &ivtbl);
-            val = ivtbl->ivptr[index];
-            ivtbl->ivptr[index] = Qundef;
-        }
+        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
 
         break;
       }
@@ -1613,7 +1595,6 @@ iterate_over_shapes_with_callback(rb_shape_t *shape, rb_ivar_foreach_callback_fu
       case SHAPE_INITIAL_CAPACITY:
       case SHAPE_CAPACITY_CHANGE:
       case SHAPE_FROZEN:
-      case SHAPE_IVAR_UNDEF:
       case SHAPE_T_OBJECT:
         iterate_over_shapes_with_callback(rb_shape_get_parent(shape), callback, itr_data);
         return;
@@ -1892,58 +1873,39 @@ check_id_type(VALUE obj, VALUE *pname,
 VALUE
 rb_obj_remove_instance_variable(VALUE obj, VALUE name)
 {
-    VALUE val = Qnil;
+    VALUE val = Qundef;
     const ID id = id_for_var(obj, name, an, instance);
 
     // Frozen check comes here because it's expected that we raise a
     // NameError (from the id_for_var check) before we raise a FrozenError
     rb_check_frozen(obj);
 
-    attr_index_t index;
-
     if (!id) {
         goto not_defined;
     }
 
+    rb_shape_t * shape = rb_shape_get_shape(obj);
+
     switch (BUILTIN_TYPE(obj)) {
       case T_CLASS:
       case T_MODULE:
         IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(id);
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-        if (rb_shape_get_iv_index(shape, id, &index)) {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape);
-            val = RCLASS_IVPTR(obj)[index];
-            RCLASS_IVPTR(obj)[index] = Qundef;
-            return val;
-        }
+        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
         break;
       case T_OBJECT: {
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-        if (rb_shape_get_iv_index(shape, id, &index)) {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape);
-            val = ROBJECT_IVPTR(obj)[index];
-            ROBJECT_IVPTR(obj)[index] = Qundef;
-            return val;
-        }
-
+        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
         break;
       }
       default: {
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-
-        if (rb_shape_get_iv_index(shape, id, &index)) {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape);
-            struct gen_ivtbl *ivtbl;
-            rb_gen_ivtbl_get(obj, id, &ivtbl);
-            val = ivtbl->ivptr[index];
-            ivtbl->ivptr[index] = Qundef;
-            return val;
-        }
-
+        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
         break;
       }
     }
 
+    if (val != Qundef) {
+        return val;
+    }
+
   not_defined:
     rb_name_err_raise("instance variable %1$s not defined",
                       obj, name);
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index 86d3c991f4..ec90e67373 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -1416,7 +1416,7 @@ vm_setivar(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_index_t i
                 rb_shape_t *dest_shape = rb_shape_get_shape_by_id(dest_shape_id);
                 shape_id_t source_shape_id = dest_shape->parent_id;
 
-                RUBY_ASSERT(dest_shape->type == SHAPE_IVAR || dest_shape->type == SHAPE_IVAR_UNDEF);
+                RUBY_ASSERT(dest_shape->type == SHAPE_IVAR);
 
                 if (shape_id == source_shape_id && dest_shape->edge_name == id) {
                     RUBY_ASSERT(dest_shape_id != INVALID_SHAPE_ID && shape_id != INVALID_SHAPE_ID);