Set max_iv_count (used for object shapes) based on inline caches

With this change, we're storing the iv name on an inline cache on
setinstancevariable instructions. This allows us to check the inline
cache to count instance variables set in initialize and give us an
estimate of iv capacity for an object.

For the purpose of estimating the number of instance variables required
for an object, we're assuming that all initialize methods will call
`super`.

This change allows us to estimate the number of instance variables
required without disassembling instruction sequences.

Co-Authored-By: Aaron Patterson <tenderlove@ruby-lang.org>

5 files changed, 31 insertions, 21 deletions

diff --git a/compile.c b/compile.c
index fb99826413..8f53955e74 100644
--- a/compile.c
+++ b/compile.c
@@ -2461,7 +2461,17 @@ iseq_set_sequence(rb_iseq_t *iseq, LINK_ANCHOR *const anchor)
                       case TS_IVC: /* inline ivar cache */
                         {
                             unsigned int ic_index = FIX2UINT(operands[j]);
-                            vm_ic_attr_index_initialize(((IVC)&body->is_entries[ic_index]), INVALID_SHAPE_ID);
+
+                            IVC cache = ((IVC)&body->is_entries[ic_index]);
+
+                            if (insn == BIN(setinstancevariable)) {
+                                cache->iv_set_name = SYM2ID(operands[j - 1]);
+                            }
+                            else {
+                                cache->iv_set_name = 0;
+                            }
+
+                            vm_ic_attr_index_initialize(cache, INVALID_SHAPE_ID);
                         }
                       case TS_ISE: /* inline storage entry: `once` insn */
                       case TS_ICVARC: /* inline cvar cache */
@@ -11529,7 +11539,17 @@ ibf_load_code(const struct ibf_load *load, rb_iseq_t *iseq, ibf_offset_t bytecod
                     code[code_index] = (VALUE)ic;
 
                     if (operand_type == TS_IVC) {
-                        vm_ic_attr_index_initialize(((IVC)code[code_index]), INVALID_SHAPE_ID);
+                        IVC cache = (IVC)ic;
+
+                        if (insn == BIN(setinstancevariable)) {
+                            ID iv_name = (ID)code[code_index - 1];
+                            cache->iv_set_name = iv_name;
+                        }
+                        else {
+                            cache->iv_set_name = 0;
+                        }
+
+                        vm_ic_attr_index_initialize(cache, INVALID_SHAPE_ID);
                     }
 
                 }
diff --git a/iseq.c b/iseq.c
index 381bf6ed20..869ed4da6a 100644
--- a/iseq.c
+++ b/iseq.c
@@ -2509,33 +2509,20 @@ rb_iseq_disasm(const rb_iseq_t *iseq)
 attr_index_t
 rb_estimate_iv_count(VALUE klass, const rb_iseq_t * initialize_iseq)
 {
-    bool calls_super = false;
-
     struct rb_id_table * iv_names = rb_id_table_create(0);
 
-    VALUE * code = ISEQ_BODY(initialize_iseq)->iseq_encoded;
-
-    for (unsigned int i = 0; i < ISEQ_BODY(initialize_iseq)->iseq_size; ) {
-        VALUE insn = code[i];
-        int original_insn = rb_vm_insn_addr2insn((const void *)insn);
+    for (unsigned int i = 0; i < ISEQ_BODY(initialize_iseq)->ivc_size; i++) {
+        IVC cache = (IVC)&ISEQ_BODY(initialize_iseq)->is_entries[i];
 
-        if (BIN(setinstancevariable) == original_insn) {
-            ID name = (ID)code[i + 1];
-            rb_id_table_insert(iv_names, name, Qtrue);
+        if (cache->iv_set_name) {
+            rb_id_table_insert(iv_names, cache->iv_set_name, Qtrue);
         }
-        else if (BIN(invokesuper) == original_insn) {
-            calls_super = true;
-        }
-
-        i += insn_len(original_insn);
     }
 
     attr_index_t count = (attr_index_t)rb_id_table_size(iv_names);
 
-    if (calls_super) {
-        VALUE superclass = rb_class_superclass(klass);
-        count += RCLASS_EXT(superclass)->max_iv_count;
-    }
+    VALUE superclass = rb_class_superclass(klass);
+    count += RCLASS_EXT(superclass)->max_iv_count;
 
     rb_id_table_free(iv_names);
 
diff --git a/mjit_c.rb b/mjit_c.rb
index e4933c3458..39ca89a0ea 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -310,6 +310,7 @@ module RubyVM::MJIT
     @iseq_inline_iv_cache_entry ||= CType::Struct.new(
       "iseq_inline_iv_cache_entry", Primitive.cexpr!("SIZEOF(struct iseq_inline_iv_cache_entry)"),
       value: [CType::Immediate.parse("uintptr_t"), Primitive.cexpr!("OFFSETOF((*((struct iseq_inline_iv_cache_entry *)NULL)), value)")],
+      iv_set_name: [self.ID, Primitive.cexpr!("OFFSETOF((*((struct iseq_inline_iv_cache_entry *)NULL)), iv_set_name)")],
     )
   end
 
diff --git a/vm_core.h b/vm_core.h
index a996c5b835..65f2741219 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -275,6 +275,7 @@ struct iseq_inline_constant_cache {
 
 struct iseq_inline_iv_cache_entry {
     uintptr_t value; // attr_index in lower bits, dest_shape_id in upper bits
+    ID iv_set_name;
 };
 
 struct iseq_inline_cvar_cache_entry {
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 470b6d0a5b..1555196a06 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -853,6 +853,7 @@ pub struct iseq_inline_constant_cache {
 #[derive(Debug, Copy, Clone)]
 pub struct iseq_inline_iv_cache_entry {
     pub value: usize,
+    pub iv_set_name: ID,
 }
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]