call kh_get before kh_put to avoid potential key lost; reported by @mirichi

2 files changed, 12 insertions, 3 deletions

diff --git a/src/class.c b/src/class.c
index 656120845..4c39aed9a 100644
--- a/src/class.c
+++ b/src/class.c
@@ -313,7 +313,10 @@ mrb_define_method_raw(mrb_state *mrb, struct RClass *c, mrb_sym mid, struct RPro
   khiter_t k;
 
   if (!h) h = c->mt = kh_init(mt, mrb);
-  k = kh_put(mt, mrb, h, mid);
+  k = kh_get(mt, mrb, h, mid);
+  if (k == kh_end(h)) {
+    k = kh_put(mt, mrb, h, mid);
+  }
   kh_value(h, k) = p;
   if (p) {
     mrb_field_write_barrier(mrb, (struct RBasic *)c, (struct RBasic *)p);
@@ -346,7 +349,10 @@ mrb_define_method_vm(mrb_state *mrb, struct RClass *c, mrb_sym name, mrb_value b
   struct RProc *p;
 
   if (!h) h = c->mt = kh_init(mt, mrb);
-  k = kh_put(mt, mrb, h, name);
+  k = kh_get(mt, mrb, h, name);
+  if (k == kh_end(h)) {
+    k = kh_put(mt, mrb, h, name);
+  }
   p = mrb_proc_ptr(body);
   kh_value(h, k) = p;
   if (p) {
diff --git a/src/variable.c b/src/variable.c
index 3dcdbf775..6503c8fb7 100644
--- a/src/variable.c
+++ b/src/variable.c
@@ -308,7 +308,10 @@ iv_put(mrb_state *mrb, iv_tbl *t, mrb_sym sym, mrb_value val)
   khash_t(iv) *h = &t->h;
   khiter_t k;
 
-  k = kh_put(iv, mrb, h, sym);
+  k = kh_get(mt, mrb, h, sym);
+  if (k == kh_end(h)) {
+    k = kh_put(mt, mrb, h, sym);
+  }
   kh_value(h, k) = val;
 }