Let memory sizes of the various IMEMO object types be reflected correctly

[Feature #15805]

Closes: https://github.com/ruby/ruby/pull/2140

4 files changed, 61 insertions, 8 deletions

diff --git a/gc.c b/gc.c
index d61db0448a..a978887ae9 100644
--- a/gc.c
+++ b/gc.c
@@ -842,6 +842,7 @@ int ruby_disable_gc = 0;
 void rb_iseq_mark(const rb_iseq_t *iseq);
 void rb_iseq_update_references(rb_iseq_t *iseq);
 void rb_iseq_free(const rb_iseq_t *iseq);
+size_t rb_iseq_memsize(const rb_iseq_t *iseq);
 void rb_vm_update_references(void *ptr);
 
 void rb_gcdebug_print_obj_condition(VALUE obj);
@@ -2159,6 +2160,40 @@ rb_imemo_tmpbuf_parser_heap(void *buf, rb_imemo_tmpbuf_t *old_heap, size_t cnt)
     return (rb_imemo_tmpbuf_t *)rb_imemo_tmpbuf_new((VALUE)buf, (VALUE)old_heap, (VALUE)cnt, 0);
 }
 
+static size_t
+imemo_memsize(VALUE obj)
+{
+    size_t size = 0;
+    switch (imemo_type(obj)) {
+      case imemo_ment:
+        size += sizeof(RANY(obj)->as.imemo.ment.def);
+        break;
+      case imemo_iseq:
+        size += rb_iseq_memsize((rb_iseq_t *)obj);
+        break;
+      case imemo_env:
+        size += RANY(obj)->as.imemo.env.env_size * sizeof(VALUE);
+        break;
+      case imemo_tmpbuf:
+        size += RANY(obj)->as.imemo.alloc.cnt * sizeof(VALUE);
+        break;
+      case imemo_ast:
+        size += rb_ast_memsize(&RANY(obj)->as.imemo.ast);
+        break;
+      case imemo_cref:
+      case imemo_svar:
+      case imemo_throw_data:
+      case imemo_ifunc:
+      case imemo_memo:
+      case imemo_parser_strterm:
+        break;
+      default:
+        /* unreachable */
+        break;
+    }
+    return size;
+}
+
 #if IMEMO_DEBUG
 VALUE
 rb_imemo_new_debug(enum imemo_type type, VALUE v1, VALUE v2, VALUE v3, VALUE v0, const char *file, int line)
@@ -3628,9 +3663,7 @@ obj_memsize_of(VALUE obj, int use_all_types)
       case T_COMPLEX:
         break;
       case T_IMEMO:
-	if (imemo_type_p(obj, imemo_tmpbuf)) {
-	    size += RANY(obj)->as.imemo.alloc.cnt * sizeof(VALUE);
-	}
+        size += imemo_memsize(obj);
 	break;
 
       case T_FLOAT:
diff --git a/iseq.c b/iseq.c
index fa4e72d2cd..9ee8d41275 100644
--- a/iseq.c
+++ b/iseq.c
@@ -361,8 +361,8 @@ param_keyword_size(const struct rb_iseq_param_keyword *pkw)
     return size;
 }
 
-static size_t
-iseq_memsize(const rb_iseq_t *iseq)
+size_t
+rb_iseq_memsize(const rb_iseq_t *iseq)
 {
     size_t size = 0; /* struct already counted as RVALUE size */
     const struct rb_iseq_constant_body *body = iseq->body;
@@ -1109,7 +1109,7 @@ iseqw_mark(void *ptr)
 static size_t
 iseqw_memsize(const void *ptr)
 {
-    return iseq_memsize((const rb_iseq_t *)ptr);
+    return rb_iseq_memsize((const rb_iseq_t *)ptr);
 }
 
 static const rb_data_type_t iseqw_data_type = {
diff --git a/node.c b/node.c
index e4c4961217..99558319f8 100644
--- a/node.c
+++ b/node.c
@@ -12,6 +12,8 @@
 #include "ruby/ruby.h"
 #include "vm_core.h"
 
+#define NODE_BUF_DEFAULT_LEN 16
+
 #define A(str) rb_str_cat2(buf, (str))
 #define AR(str) rb_str_concat(buf, (str))
 
@@ -1122,9 +1124,9 @@ struct node_buffer_struct {
 static node_buffer_t *
 rb_node_buffer_new(void)
 {
-    node_buffer_t *nb = xmalloc(sizeof(node_buffer_t) + offsetof(node_buffer_elem_t, buf) + 16 * sizeof(NODE));
+    node_buffer_t *nb = xmalloc(sizeof(node_buffer_t) + offsetof(node_buffer_elem_t, buf) + NODE_BUF_DEFAULT_LEN * sizeof(NODE));
     nb->idx = 0;
-    nb->len = 16;
+    nb->len = NODE_BUF_DEFAULT_LEN;
     nb->head = nb->last = (node_buffer_elem_t*) &nb[1];
     nb->head->next = NULL;
     nb->mark_ary = rb_ary_tmp_new(0);
@@ -1193,6 +1195,23 @@ rb_ast_free(rb_ast_t *ast)
     }
 }
 
+size_t
+rb_ast_memsize(const rb_ast_t *ast)
+{
+    size_t size = 0;
+    node_buffer_t *nb = ast->node_buffer;
+
+    if (nb) {
+        size += sizeof(node_buffer_t) + offsetof(node_buffer_elem_t, buf) + NODE_BUF_DEFAULT_LEN * sizeof(NODE);
+        node_buffer_elem_t *nbe = nb->head;
+        while (nbe != nb->last) {
+            nbe = nbe->next;
+            size += offsetof(node_buffer_elem_t, buf) + nb->len * sizeof(NODE);
+        }
+    }
+    return size;
+}
+
 void
 rb_ast_dispose(rb_ast_t *ast)
 {
diff --git a/node.h b/node.h
index a1f01ed088..3741e649ca 100644
--- a/node.h
+++ b/node.h
@@ -406,6 +406,7 @@ rb_ast_t *rb_ast_new(void);
 void rb_ast_mark(rb_ast_t*);
 void rb_ast_dispose(rb_ast_t*);
 void rb_ast_free(rb_ast_t*);
+size_t rb_ast_memsize(const rb_ast_t*);
 void rb_ast_add_mark_object(rb_ast_t*, VALUE);
 NODE *rb_ast_newnode(rb_ast_t*);
 void rb_ast_delete_node(rb_ast_t*, NODE *n);