Merge pull request #5032 from RoryO/add-objspace-memsize-of

Add ObjectSpace.memsize_of

3 files changed, 394 insertions, 0 deletions

diff --git a/mrbgems/mruby-objectspace/mrbgem.rake b/mrbgems/mruby-objectspace/mrbgem.rake
index fa35136a1..101e24275 100644
--- a/mrbgems/mruby-objectspace/mrbgem.rake
+++ b/mrbgems/mruby-objectspace/mrbgem.rake
@@ -2,4 +2,8 @@ MRuby::Gem::Specification.new('mruby-objectspace') do |spec|
   spec.license = 'MIT'
   spec.author  = 'mruby developers'
   spec.summary = 'ObjectSpace class'
+
+  spec.add_test_dependency('mruby-metaprog')
+  spec.add_test_dependency('mruby-method')
+  spec.add_test_dependency('mruby-fiber')
 end
diff --git a/mrbgems/mruby-objectspace/src/mruby_objectspace.c b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
index b89fb0580..7892c6a1b 100644
--- a/mrbgems/mruby-objectspace/src/mruby_objectspace.c
+++ b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
@@ -2,6 +2,14 @@
 #include <mruby/gc.h>
 #include <mruby/hash.h>
 #include <mruby/class.h>
+#include <mruby/object.h>
+#include <mruby/numeric.h>
+#include <mruby/string.h>
+#include <mruby/array.h>
+#include <mruby/variable.h>
+#include <mruby/proc.h>
+#include <mruby/value.h>
+#include <mruby/range.h>
 
 struct os_count_struct {
   mrb_int total;
@@ -168,12 +176,250 @@ os_each_object(mrb_state *mrb, mrb_value self)
   return mrb_fixnum_value(d.count);
 }
 
