Re-embed when removing Object instance variables

Objects with the same shape must always have the same "embeddedness"
(either embedded or heap allocated) because YJIT assumes so. However,
using remove_instance_variable, it's possible that some objects are
embedded and some are heap allocated because it does not re-embed heap
allocated objects.

This commit changes remove_instance_variable to re-embed Object
instance variables when it becomes small enough.

common.mk

@@ -15849,6 +15849,7 @@ shape.$(OBJEXT): $(top_srcdir)/internal/compilers.h
 shape.$(OBJEXT): $(top_srcdir)/internal/error.h
 shape.$(OBJEXT): $(top_srcdir)/internal/gc.h
 shape.$(OBJEXT): $(top_srcdir)/internal/imemo.h
+shape.$(OBJEXT): $(top_srcdir)/internal/object.h
 shape.$(OBJEXT): $(top_srcdir)/internal/serial.h
 shape.$(OBJEXT): $(top_srcdir)/internal/static_assert.h
 shape.$(OBJEXT): $(top_srcdir)/internal/string.h

gc.c

@@ -2960,12 +2960,6 @@ rb_newobj(void)
     return newobj_of(GET_RACTOR(), 0, T_NONE, 0, 0, 0, FALSE, RVALUE_SIZE);
 }
 
-static size_t
-rb_obj_embedded_size(uint32_t numiv)
-{
-    return offsetof(struct RObject, as.ary) + (sizeof(VALUE) * numiv);
-}
-
 static VALUE
 rb_class_instance_allocate_internal(VALUE klass, VALUE flags, bool wb_protected)
 {

internal/object.h

@@ -11,6 +11,7 @@
 #include "ruby/ruby.h"          /* for VALUE */
 
 /* object.c */
+size_t rb_obj_embedded_size(uint32_t numiv);
 VALUE rb_class_search_ancestor(VALUE klass, VALUE super);
 NORETURN(void rb_undefined_alloc(VALUE klass));
 double rb_num_to_dbl(VALUE val);

object.c

@@ -92,6 +92,12 @@ static VALUE rb_cFalseClass_to_s;
 
 /*! \endcond */
 
+size_t
+rb_obj_embedded_size(uint32_t numiv)
+{
+    return offsetof(struct RObject, as.ary) + (sizeof(VALUE) * numiv);
+}
+
 VALUE
 rb_obj_hide(VALUE obj)
 {

shape.c

@@ -6,6 +6,7 @@
 #include "internal/class.h"
 #include "internal/error.h"
 #include "internal/gc.h"
+#include "internal/object.h"
 #include "internal/symbol.h"
 #include "internal/variable.h"
 #include "variable.h"
@@ -642,6 +643,17 @@ rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape, VALUE
         memmove(&ivptr[removed_shape->next_iv_index - 1], &ivptr[removed_shape->next_iv_index],
                 ((new_shape->next_iv_index + 1) - removed_shape->next_iv_index) * sizeof(VALUE));
 
+        // Re-embed objects when instances become small enough
+        // This is necessary because YJIT assumes that objects with the same shape
+        // have the same embeddedness for efficiency (avoid extra checks)
+        if (BUILTIN_TYPE(obj) == T_OBJECT &&
+                !RB_FL_TEST_RAW(obj, ROBJECT_EMBED) &&
+                rb_obj_embedded_size(new_shape->next_iv_index) <= rb_gc_obj_slot_size(obj)) {
+            RB_FL_SET_RAW(obj, ROBJECT_EMBED);
+            memcpy(ROBJECT_IVPTR(obj), ivptr, new_shape->next_iv_index * sizeof(VALUE));
+            xfree(ivptr);
+        }
+
         rb_shape_set_shape(obj, new_shape);
     }
     return true;

test/ruby/test_object.rb

@@ -355,6 +355,41 @@ class TestObject < Test::Unit::TestCase
     end
   end
 
+  def test_remove_instance_variable_re_embed
+    require "objspace"
+
+    c = Class.new do
+      def a = @a
+
+      def b = @b
+
+      def c = @c
+    end
+
+    o1 = c.new
+    o2 = c.new
+
+    o1.instance_variable_set(:@foo, 5)
+    o1.instance_variable_set(:@a, 0)
+    o1.instance_variable_set(:@b, 1)
+    o1.instance_variable_set(:@c, 2)
+    refute_includes ObjectSpace.dump(o1), '"embedded":true'
+    o1.remove_instance_variable(:@foo)
+    assert_includes ObjectSpace.dump(o1), '"embedded":true'
+
+    o2.instance_variable_set(:@a, 0)
+    o2.instance_variable_set(:@b, 1)
+    o2.instance_variable_set(:@c, 2)
+    assert_includes ObjectSpace.dump(o2), '"embedded":true'
+
+    assert_equal(0, o1.a)
+    assert_equal(1, o1.b)
+    assert_equal(2, o1.c)
+    assert_equal(0, o2.a)
+    assert_equal(1, o2.b)
+    assert_equal(2, o2.c)
+  end
+
   def test_convert_string
     o = Object.new
     def o.to_s; 1; end