Merge pull request #5951 from dearblue/mrb_yield

Safely retrieve `self` from proc in `mrb_yield()`

include/mruby/internal.h

@@ -117,6 +117,7 @@ struct REnv *mrb_env_new(mrb_state *mrb, struct mrb_context *c, mrb_callinfo *ci
 void mrb_proc_merge_lvar(mrb_state *mrb, mrb_irep *irep, struct REnv *env, int num, const mrb_sym *lv, const mrb_value *stack);
 mrb_value mrb_proc_local_variables(mrb_state *mrb, const struct RProc *proc);
 const struct RProc *mrb_proc_get_caller(mrb_state *mrb, struct REnv **env);
+mrb_value mrb_proc_get_self(mrb_state *mrb, struct RProc *p, struct RClass **target_class_p);
 #endif
 
 /* range */

src/proc.c

@@ -193,6 +193,29 @@ mrb_proc_cfunc_env_get(mrb_state *mrb, mrb_int idx)
   return e->stack[idx];
 }
 
+mrb_value
+mrb_proc_get_self(mrb_state *mrb, struct RProc *p, struct RClass **target_class_p)
+{
+  if (MRB_PROC_CFUNC_P(p)) {
+    *target_class_p = mrb->object_class;
+    return mrb_nil_value();
+  }
+  else {
+    struct REnv *e = p->e.env;
+
+    if (!e || e->tt != MRB_TT_ENV) {
+      *target_class_p = mrb->object_class;
+      return mrb_top_self(mrb);
+    }
+    else if (MRB_ENV_LEN(e) < 1) {
+      mrb_raise(mrb, E_ARGUMENT_ERROR, "self is lost (probably ran out of memory when the block became independent)");
+    }
+
+    *target_class_p = e->c;
+    return e->stack[0];
+  }
+}
+
 void
 mrb_proc_copy(mrb_state *mrb, struct RProc *a, struct RProc *b)
 {

src/vm.c

@@ -1004,16 +1004,20 @@ MRB_API mrb_value
 mrb_yield_argv(mrb_state *mrb, mrb_value b, mrb_int argc, const mrb_value *argv)
 {
   struct RProc *p = mrb_proc_ptr(b);
+  struct RClass *tc;
+  mrb_value self = mrb_proc_get_self(mrb, p, &tc);
 
-  return mrb_yield_with_class(mrb, b, argc, argv, MRB_PROC_ENV(p)->stack[0], MRB_PROC_TARGET_CLASS(p));
+  return mrb_yield_with_class(mrb, b, argc, argv, self, tc);
 }
 
 MRB_API mrb_value
 mrb_yield(mrb_state *mrb, mrb_value b, mrb_value arg)
 {
   struct RProc *p = mrb_proc_ptr(b);
+  struct RClass *tc;
+  mrb_value self = mrb_proc_get_self(mrb, p, &tc);
 
-  return mrb_yield_with_class(mrb, b, 1, &arg, MRB_PROC_ENV(p)->stack[0], MRB_PROC_TARGET_CLASS(p));
+  return mrb_yield_with_class(mrb, b, 1, &arg, self, tc);
 }
 
 mrb_value