+static void os_memsize_of_object(mrb_state*,mrb_value,mrb_value,mrb_int*);
+
+struct os_memsize_cb_data {
+  mrb_int *t;
+  mrb_value recurse;
+};
+
+static int
+os_memsize_ivar_cb(mrb_state *mrb, mrb_sym _name, mrb_value obj, void *data)
+{
+  struct os_memsize_cb_data *cb_data = (struct os_memsize_cb_data *)(data);
+  os_memsize_of_object(mrb, obj, cb_data->recurse, cb_data->t);
+  return 0;
+}
+
+static void
+os_memsize_of_ivars(mrb_state* mrb, mrb_value obj, mrb_value recurse, mrb_int *t)
+{
+  (*t) += mrb_obj_iv_tbl_memsize(mrb, obj);
+  if(!mrb_nil_p(recurse)) {
+    struct os_memsize_cb_data cb_data = {t, recurse};
+    mrb_iv_foreach(mrb, obj, os_memsize_ivar_cb, &cb_data);
+  }
+}
+
+static void
+os_memsize_of_irep(mrb_state* state, struct mrb_irep *irep, mrb_int* t)
+{
+  mrb_int i;
+  (*t) += (irep->slen * sizeof(mrb_sym)) +
+          (irep->plen * sizeof(mrb_code)) +
+          (irep->ilen * sizeof(mrb_code));
+
+  for(i = 0; i < irep->rlen; i++) {
+    os_memsize_of_irep(state, irep->reps[i], t);
+  }
+}
+
+static void
+os_memsize_of_method(mrb_state* mrb, mrb_value method_obj, mrb_int* t)
+{
+  mrb_value proc_value = mrb_obj_iv_get(mrb, mrb_obj_ptr(method_obj),
+                                        mrb_intern_lit(mrb, "_proc"));
+  struct RProc *proc = mrb_proc_ptr(proc_value);
+
+  (*t) += sizeof(struct RProc);
+  if(!MRB_PROC_CFUNC_P(proc)) os_memsize_of_irep(mrb, proc->body.irep, t);
+}
+
+static void
+os_memsize_of_methods(mrb_state* mrb, mrb_value obj, mrb_int* t)
+{
+  mrb_value method_list;
+  mrb_int i;
+  if(!mrb_respond_to(mrb, obj, mrb_intern_lit(mrb, "instance_methods"))) return;
+  method_list = mrb_funcall(mrb, obj, "instance_methods", 1, mrb_false_value());
+  for(i = 0; i < RARRAY_LEN(method_list); i++) {
+    mrb_value method = mrb_funcall(mrb, obj, "instance_method", 1,
+                                   mrb_ary_ref(mrb, method_list, i));
+    os_memsize_of_method(mrb, method, t);
+  }
+}
+
+static void
+os_memsize_of_object(mrb_state* mrb, mrb_value obj, mrb_value recurse, mrb_int* t)
+{
+  if(!mrb_nil_p(recurse)) {
+    const mrb_value obj_id = mrb_fixnum_value(mrb_obj_id(obj));
+    if(mrb_hash_key_p(mrb, recurse, obj_id)) return;
+    mrb_hash_set(mrb, recurse, obj_id, mrb_true_value());
+  }
+
+  switch(obj.tt) {
+    case MRB_TT_STRING:
+      (*t) += RSTRING_LEN(obj);
+      break;
+    case MRB_TT_CLASS:
+    case MRB_TT_MODULE:
+    case MRB_TT_EXCEPTION:
+    case MRB_TT_SCLASS:
+    case MRB_TT_ICLASS:
+    case MRB_TT_OBJECT: {
+      (*t) += mrb_objspace_page_slot_size();
+      os_memsize_of_ivars(mrb, obj, recurse, t);
+      if(mrb_obj_is_kind_of(mrb, obj, mrb_class_get(mrb, "UnboundMethod"))) {
+        os_memsize_of_method(mrb, obj, t);
+      }
+      else {
+        os_memsize_of_methods(mrb, obj, t);
+      }
+      break;
+    }
+    case MRB_TT_HASH: {
+      (*t) += mrb_objspace_page_slot_size() +
+              mrb_os_memsize_of_hash_table(obj);
+      if(!mrb_nil_p(recurse)) {
+        os_memsize_of_object(mrb, mrb_hash_keys(mrb, obj), recurse, t);
+        os_memsize_of_object(mrb, mrb_hash_values(mrb, obj), recurse, t);
+      }
+      break;
+    }
+    case MRB_TT_ARRAY: {
+      mrb_int len, i;
+      len = RARRAY_LEN(obj);
+      /* Arrays that do not fit within an RArray perform a heap allocation
+      *  storing an array of pointers to the original objects*/
+      (*t) += mrb_objspace_page_slot_size();
+      if(len > MRB_ARY_EMBED_LEN_MAX) (*t) += sizeof(mrb_value *) * len;
+
+      if(!mrb_nil_p(recurse)) {
+        for(i = 0; i < len; i++) {
+          os_memsize_of_object(mrb, ARY_PTR(mrb_ary_ptr(obj))[i], recurse, t);
+        }
+      }
+      break;
+    }
+    case MRB_TT_PROC: {
+      struct RProc* proc = mrb_proc_ptr(obj);
+      (*t) += mrb_objspace_page_slot_size();
+      (*t) += MRB_ENV_LEN(proc->e.env) * sizeof(mrb_value);
+      if(!MRB_PROC_CFUNC_P(proc)) os_memsize_of_irep(mrb, proc->body.irep, t);
+      break;
+    }
+    case MRB_TT_DATA:
+      (*t) += mrb_objspace_page_slot_size();
+      if(mrb_respond_to(mrb, obj, mrb_intern_lit(mrb, "memsize"))) {
+        (*t) += mrb_fixnum(mrb_funcall(mrb, obj, "memsize", 0));
+      }
+      break;
+    #ifndef MRB_WITHOUT_FLOAT
+    case MRB_TT_FLOAT:
+      #ifdef MRB_WORD_BOXING
+      (*t) += mrb_objspace_page_slot_size() +
+              sizeof(struct RFloat);
+      #endif
+      break;
+    #endif
+    case MRB_TT_RANGE:
+    #ifndef MRB_RANGE_EMBED
+      (*t) += mrb_objspace_page_slot_size() +
+              sizeof(struct mrb_range_edges);
+    #endif
+      break;
+    case MRB_TT_FIBER: {
+      struct RFiber* f = (struct RFiber *)mrb_ptr(obj);
+      mrb_callinfo *ci_p = f->cxt->cibase;
+      ptrdiff_t stack_size = f->cxt->stend - f->cxt->stbase;
+      ptrdiff_t ci_size = f->cxt->ciend - f->cxt->cibase;
+      mrb_int i = 0;
+
+      while(ci_p < f->cxt->ciend) {
+        if(ci_p->proc) os_memsize_of_irep(mrb, ci_p->proc->body.irep, t);
+        ci_p++;
+      }
+
+      if(f->cxt->esize) {
+        for(i = 0; i <= f->cxt->esize; i++) {
+          os_memsize_of_irep(mrb, f->cxt->ensure[i]->body.irep, t);
+        }
+      }
+
+      (*t) += mrb_objspace_page_slot_size() +
+        sizeof(struct RFiber) +
+        sizeof(struct mrb_context) +
+        sizeof(struct RProc *) * f->cxt->esize +
+        sizeof(uint16_t *) * f->cxt->rsize +
+        stack_size +
+        ci_size;
+      break;
+    }
+    case MRB_TT_ISTRUCT:
+      (*t) += mrb_objspace_page_slot_size();
+      break;
+    /*  zero heap size types.
+     *  immediate VM stack values, contained within mrb_state, or on C stack */
+    case MRB_TT_TRUE:
+    case MRB_TT_FALSE:
+    case MRB_TT_FIXNUM:
+    case MRB_TT_BREAK:
+    case MRB_TT_CPTR:
+    case MRB_TT_SYMBOL:
+    case MRB_TT_FREE:
+    case MRB_TT_UNDEF:
+    case MRB_TT_ENV:
+    /* never used, silences compiler warning
+     * not having a default: clause lets the compiler tell us when there is a new
+     * TT not accounted for */
+    case MRB_TT_MAXDEFINE:
+      break;
+  }
+}
+
+/*
+ *  call-seq:
+ *    ObjectSpace.memsize_of(obj, recurse: false) -> Numeric
+ *
+ *  Returns the amount of heap memory allocated for object in size_t units.
+ *
+ *  The return value depends on the definition of size_t on that platform,
+ *  therefore the value is not comparable across platform types.
+ *
+ *  Immediate values such as integers, booleans, symbols and unboxed float numbers
+ *  return 0. Additionally special objects which are small enough to fit inside an
+ *  object pointer, termed embedded objects, will return the size of the object pointer.
+ *  Strings and arrays below a compile-time defined size may be embedded.
+ *
+ *  Setting recurse: true descends into instance variables, array members,
+ *  hash keys and hash values recursively, calculating the child objects and adding to
+ *  the final sum. It avoids infinite recursion and over counting objects by
+ *  internally tracking discovered object ids.
+ *
+ *  MRB_TT_DATA objects aren't calculated beyond their original page slot. However,
+ *  if the object implements a memsize method it will call that method and add the
+ *  return value to the total. This provides an opportunity for C based data structures
+ *  to report their memory usage.
+ *
+ */
+
+static mrb_value
+os_memsize_of(mrb_state *mrb, mrb_value self)
+{
+  mrb_int total;
+  mrb_value obj;
+  mrb_value recurse;
+  const char *kw_names[1] = { "recurse" };
+  mrb_value kw_values[1];
+  const mrb_kwargs kwargs = { 1, kw_values, kw_names, 0, NULL };
+
+  mrb_get_args(mrb, "o:", &obj, &kwargs);
+  recurse = mrb_obj_eq(mrb, kw_values[0], mrb_true_value())? mrb_hash_new(mrb) : mrb_nil_value();
+
+  total = 0;
+  os_memsize_of_object(mrb, obj, recurse, &total);
+
+  return mrb_fixnum_value(total);
+}
+
 void
 mrb_mruby_objectspace_gem_init(mrb_state *mrb)
 {
   struct RClass *os = mrb_define_module(mrb, "ObjectSpace");
   mrb_define_class_method(mrb, os, "count_objects", os_count_objects, MRB_ARGS_OPT(1));
   mrb_define_class_method(mrb, os, "each_object", os_each_object, MRB_ARGS_OPT(1));
+  mrb_define_class_method(mrb, os, "memsize_of", os_memsize_of, MRB_ARGS_REQ(1)|MRB_ARGS_OPT(1));
 }
 
 void
