3 files changed, 148 insertions, 87 deletions

diff --git a/mrbgems/mruby-objectspace/mrbgem.rake b/mrbgems/mruby-objectspace/mrbgem.rake
index 6a6a3e778..fa35136a1 100644
--- a/mrbgems/mruby-objectspace/mrbgem.rake
+++ b/mrbgems/mruby-objectspace/mrbgem.rake
@@ -1,4 +1,5 @@
 MRuby::Gem::Specification.new('mruby-objectspace') do |spec|
   spec.license = 'MIT'
   spec.author  = 'mruby developers'
+  spec.summary = 'ObjectSpace class'
 end
diff --git a/mrbgems/mruby-objectspace/src/mruby_objectspace.c b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
index ccca6bcbc..7d3b0b74e 100644
--- a/mrbgems/mruby-objectspace/src/mruby_objectspace.c
+++ b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
@@ -1,25 +1,27 @@
 #include <mruby.h>
 #include <mruby/gc.h>
 #include <mruby/hash.h>
-#include <mruby/value.h>
+#include <mruby/class.h>
 
 struct os_count_struct {
-  size_t total;
-  size_t freed;
-  size_t counts[MRB_TT_MAXDEFINE+1];
+  mrb_int total;
+  mrb_int freed;
+  mrb_int counts[MRB_TT_MAXDEFINE+1];
 };
 
-void
-os_count_object_type(mrb_state *mrb, struct RBasic* obj, void *data)
+static void
+os_count_object_type(mrb_state *mrb, struct RBasic *obj, void *data)
 {
-  struct os_count_struct* obj_count;
-  obj_count = (struct os_count_struct*)(data);
+  struct os_count_struct *obj_count;
+  obj_count = (struct os_count_struct*)data;
+
+  obj_count->total++;
 
   if (is_dead(mrb, obj)) {
     obj_count->freed++;
-  } else {
+  }
+  else {
     obj_count->counts[obj->tt]++;
-    obj_count->total++;
   }
 }
 
@@ -33,8 +35,8 @@ os_count_object_type(mrb_state *mrb, struct RBasic* obj, void *data)
  *  {
  *    :TOTAL=>10000,
  *    :FREE=>3011,
- *    :MRB_TT_OBJECT=>6,
- *    :MRB_TT_CLASS=>404,
+ *    :T_OBJECT=>6,
+ *    :T_CLASS=>404,
  *    # ...
  *  }
  *
@@ -43,11 +45,11 @@ os_count_object_type(mrb_state *mrb, struct RBasic* obj, void *data)
  *
  */
 
