1 files changed, 58 insertions, 2 deletions

diff --git a/src/gc.c b/src/gc.c
index 8bd8243f1..885368f37 100644
--- a/src/gc.c
+++ b/src/gc.c
@@ -385,6 +385,7 @@ gc_protect(mrb_state *mrb, struct RBasic *p)
   mrb->arena[mrb->arena_idx++] = p;
 }
 
+/* mrb_gc_protect() leaves the object in the arena */
 MRB_API void
 mrb_gc_protect(mrb_state *mrb, mrb_value obj)
 {
@@ -392,6 +393,53 @@ mrb_gc_protect(mrb_state *mrb, mrb_value obj)
   gc_protect(mrb, mrb_basic_ptr(obj));
 }
 
+#define GC_ROOT_NAME "_gc_root_"
+
+/* mrb_gc_register() keeps the object from GC.
+
+   Register your object when it's exported to C world,
+   without reference from Ruby world, e.g. callback
+   arguments.  Don't forget to remove the obejct using
+   mrb_gc_unregister, otherwise your object will leak.
+*/
+
+MRB_API void
+mrb_gc_register(mrb_state *mrb, mrb_value obj)
+{
+  mrb_sym root = mrb_intern_lit(mrb, GC_ROOT_NAME);
+  mrb_value table = mrb_gv_get(mrb, root);
+
+  if (mrb_nil_p(table) || mrb_type(table) != MRB_TT_ARRAY) {
+    table = mrb_ary_new(mrb);
+    mrb_gv_set(mrb, root, table);
+  }
+  mrb_ary_push(mrb, table, obj);
+}
+
+/* mrb_gc_unregister() removes the object from GC root. */
+MRB_API void
+mrb_gc_unregister(mrb_state *mrb, mrb_value obj)
+{
+  mrb_sym root = mrb_intern_lit(mrb, GC_ROOT_NAME);
+  mrb_value table = mrb_gv_get(mrb, root);
+  struct RArray *a;
+  mrb_int i, j;
+
+  if (mrb_nil_p(table)) return;
+  if (mrb_type(table) != MRB_TT_ARRAY) {
+    mrb_gv_set(mrb, root, mrb_nil_value());
+    return;
+  }
+  a = mrb_ary_ptr(table);
+  mrb_ary_modify(mrb, a);
+  for (i=j=0; i<a->len; i++) {
+    if (!mrb_obj_eq(mrb, a->ptr[i], obj)) {
+      a->ptr[j++] = a->ptr[i];
+    }
+  }
+  a->len = j;
+}
+
 MRB_API struct RBasic*
 mrb_obj_alloc(mrb_state *mrb, enum mrb_vtype ttype, struct RClass *cls)
 {
@@ -498,7 +546,12 @@ gc_mark_children(mrb_state *mrb, struct RBasic *obj)
   mrb_gc_mark(mrb, (struct RBasic*)obj->c);
   switch (obj->tt) {
   case MRB_TT_ICLASS:
-    mrb_gc_mark(mrb, (struct RBasic*)((struct RClass*)obj)->super);
+    {
+      struct RClass *c = (struct RClass*)obj;
+      if (MRB_FLAG_TEST(c, MRB_FLAG_IS_ORIGIN))
+        mrb_gc_mark_mt(mrb, c);
+      mrb_gc_mark(mrb, (struct RBasic*)((struct RClass*)obj)->super);
+    }
     break;
 
   case MRB_TT_CLASS:
@@ -624,7 +677,10 @@ obj_free(mrb_state *mrb, struct RBasic *obj)
     mrb_gc_free_mt(mrb, (struct RClass*)obj);
     mrb_gc_free_iv(mrb, (struct RObject*)obj);
     break;
-
+  case MRB_TT_ICLASS:
+    if (MRB_FLAG_TEST(obj, MRB_FLAG_IS_ORIGIN))
+      mrb_gc_free_mt(mrb, (struct RClass*)obj);
+    break;
   case MRB_TT_ENV:
     {
       struct REnv *e = (struct REnv*)obj;