diff --git a/mrbgems/mruby-objectspace/test/objectspace.rb b/mrbgems/mruby-objectspace/test/objectspace.rb
index 9c44c2157..610cdfbfa 100644
--- a/mrbgems/mruby-objectspace/test/objectspace.rb
+++ b/mrbgems/mruby-objectspace/test/objectspace.rb
@@ -1,3 +1,4 @@
+# coding: utf-8
 assert('ObjectSpace.count_objects') do
   h = {}
   f = Fiber.new {} if Object.const_defined?(:Fiber)
@@ -58,3 +59,146 @@ end
 assert 'Check class pointer of ObjectSpace.each_object.' do
   assert_nothing_raised { ObjectSpace.each_object { |obj| !obj } }
 end
+
+assert 'ObjectSpace.memsize_of' do
+  # immediate literals
+  int_size = ObjectSpace.memsize_of 1
+  assert_equal int_size, 0, 'int zero'
+
+  sym_size = ObjectSpace.memsize_of :foo
+  assert_equal sym_size, 0, 'sym zero'
+
+  assert_equal ObjectSpace.memsize_of(true), int_size
+  assert_equal ObjectSpace.memsize_of(false), int_size
+
+  float_size = if Object.const_defined? :Float
+                 ObjectSpace.memsize_of 1.0
+               else
+                 nil
+               end
+
+  # need some way of asking if floats are boxed
+  assert_equal float_size, 0 if float_size
+
+  assert_not_equal ObjectSpace.memsize_of('a'), 0, 'memsize of str'
+
+  if __ENCODING__ == "UTF-8"
+    assert_not_equal ObjectSpace.memsize_of("こんにちは世界"), 0, 'memsize of utf8 str'
+  end
+
+  assert_not_equal ObjectSpace.memsize_of(0..1), 0, 'range not zero'
+
+  # class defs
+  class_obj_size = ObjectSpace.memsize_of Class
+  assert_not_equal class_obj_size, 0, 'Class obj not zero'
+
+  empty_class_def_size = ObjectSpace.memsize_of Class.new
+  assert_not_equal empty_class_def_size, 0, 'Class def not zero'
+
+  class_without_methods = Class.new do
+    @a = 1
+    @b = 2
+  end
+  class_total_size = empty_class_def_size + (int_size * 2)
+  assert_equal ObjectSpace.memsize_of(class_without_methods), class_total_size, 'class without methods size'
+
+  module_without_methods = Module.new do
+    @a = 1
+    @b = 2
+  end
+  module_total_size = empty_class_def_size + (int_size * 2)
+  assert_equal ObjectSpace.memsize_of(module_without_methods), module_total_size, 'module without methods size'
+
+  proc_size = ObjectSpace.memsize_of Proc.new { x = 1; x }
+  assert_not_equal proc_size, 0
+
+  class_with_methods = Class.new do
+    def foo
+      a = 0
+      a + 1
+    end
+  end
+
+  m_size = ObjectSpace.memsize_of class_with_methods.instance_method(:foo)
+  assert_not_equal m_size, 0, 'method size not zero'
+
+  # collections
+  empty_array_size = ObjectSpace.memsize_of []
+  assert_not_equal empty_array_size, 0, 'empty array size not zero'
+  assert_operator empty_array_size, :<, ObjectSpace.memsize_of(Array.new(16)), 'large array size greater than embed'
+
+  # fiber
+  empty_fiber_size = ObjectSpace.memsize_of(Fiber.new {})
+  assert_not_equal empty_fiber_size, 0, 'empty fiber not zero'
+  assert_operator empty_fiber_size, :<, ObjectSpace.memsize_of(Fiber.new { yield; 1 }), 'Fiber code size growth'
+
+  #hash
+  assert_not_equal ObjectSpace.memsize_of({}), 0, 'empty hash size not zero'
+
+  # recursion
+  foo_str = 'foo' * 10
+  bar_str = 'bar' * 10
+  caz_str = 'caz' * 10
+  fbc_ary = [foo_str, bar_str, caz_str]
+  assert_operator ObjectSpace.memsize_of(fbc_ary),
+                  :<,
+                  ObjectSpace.memsize_of(fbc_ary, recurse: true),
+                  'basic array recursion'
+
+  big_ary = [ 'a' * 10,
+              [ 'b' * 10,
+                [ 'c' * 10,
+                  [ 'd' * 10,
+                    [ 'e' * 10,
+                      [ 'f' * 10,
+                        ['g' * 10]
+                      ] * 10,
+                    ] * 10,
+                  ] * 10,
+                ] * 10,
+              ] * 10,
+            ] * 10
+  assert_operator ObjectSpace.memsize_of(big_ary),
+                  :<,
+                  ObjectSpace.memsize_of(big_ary, recurse: true),
+                  'large array recursion'
+
+  assert_nothing_raised 'infinite array recursion' do
+    ObjectSpace.memsize_of(fbc_ary.push(fbc_ary))
+  end
+
+  basic_hsh = {a: [foo_str, bar_str], b: caz_str, 'c' => {}}
+  assert_operator ObjectSpace.memsize_of(basic_hsh),
+                  :<,
+                  ObjectSpace.memsize_of(basic_hsh, recurse: true),
+                  'hash recursion with basic keys'
+
+  weird_keys = {big_ary => foo_str}
+  assert_operator ObjectSpace.memsize_of(weird_keys),
+                  :<,
+                  ObjectSpace.memsize_of(weird_keys, recurse: true),
+                  'hash recursion with collection as key'
+
+  basic_hsh.store('d', basic_hsh)
+  assert_nothing_raised 'hash value recursion' do
+    ObjectSpace.memsize_of basic_hsh, recurse: true
+  end
+
+  foo_klass = Class.new do
+    def bar= b
+      @bar = b
+    end
+  end
+
+  fk_one = foo_klass.new
+  fk_one.bar = fbc_ary
+  assert_operator ObjectSpace.memsize_of(fk_one),
+                  :<,
+                  ObjectSpace.memsize_of(fk_one, recurse: true),
+                  'basic ivar recursion'
+
+  fk_one.bar = fk_one
+  assert_nothing_raised 'ivar infinite recursion' do
+    ObjectSpace.memsize_of(fk_one, recurse: true)
+  end
+end