-mrb_value
+static mrb_value
 os_count_objects(mrb_state *mrb, mrb_value self)
 {
-  struct os_count_struct obj_count;
-  size_t i;
+  struct os_count_struct obj_count = { 0 };
+  enum mrb_vtype i;
   mrb_value hash;
 
   if (mrb_get_args(mrb, "|H", &hash) == 0) {
@@ -58,43 +60,38 @@ os_count_objects(mrb_state *mrb, mrb_value self)
     mrb_hash_clear(mrb, hash);
   }
 
-  for (i = 0; i <= MRB_TT_MAXDEFINE; i++) {
-    obj_count.counts[i] = 0;
-  }
-  obj_count.total = 0;
-  obj_count.freed = 0;
-
   mrb_objspace_each_objects(mrb, os_count_object_type, &obj_count);
 
-  mrb_hash_set(mrb, hash, mrb_symbol_value(mrb_intern_cstr(mrb, "TOTAL")), mrb_fixnum_value(obj_count.total));
-  mrb_hash_set(mrb, hash, mrb_symbol_value(mrb_intern_cstr(mrb, "FREE")), mrb_fixnum_value(obj_count.freed));
+  mrb_hash_set(mrb, hash, mrb_symbol_value(mrb_intern_lit(mrb, "TOTAL")), mrb_fixnum_value(obj_count.total));
+  mrb_hash_set(mrb, hash, mrb_symbol_value(mrb_intern_lit(mrb, "FREE")), mrb_fixnum_value(obj_count.freed));
 
-  for (i = 0; i < MRB_TT_MAXDEFINE; i++) {
+  for (i = MRB_TT_FALSE; i < MRB_TT_MAXDEFINE; i++) {
     mrb_value type;
     switch (i) {
-#define COUNT_TYPE(t) case (t): type = mrb_symbol_value(mrb_intern_cstr(mrb, #t)); break;
-      COUNT_TYPE(MRB_TT_FALSE);
-      COUNT_TYPE(MRB_TT_FREE);
-      COUNT_TYPE(MRB_TT_TRUE);
-      COUNT_TYPE(MRB_TT_FIXNUM);
-      COUNT_TYPE(MRB_TT_SYMBOL);
-      COUNT_TYPE(MRB_TT_UNDEF);
-      COUNT_TYPE(MRB_TT_FLOAT);
-      COUNT_TYPE(MRB_TT_VOIDP);
-      COUNT_TYPE(MRB_TT_OBJECT);
-      COUNT_TYPE(MRB_TT_CLASS);
-      COUNT_TYPE(MRB_TT_MODULE);
-      COUNT_TYPE(MRB_TT_ICLASS);
-      COUNT_TYPE(MRB_TT_SCLASS);
-      COUNT_TYPE(MRB_TT_PROC);
-      COUNT_TYPE(MRB_TT_ARRAY);
-      COUNT_TYPE(MRB_TT_HASH);
-      COUNT_TYPE(MRB_TT_STRING);
-      COUNT_TYPE(MRB_TT_RANGE);
-      COUNT_TYPE(MRB_TT_EXCEPTION);
-      COUNT_TYPE(MRB_TT_FILE);
-      COUNT_TYPE(MRB_TT_ENV);
-      COUNT_TYPE(MRB_TT_DATA);
+#define COUNT_TYPE(t) case (MRB_T ## t): type = mrb_symbol_value(mrb_intern_lit(mrb, #t)); break;
+      COUNT_TYPE(T_FALSE);
+      COUNT_TYPE(T_FREE);
+      COUNT_TYPE(T_TRUE);
+      COUNT_TYPE(T_FIXNUM);
+      COUNT_TYPE(T_SYMBOL);
+      COUNT_TYPE(T_UNDEF);
+      COUNT_TYPE(T_FLOAT);
+      COUNT_TYPE(T_CPTR);
+      COUNT_TYPE(T_OBJECT);
+      COUNT_TYPE(T_CLASS);
+      COUNT_TYPE(T_MODULE);
+      COUNT_TYPE(T_ICLASS);
+      COUNT_TYPE(T_SCLASS);
+      COUNT_TYPE(T_PROC);
+      COUNT_TYPE(T_ARRAY);
+      COUNT_TYPE(T_HASH);
+      COUNT_TYPE(T_STRING);
+      COUNT_TYPE(T_RANGE);
+      COUNT_TYPE(T_EXCEPTION);
+      COUNT_TYPE(T_FILE);
+      COUNT_TYPE(T_ENV);
+      COUNT_TYPE(T_DATA);
+      COUNT_TYPE(T_FIBER);
 #undef COUNT_TYPE
     default:
       type = mrb_fixnum_value(i); break;
@@ -106,12 +103,57 @@ os_count_objects(mrb_state *mrb, mrb_value self)
   return hash;
 }
 
+struct os_each_object_data {
+  mrb_value block;
+  struct RClass *target_module;
+  mrb_int count;
+};
+
+static void
+os_each_object_cb(mrb_state *mrb, struct RBasic *obj, void *ud)
+{
+  struct os_each_object_data *d = (struct os_each_object_data*)ud;
+
+  /* filter dead objects */
+  if (is_dead(mrb, obj)) {
+    return;
+  }
+
+  /* filter class kind if target module defined */
+  if (d->target_module && !mrb_obj_is_kind_of(mrb, mrb_obj_value(obj), d->target_module)) {
+    return;
+  }
+
+  mrb_yield(mrb, d->block, mrb_obj_value(obj));
+  ++d->count;
+}
+
+static mrb_value
+os_each_object(mrb_state *mrb, mrb_value self)
+{
+  mrb_value cls = mrb_nil_value();
+  struct os_each_object_data d;
+  mrb_get_args(mrb, "&|C", &d.block, &cls);
+
+  if (mrb_nil_p(d.block)) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "Expected block in ObjectSpace.each_object.");
+  }
+
+  d.target_module = mrb_nil_p(cls) ? NULL : mrb_class_ptr(cls);
+  d.count = 0;
+  mrb_objspace_each_objects(mrb, os_each_object_cb, &d);
+  return mrb_fixnum_value(d.count);
+}
+
 void
-mrb_mruby_objectspace_gem_init(mrb_state* mrb) {
+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_ANY());
+  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));
 }
 
 void
