Merge branch 'module-prepend' of https://github.com/polyfox/mruby into polyfox-module-prepend

4 files changed, 168 insertions, 46 deletions

diff --git a/src/class.c b/src/class.c
index 33fb61211..462ab40b5 100644
--- a/src/class.c
+++ b/src/class.c
@@ -76,7 +76,8 @@ prepare_singleton_class(mrb_state *mrb, struct RBasic *o)
 
   if (o->c->tt == MRB_TT_SCLASS) return;
   sc = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_SCLASS, mrb->class_class);
-  sc->mt = 0;
+  sc->origin = sc;
+  sc->mt = kh_init(mt, mrb);
   sc->iv = 0;
   if (o->tt == MRB_TT_CLASS) {
     c = (struct RClass*)o;
@@ -194,6 +195,7 @@ define_class(mrb_state *mrb, mrb_sym name, struct RClass *super, struct RClass *
 
   if (mrb_const_defined_at(mrb, mrb_obj_value(outer), name)) {
     c = class_from_sym(mrb, outer, name);
+    c = c->origin;
     if (super && mrb_class_real(c->super) != super) {
       mrb_raisef(mrb, E_TYPE_ERROR, "superclass mismatch for Class %S (%S not %S)",
                  mrb_sym2str(mrb, name),
@@ -323,8 +325,9 @@ mrb_define_class_under(mrb_state *mrb, struct RClass *outer, const char *name, s
 MRB_API void
 mrb_define_method_raw(mrb_state *mrb, struct RClass *c, mrb_sym mid, struct RProc *p)
 {
-  khash_t(mt) *h = c->mt;
+  khash_t(mt) *h;
   khiter_t k;
+  h = c->origin->mt;
 
   if (!h) h = c->mt = kh_init(mt, mrb);
   k = kh_put(mt, mrb, h, mid);
@@ -806,6 +809,7 @@ boot_defclass(mrb_state *mrb, struct RClass *super)
   struct RClass *c;
 
   c = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_CLASS, mrb->class_class);
+  c->origin = c;
   if (super) {
     c->super = super;
     mrb_field_write_barrier(mrb, (struct RBasic*)c, (struct RBasic*)super);
@@ -817,47 +821,133 @@ boot_defclass(mrb_state *mrb, struct RClass *super)
   return c;
 }
 
-MRB_API void
-mrb_include_module(mrb_state *mrb, struct RClass *c, struct RClass *m)
+static void
+boot_initmod(mrb_state *mrb, struct RClass *mod)
+{
+  mod->origin = mod;
+  mod->mt = kh_init(mt, mrb);
+}
+
+static struct RClass*
+include_class_new(mrb_state *mrb, struct RClass *m, struct RClass *super)
+{
+  struct RClass *ic = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_ICLASS, mrb->class_class);
+  if (m->tt == MRB_TT_ICLASS) {
+    m = m->c;
+  }
+  ic->origin = ic;
+  ic->iv = m->iv;
+  ic->mt = m->origin->mt;
+  ic->super = super;
+  if (m->tt == MRB_TT_ICLASS) {
+    ic->c = m->c;
+  } else {
+    ic->c = m;
+  }
+  return ic;
+}
+
+static int
+include_module_at(mrb_state *mrb, struct RClass *c, struct RClass *ins_pos, struct RClass *m, int search_super)
 {
-  struct RClass *ins_pos;
+  struct RClass *p, *ic;
+  void *klass_mt = c->origin->mt;
 
-  ins_pos = c;
   while (m) {
-    struct RClass *p = c, *ic;
     int superclass_seen = 0;
 
-    if (c->mt && c->mt == m->mt) {
-      mrb_raise(mrb, E_ARGUMENT_ERROR, "cyclic include detected");
-    }
-    while (p) {
-      if (c != p && p->tt == MRB_TT_CLASS) {
-        superclass_seen = 1;
-      }
-      else if (p->mt == m->mt) {
-        if (p->tt == MRB_TT_ICLASS && !superclass_seen) {
-          ins_pos = p;
+    if (m->origin != m)
+      goto skip;
+
+    if (klass_mt && klass_mt == m->mt)
+      return -1;
+
+    p = c->super;
+    while(p) {
+      if (p->tt == MRB_TT_ICLASS) {
+        if (p->mt == m->mt) {
+          if (!superclass_seen) {
+            ins_pos = p; // move insert point
+          }
+          goto skip;
         }
-        goto skip;
+      } else if (p->tt == MRB_TT_CLASS) {
+        if (!search_super) break;
+        superclass_seen = 1;
       }
       p = p->super;
     }
-    ic = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_ICLASS, mrb->class_class);
-    if (m->tt == MRB_TT_ICLASS) {
-      ic->c = m->c;
-    }
-    else {
-      ic->c = m;
-    }
-    ic->mt = m->mt;
-    ic->iv = m->iv;
-    ic->super = ins_pos->super;
+
+    ic = include_class_new(mrb, m, ins_pos->super);
     ins_pos->super = ic;
-    mrb_field_write_barrier(mrb, (struct RBasic*)ins_pos, (struct RBasic*)ic);
+    mrb_field_write_barrier(mrb, (struct RBasic*)ins_pos, (struct RBasic*)ins_pos->super);
     ins_pos = ic;
   skip:
     m = m->super;
   }
+  return 0;
+}
+
+MRB_API void
+mrb_include_module(mrb_state *mrb, struct RClass *c, struct RClass *m)
+{
+  int changed = include_module_at(mrb, c, c->origin, m, 1);
+  if (changed < 0) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "cyclic include detected");
+  }
+}
+
+MRB_API void
+mrb_prepend_module(mrb_state *mrb, struct RClass *c, struct RClass *m)
+{
+  struct RClass *origin;
+  int changed = 0;
+
+  origin = c->origin;
+  if (origin == c) {
+    origin = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_ICLASS, c);
+    origin->flags |= MRB_FLAG_IS_ORIGIN;
+    origin->origin = origin;
+    origin->super = c->super;
+    c->super = origin;
+    c->origin = origin;
+    origin->mt = c->mt;
+    c->mt = kh_init(mt, mrb);
+    mrb_field_write_barrier(mrb, (struct RBasic*)c, (struct RBasic*)c->origin);
+  }
+  changed = include_module_at(mrb, c, c, m, 0);
+  if (changed < 0) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "cyclic prepend detected");
+  }
+}
+
+static mrb_value
+mrb_mod_prepend_features(mrb_state *mrb, mrb_value mod)
+{
+  mrb_value klass;
+
+  mrb_check_type(mrb, mod, MRB_TT_MODULE);
+  mrb_get_args(mrb, "C", &klass);
+  mrb_prepend_module(mrb, mrb_class_ptr(klass), mrb_class_ptr(mod));
+  return mod;
+}
+
+static mrb_value
+mrb_mod_prepend(mrb_state *mrb, mrb_value klass)
+{
+  mrb_value *argv;
+  mrb_int argc, i;
+
+  mrb_get_args(mrb, "*", &argv, &argc);
+  for (i=0; i<argc; i++) {
+    mrb_check_type(mrb, argv[i], MRB_TT_MODULE);
+  }
+  while (argc--) {
+    mrb_funcall(mrb, argv[argc], "prepend_features", 1, klass);
+    mrb_funcall(mrb, argv[argc], "prepended", 1, klass);
+  }
+
+  return klass;
 }
 
 static mrb_value
@@ -931,15 +1021,12 @@ mrb_mod_ancestors(mrb_state *mrb, mrb_value self)
 {
   mrb_value result;
   struct RClass *c = mrb_class_ptr(self);
-
   result = mrb_ary_new(mrb);
-  mrb_ary_push(mrb, result, mrb_obj_value(c));
-  c = c->super;
   while (c) {
     if (c->tt == MRB_TT_ICLASS) {
       mrb_ary_push(mrb, result, mrb_obj_value(c->c));
     }
-    else if (c->tt != MRB_TT_SCLASS) {
+    else if (c->origin == c) {
       mrb_ary_push(mrb, result, mrb_obj_value(c));
     }
     c = c->super;
@@ -964,11 +1051,14 @@ mrb_mod_included_modules(mrb_state *mrb, mrb_value self)
 {
   mrb_value result;
   struct RClass *c = mrb_class_ptr(self);
+  struct RClass *origin = c->origin;
 
   result = mrb_ary_new(mrb);
   while (c) {
-    if (c->tt == MRB_TT_ICLASS) {
-      mrb_ary_push(mrb, result, mrb_obj_value(c->c));
+    if (c != origin && c->tt == MRB_TT_ICLASS) {
+      if (c->c->tt == MRB_TT_MODULE) {
+        mrb_ary_push(mrb, result, mrb_obj_value(c->c));
+      }
     }
     c = c->super;
   }
@@ -980,10 +1070,11 @@ static mrb_value
 mrb_mod_initialize(mrb_state *mrb, mrb_value mod)
 {
   mrb_value b;
-
-  mrb_get_args(mrb, "&", &b);
+  struct RClass *m = mrb_class_ptr(mod);
+  boot_initmod(mrb, m); // bootstrap a newly initialized module
+  mrb_get_args(mrb, "|&", &b);
   if (!mrb_nil_p(b)) {
-    mrb_yield_with_class(mrb, b, 1, &mod, mod, mrb_class_ptr(mod));
+    mrb_yield_with_class(mrb, b, 1, &mod, mod, m);
   }
   return mod;
 }
@@ -1300,9 +1391,9 @@ mrb_class_superclass(mrb_state *mrb, mrb_value klass)
   struct RClass *c;
 
   c = mrb_class_ptr(klass);
-  c = c->super;
+  c = c->origin->super;
   while (c && c->tt == MRB_TT_ICLASS) {
-    c = c->super;
+    c = c->origin->super;
   }
   if (!c) return mrb_nil_value();
   return mrb_obj_value(c);
@@ -1540,8 +1631,7 @@ MRB_API struct RClass*
 mrb_module_new(mrb_state *mrb)
 {
   struct RClass *m = (struct RClass*)mrb_obj_alloc(mrb, MRB_TT_MODULE, mrb->module_class);
-  m->mt = kh_init(mt, mrb);
-
+  boot_initmod(mrb, m);
   return m;
 }
 
@@ -1900,13 +1990,14 @@ static void
 remove_method(mrb_state *mrb, mrb_value mod, mrb_sym mid)
 {
   struct RClass *c = mrb_class_ptr(mod);
-  khash_t(mt) *h = c->mt;
+  khash_t(mt) *h = c->origin->mt;
   khiter_t k;
 
   if (h) {
     k = kh_get(mt, mrb, h, mid);
     if (k != kh_end(h)) {
       kh_del(mt, mrb, h, k);
+      mrb_funcall(mrb, mod, "method_removed", 1, mrb_symbol_value(mid));
       return;
     }
   }
@@ -2140,6 +2231,9 @@ mrb_init_class(mrb_state *mrb)
   mrb_define_method(mrb, mod, "class_variable_set",      mrb_mod_cvar_set,         MRB_ARGS_REQ(2)); /* 15.2.2.4.18 */
   mrb_define_method(mrb, mod, "extend_object",           mrb_mod_extend_object,    MRB_ARGS_REQ(1)); /* 15.2.2.4.25 */
   mrb_define_method(mrb, mod, "extended",                mrb_bob_init,             MRB_ARGS_REQ(1)); /* 15.2.2.4.26 */
+  mrb_define_method(mrb, mod, "prepend",                 mrb_mod_prepend,          MRB_ARGS_ANY());
+  mrb_define_method(mrb, mod, "prepended",               mrb_bob_init,             MRB_ARGS_REQ(1));
+  mrb_define_method(mrb, mod, "prepend_features",        mrb_mod_prepend_features, MRB_ARGS_REQ(1));
   mrb_define_method(mrb, mod, "include",                 mrb_mod_include,          MRB_ARGS_ANY());  /* 15.2.2.4.27 */
   mrb_define_method(mrb, mod, "include?",                mrb_mod_include_p,        MRB_ARGS_REQ(1)); /* 15.2.2.4.28 */
   mrb_define_method(mrb, mod, "append_features",         mrb_mod_append_features,  MRB_ARGS_REQ(1)); /* 15.2.2.4.10 */
@@ -2156,6 +2250,7 @@ mrb_init_class(mrb_state *mrb)
   mrb_define_method(mrb, mod, "public",                  mrb_mod_dummy_visibility, MRB_ARGS_ANY());  /* 15.2.2.4.38 */
   mrb_define_method(mrb, mod, "remove_class_variable",   mrb_mod_remove_cvar,      MRB_ARGS_REQ(1)); /* 15.2.2.4.39 */
   mrb_define_method(mrb, mod, "remove_method",           mrb_mod_remove_method,    MRB_ARGS_ANY());  /* 15.2.2.4.41 */
+  mrb_define_method(mrb, mod, "method_removed",          mrb_bob_init,             MRB_ARGS_REQ(1));
   mrb_define_method(mrb, mod, "attr_reader",             mrb_mod_attr_reader,      MRB_ARGS_ANY());  /* 15.2.2.4.13 */
   mrb_define_method(mrb, mod, "attr_writer",             mrb_mod_attr_writer,      MRB_ARGS_ANY());  /* 15.2.2.4.14 */
   mrb_define_method(mrb, mod, "to_s",                    mrb_mod_to_s,             MRB_ARGS_NONE());
diff --git a/src/gc.c b/src/gc.c
index 8bd8243f1..15e1bd423 100644
--- a/src/gc.c
+++ b/src/gc.c
@@ -498,7 +498,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 +629,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;
diff --git a/src/kernel.c b/src/kernel.c
index b5b13f874..36ad683ee 100644
--- a/src/kernel.c
+++ b/src/kernel.c
@@ -248,6 +248,11 @@ mrb_singleton_class_clone(mrb_state *mrb, mrb_value obj)
       clone->c = mrb_singleton_class_clone(mrb, mrb_obj_value(klass));
     }
 
+    if (klass->origin != klass)
+      clone->origin = klass->origin;
+    else
+      clone->origin = clone;
+
     clone->super = klass->super;
     if (klass->iv) {
       mrb_iv_copy(mrb, mrb_obj_value(clone), mrb_obj_value(klass));
@@ -269,6 +274,13 @@ copy_class(mrb_state *mrb, mrb_value dst, mrb_value src)
 {
   struct RClass *dc = mrb_class_ptr(dst);
   struct RClass *sc = mrb_class_ptr(src);
+  /* if the origin is not the same as the class, then the origin and
+     the current class need to be copied */
+  if (sc->origin != sc) {
+    dc->origin = mrb_class_ptr(mrb_obj_dup(mrb, mrb_obj_value(sc->origin)));
+  } else {
+    dc->origin = dc;
+  }
   dc->mt = kh_copy(mt, mrb, sc->mt);
   dc->super = sc->super;
 }
@@ -641,13 +653,19 @@ mrb_class_instance_method_list(mrb_state *mrb, mrb_bool recur, struct RClass* kl
 {
   khint_t i;
   mrb_value ary;
+  mrb_bool prepended;
   struct RClass* oldklass;
   khash_t(st)* set = kh_init(st, mrb);
 
+  if (!recur && klass->origin != klass) {
+    klass = klass->origin;
+    prepended = 1;
+  }
+
   oldklass = 0;
   while (klass && (klass != oldklass)) {
     method_entry_loop(mrb, klass, set);
-    if ((klass->tt == MRB_TT_ICLASS) ||
+    if ((klass->tt == MRB_TT_ICLASS && !prepended) ||
         (klass->tt == MRB_TT_SCLASS)) {
     }
     else {
diff --git a/src/object.c b/src/object.c
index f8f41bfe8..c834ee04f 100644
--- a/src/object.c
+++ b/src/object.c
@@ -487,6 +487,7 @@ mrb_obj_is_kind_of(mrb_state *mrb, mrb_value obj, struct RClass *c)
       mrb_raise(mrb, E_TYPE_ERROR, "class or module required");
   }
 
+  c = c->origin;
   while (cl) {
     if (cl == c || cl->mt == c->mt)
       return TRUE;