From 12e3b07455fea0e99e6aaf1893a7883fb2b0197e Mon Sep 17 00:00:00 2001
From: Peter Zhu <peter@peterzhu.ca>
Date: Wed, 6 Dec 2023 07:37:57 -0500
Subject: [PATCH] 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                |  1 +
 gc.c                     |  6 ------
 internal/object.h        |  1 +
 object.c                 |  6 ++++++
 shape.c                  | 12 ++++++++++++
 test/ruby/test_object.rb | 35 +++++++++++++++++++++++++++++++++++
 6 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/common.mk b/common.mk
index 0e18bfe35bf581..928eb0a02bf214 100644
--- a/common.mk
+++ b/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
diff --git a/gc.c b/gc.c
index d94bd970e55806..993d81c2d6e656 100644
--- a/gc.c
+++ b/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)
 {
diff --git a/internal/object.h b/internal/object.h
index 58e989562a182e..06595bdd91a329 100644
--- a/internal/object.h
+++ b/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);
diff --git a/object.c b/object.c
index 2bcfbc4c3c70bd..81d1ed6f6fac5c 100644
--- a/object.c
+++ b/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)
 {
diff --git a/shape.c b/shape.c
index b19a6e2cc5fa07..2d5c2109c7d333 100644
--- a/shape.c
+++ b/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;
diff --git a/test/ruby/test_object.rb b/test/ruby/test_object.rb
index 3c5b6424ba125c..acc913b9c0ad2d 100644
--- a/test/ruby/test_object.rb
+++ b/test/ruby/test_object.rb
@@ -355,6 +355,41 @@ def test_remove_instance_variable
     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