4 files changed, 46 insertions, 1 deletions
diff --git a/gc.c b/gc.c
index e7a1b6fb4f..dc35f77f34 100644
--- a/gc.c
+++ b/gc.c
@@ -3089,7 +3089,7 @@ rb_imemo_new_debug(enum imemo_type type, VALUE v1, VALUE v2, VALUE v3, VALUE v0,
 }
 #endif
 
-VALUE
+MJIT_FUNC_EXPORTED VALUE
 rb_class_allocate_instance(VALUE klass)
 {
     return rb_class_instance_allocate_internal(klass, T_OBJECT | ROBJECT_EMBED, RGENGC_WB_PROTECTED_OBJECT);
diff --git a/iseq.c b/iseq.c
index a2e79cd55b..b366f03823 100644
--- a/iseq.c
+++ b/iseq.c
@@ -2491,6 +2491,45 @@ rb_iseq_disasm(const rb_iseq_t *iseq)
 }
 
 /*
+ * Estimates the number of instance variables that will be set on
+ * a given `class` with the initialize method defined in
+ * `initialize_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);
+
+        if (BIN(setinstancevariable) == original_insn) {
+            ID name = (ID)code[i + 1];
+            rb_id_table_insert(iv_names, 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;
+    }
+
+    return count;
+}
+
+/*
  *  call-seq:
  *     iseq.disasm -> str
  *     iseq.disassemble -> str
diff --git a/vm_core.h b/vm_core.h
index 3364089f79..9c016de8fd 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -1158,6 +1158,7 @@ rb_iseq_t *rb_iseq_new_with_callback(const struct rb_iseq_new_with_callback_call
 
 VALUE rb_iseq_disasm(const rb_iseq_t *iseq);
 int rb_iseq_disasm_insn(VALUE str, const VALUE *iseqval, size_t pos, const rb_iseq_t *iseq, VALUE child);
+attr_index_t rb_estimate_iv_count(VALUE klass, const rb_iseq_t * initialize_iseq);
 
 VALUE rb_iseq_coverage(const rb_iseq_t *iseq);
 
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index 76bdbb86de..073c515d3c 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -4946,6 +4946,11 @@ vm_define_method(const rb_execution_context_t *ec, VALUE obj, ID id, VALUE iseqv
     }
 
     rb_add_method_iseq(klass, id, (const rb_iseq_t *)iseqval, cref, visi);
+    // Set max_iv_count on klasses based on number of ivar sets that are in the initialize method
+    if (id == rb_intern("initialize") && klass != rb_cObject &&  RB_TYPE_P(klass, T_CLASS) && (rb_get_alloc_func(klass) == rb_class_allocate_instance)) {
+
+        RCLASS_EXT(klass)->max_iv_count = rb_estimate_iv_count(klass, (const rb_iseq_t *)iseqval);
+    }
 
     if (!is_singleton && vm_scope_module_func_check(ec)) {
         klass = rb_singleton_class(klass);