4 files changed, 62 insertions, 0 deletions

diff --git a/include/mruby/class.h b/include/mruby/class.h
index 28ba6b1b7..3c4915dc4 100644
--- a/include/mruby/class.h
+++ b/include/mruby/class.h
@@ -41,6 +41,8 @@ mrb_class(mrb_state *mrb, mrb_value v)
     return mrb->float_class;
   case MRB_TT_CPTR:
     return mrb->object_class;
+  case MRB_TT_ENV:
+    return NULL;
   default:
     return mrb_obj_ptr(v)->c;
   }
diff --git a/mrbgems/mruby-objectspace/src/mruby_objectspace.c b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
index 538959f2a..31139c429 100644
--- a/mrbgems/mruby-objectspace/src/mruby_objectspace.c
+++ b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
@@ -1,6 +1,7 @@
 #include <mruby.h>
 #include <mruby/gc.h>
 #include <mruby/hash.h>
+#include <mruby/class.h>
 
 struct os_count_struct {
   mrb_int total;
@@ -100,11 +101,54 @@ 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)
 {
   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));
 }
 
 void
diff --git a/mrbgems/mruby-objectspace/test/objectspace.rb b/mrbgems/mruby-objectspace/test/objectspace.rb
index 612137019..a619ffe37 100644
--- a/mrbgems/mruby-objectspace/test/objectspace.rb
+++ b/mrbgems/mruby-objectspace/test/objectspace.rb
@@ -35,3 +35,18 @@ assert('ObjectSpace.count_objects') do
   assert_equal(h[:MRB_TT_HASH], h_before[:MRB_TT_HASH] + 1000)
   assert_equal(h_after[:MRB_TT_HASH], h_before[:MRB_TT_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
diff --git a/src/proc.c b/src/proc.c
index 1ae1d495e..93354dd16 100644
--- a/src/proc.c
+++ b/src/proc.c
@@ -70,6 +70,7 @@ mrb_proc_new_cfunc(mrb_state *mrb, mrb_func_t func)
   p = (struct RProc*)mrb_obj_alloc(mrb, MRB_TT_PROC, mrb->proc_class);
   p->body.func = func;
   p->flags |= MRB_PROC_CFUNC;
+  p->env = 0;
 
   return p;
 }