Allow `mrb_objspace_each_objects()` to break iteration; ref #3359

3 files changed, 13 insertions, 8 deletions

diff --git a/include/mruby/gc.h b/include/mruby/gc.h
index 81950f169..dd161efa1 100644
--- a/include/mruby/gc.h
+++ b/include/mruby/gc.h
@@ -17,7 +17,9 @@ MRB_BEGIN_DECL
 
 struct mrb_state;
 
-typedef void (mrb_each_object_callback)(struct mrb_state *mrb, struct RBasic *obj, void *data);
+#define MRB_EACH_OBJ_OK 0
+#define MRB_EACH_OBJ_BREAK 1
+typedef int (mrb_each_object_callback)(struct mrb_state *mrb, struct RBasic *obj, void *data);
 void mrb_objspace_each_objects(struct mrb_state *mrb, mrb_each_object_callback *callback, void *data);
 MRB_API void mrb_free_context(struct mrb_state *mrb, struct mrb_context *c);
 
diff --git a/mrbgems/mruby-objectspace/src/mruby_objectspace.c b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
index d0a8effd0..3887091d3 100644
--- a/mrbgems/mruby-objectspace/src/mruby_objectspace.c
+++ b/mrbgems/mruby-objectspace/src/mruby_objectspace.c
@@ -9,7 +9,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 +23,7 @@ os_count_object_type(mrb_state *mrb, struct RBasic *obj, void *data)
   else {
     obj_count->counts[obj->tt]++;
   }
+  return MRB_EACH_OBJ_OK;
 }
 
 /*
@@ -109,35 +110,36 @@ struct os_each_object_data {
   mrb_int count;
 };
 
-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;
+    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;
 }
 
 /*
diff --git a/src/gc.c b/src/gc.c
index 81f47eed7..833d35c71 100644
--- a/src/gc.c
+++ b/src/gc.c
@@ -1498,7 +1498,8 @@ gc_each_objects(mrb_state *mrb, mrb_gc *gc, mrb_each_object_callback *callback,
     p = objects(page);
     pend = p + MRB_HEAP_PAGE_SIZE;
     for (;p < pend; p++) {
-      (*callback)(mrb, &p->as.basic, data);
+      if ((*callback)(mrb, &p->as.basic, data) == MRB_EACH_OBJ_BREAK)
+        break;
     }
 
     page = page->next;