-mrb_mruby_objectspace_gem_final(mrb_state* mrb) {
+mrb_mruby_objectspace_gem_final(mrb_state *mrb)
+{
 }
diff --git a/mrbgems/mruby-objectspace/test/objectspace.rb b/mrbgems/mruby-objectspace/test/objectspace.rb
index 37137eb04..f792c165d 100644
--- a/mrbgems/mruby-objectspace/test/objectspace.rb
+++ b/mrbgems/mruby-objectspace/test/objectspace.rb
@@ -1,38 +1,56 @@
 assert('ObjectSpace.count_objects') do
-	h = {}
-	ObjectSpace.count_objects(h)
-	assert_kind_of(Hash, h)
-	assert_true(h.keys.all? {|x| x.is_a?(Symbol) || x.is_a?(Integer) })
-	assert_true(h.values.all? {|x| x.is_a?(Integer) })
-
-	assert_true(h.has_key?(:TOTAL))
-	assert_true(h.has_key?(:FREE))
-
-	h = ObjectSpace.count_objects
-	assert_kind_of(Hash, h)
-	assert_true(h.keys.all? {|x| x.is_a?(Symbol) || x.is_a?(Integer) })
-	assert_true(h.values.all? {|x| x.is_a?(Integer) })
-
-	assert_raise(TypeError) { ObjectSpace.count_objects(1) }
-
-	h0 = {:MRB_TT_FOO=>1000}
-	h = ObjectSpace.count_objects(h0)
-	assert_false(h0.has_key?(:MRB_TT_FOO))
-
-	GC.start
-	h_after = {}
-	h_before = ObjectSpace.count_objects
-
-	objs = []
-	1000.times do
-		objs << {}
-	end
-	objs = nil
-	ObjectSpace.count_objects(h)
-	GC.start
-	ObjectSpace.count_objects(h_after)
-
-	assert_equal(h_before[:MRB_TT_HASH] + 1000, h[:MRB_TT_HASH])
-	assert_equal(h_before[:MRB_TT_HASH], h_after[:MRB_TT_HASH])
-
-end
\ No newline at end of file
+  h = {}
+  f = Fiber.new {} if Object.const_defined? :Fiber
+  ObjectSpace.count_objects(h)
+  assert_kind_of(Hash, h)
+  assert_true(h.keys.all? {|x| x.is_a?(Symbol) || x.is_a?(Integer) })
+  assert_true(h.values.all? {|x| x.is_a?(Integer) })
+
+  assert_true(h.has_key?(:TOTAL))
+  assert_true(h.has_key?(:FREE))
+  assert_true(h.has_key?(:T_FIBER)) if Object.const_defined? :Fiber
+
+  assert_equal(h[:TOTAL] * 2, h.values.reduce(:+))
+
+  h = ObjectSpace.count_objects
+  assert_kind_of(Hash, h)
+  assert_true(h.keys.all? {|x| x.is_a?(Symbol) || x.is_a?(Integer) })
+  assert_true(h.values.all? {|x| x.is_a?(Integer) })
+
+  assert_raise(TypeError) { ObjectSpace.count_objects(1) }
+
+  h0 = {:T_FOO=>1000}
+  h = ObjectSpace.count_objects(h0)
+  assert_false(h0.has_key?(:T_FOO))
+
+  GC.start
+  h_after = {}
+  h_before = ObjectSpace.count_objects
+
+  objs = []
+  1000.times do
+    objs << {}
+  end
+  objs = nil
+  ObjectSpace.count_objects(h)
+  GC.start
+  ObjectSpace.count_objects(h_after)
+
+  assert_equal(h[:T_HASH], h_before[:T_HASH] + 1000)
+  assert_equal(h_after[:T_HASH], h_before[:T_HASH])
+end
+
+assert('ObjectSpace.each_object') do
+  objs = []
+  objs_count = ObjectSpace.each_object { |obj|
+    objs << obj
+  }
+  assert_equal objs.length, objs_count
+
+  arys = []
+  arys_count = ObjectSpace.each_object(Array) { |obj|
+    arys << obj
+  }
+  assert_equal arys.length, arys_count
+  assert_true arys.length < objs.length
+end