3 files changed, 13 insertions, 16 deletions

diff --git a/include/mruby/class.h b/include/mruby/class.h
index c88079e59..b97bc3991 100644
--- a/include/mruby/class.h
+++ b/include/mruby/class.h
@@ -100,8 +100,7 @@ void mrb_mc_clear_by_class(mrb_state *mrb, struct RClass* c);
 #endif
 
 /* return non zero to break the loop */
-struct mt_elem;
-typedef int (mrb_mt_foreach_func)(mrb_state*,mrb_sym,struct mt_elem*,void*);
+typedef int (mrb_mt_foreach_func)(mrb_state*,mrb_sym,mrb_method_t,void*);
 MRB_API void mrb_mt_foreach(mrb_state*, struct RClass*, mrb_mt_foreach_func*, void*);
 
 MRB_END_DECL
diff --git a/mrbgems/mruby-metaprog/src/metaprog.c b/mrbgems/mruby-metaprog/src/metaprog.c
index ce198c9b9..8bc9ab038 100644
--- a/mrbgems/mruby-metaprog/src/metaprog.c
+++ b/mrbgems/mruby-metaprog/src/metaprog.c
@@ -171,28 +171,17 @@ mrb_local_variables(mrb_state *mrb, mrb_value self)
 KHASH_DECLARE(st, mrb_sym, char, FALSE)
 KHASH_DEFINE(st, mrb_sym, char, FALSE, kh_int_hash_func, kh_int_hash_equal)
 
-union mt_ptr {
-  struct RProc *proc;
-  mrb_func_t func;
-};
-
-struct mt_elem {
-  union mt_ptr ptr;
-  size_t func_p:1;
-  mrb_sym key:sizeof(mrb_sym)*8-1;
-};
-
 struct mt_set {
   khash_t(st) *set;
   khash_t(st) *undef;
 };
 
 static int
-method_entry_i(mrb_state *mrb, mrb_sym mid, struct mt_elem *e, void *p)
+method_entry_i(mrb_state *mrb, mrb_sym mid, mrb_method_t m, void *p)
 {
   struct mt_set *s = (struct mt_set*)p;
 
-  if (e->ptr.proc == 0) {
+  if (MRB_METHOD_UNDEF_P(m)) {
     if (s->undef) {
       kh_put(st, mrb, s->undef, mid);
     }
diff --git a/src/class.c b/src/class.c
index f00d3c782..26366aeb9 100644
--- a/src/class.c
+++ b/src/class.c
@@ -230,7 +230,16 @@ mrb_mt_foreach(mrb_state *mrb, struct RClass *c, mrb_mt_foreach_func *fn, void *
     struct mt_elem *slot = &t->table[i];
 
     if (slot->key) {
-      if (fn(mrb, slot->key, slot, p) != 0)
+      mrb_method_t m;
+
+      if (slot->func_p) {
+        MRB_METHOD_FROM_FUNC(m, slot->ptr.func);
+      }
+      else {
+        MRB_METHOD_FROM_PROC(m, slot->ptr.proc);
+      }
+
+      if (fn(mrb, slot->key, m, p) != 0)
         return;
     }
   }