1 files changed, 27 insertions, 24 deletions

diff --git a/mrbgems/mruby-objectspace/src/mruby_objectspace.c b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
index d0a8effd0..76c804fbb 100644
--- a/mrbgems/mruby-objectspace/src/mruby_objectspace.c
+++ b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
@@ -2,6 +2,15 @@
 #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>
+#include <mruby/presym.h>
 
 struct os_count_struct {
   mrb_int total;
@@ -9,7 +18,7 @@ struct os_count_struct {
   mrb_int counts[MRB_TT_MAXDEFINE+1];
 };
 
-static void
+static int
 os_count_object_type(mrb_state *mrb, struct RBasic *obj, void *data)
 {
   struct os_count_struct *obj_count;
@@ -23,6 +32,7 @@ os_count_object_type(mrb_state *mrb, struct RBasic *obj, void *data)
   else {
     obj_count->counts[obj->tt]++;
   }
+  return MRB_EACH_OBJ_OK;
 }
 
 /*
@@ -56,14 +66,14 @@ os_count_objects(mrb_state *mrb, mrb_value self)
     hash = mrb_hash_new(mrb);
   }
 
-  if (!mrb_test(mrb_hash_empty_p(mrb, hash))) {
+  if (!mrb_hash_empty_p(mrb, hash)) {
     mrb_hash_clear(mrb, hash);
   }
 
   mrb_objspace_each_objects(mrb, os_count_object_type, &obj_count);
 
-  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));
+  mrb_hash_set(mrb, hash, mrb_symbol_value(MRB_SYM(TOTAL)), mrb_fixnum_value(obj_count.total));
+  mrb_hash_set(mrb, hash, mrb_symbol_value(MRB_SYM(FREE)), mrb_fixnum_value(obj_count.freed));
 
   for (i = MRB_TT_FALSE; i < MRB_TT_MAXDEFINE; i++) {
     mrb_value type;
@@ -72,7 +82,7 @@ os_count_objects(mrb_state *mrb, mrb_value self)
       COUNT_TYPE(T_FALSE);
       COUNT_TYPE(T_FREE);
       COUNT_TYPE(T_TRUE);
-      COUNT_TYPE(T_FIXNUM);
+      COUNT_TYPE(T_INTEGER);
       COUNT_TYPE(T_SYMBOL);
       COUNT_TYPE(T_UNDEF);
       COUNT_TYPE(T_FLOAT);
@@ -88,10 +98,10 @@ os_count_objects(mrb_state *mrb, mrb_value self)
       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);
+      COUNT_TYPE(T_ISTRUCT);
 #undef COUNT_TYPE
     default:
       type = mrb_fixnum_value(i); break;
@@ -104,40 +114,41 @@ os_count_objects(mrb_state *mrb, mrb_value self)
 }
 
 struct os_each_object_data {
-  mrb_value block;
   struct RClass *target_module;
   mrb_int count;
+  mrb_value block;
 };
 
-static void
+static int
 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 (mrb_object_dead_p(mrb, obj)) {
-    return;
+    return MRB_EACH_OBJ_OK;
   }
 
   /* filter internal objects */
   switch (obj->tt) {
-  case MRB_TT_ENV:
-  case MRB_TT_ICLASS:
-    return;
+  case MRB_TT_FREE: case MRB_TT_ENV:
+  case MRB_TT_BREAK: case MRB_TT_ICLASS:
+    return MRB_EACH_OBJ_OK;
   default:
     break;
   }
 
   /* filter half baked (or internal) objects */
-  if (!obj->c) return;
+  if (!obj->c) return MRB_EACH_OBJ_OK;
 
   /* 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;
+    return MRB_EACH_OBJ_OK;
   }
 
   mrb_yield(mrb, d->block, mrb_obj_value(obj));
   ++d->count;
+  return MRB_EACH_OBJ_OK;
 }
 
 /*
@@ -157,16 +168,8 @@ os_each_object_cb(mrb_state *mrb, struct RBasic *obj, void *ud)
 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;
+  struct os_each_object_data d = {0};
+  mrb_get_args(mrb, "&!|c", &d.block, &d.target_module);
   mrb_objspace_each_objects(mrb, os_each_object_cb, &d);
   return mrb_fixnum_value(d.count);
 }