1 files changed, 103 insertions, 93 deletions

diff --git a/enumerator.c b/enumerator.c
index b33c171371..193a865dbc 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -85,12 +85,16 @@
  *   puts e.next   # => 3
  *   puts e.next   # raises StopIteration
  *
- * +next+, +next_values+, +peek+ and +peek_values+ are the only methods
- * which use external iteration (and Array#zip(Enumerable-not-Array) which uses +next+).
+ * +next+, +next_values+, +peek+, and +peek_values+ are the only methods
+ * which use external iteration (and Array#zip(Enumerable-not-Array) which uses +next+ internally).
  *
  * These methods do not affect other internal enumeration methods,
  * unless the underlying iteration method itself has side-effect, e.g. IO#each_line.
  *
+ * FrozenError will be raised if these methods are called against a frozen enumerator.
+ * Since +rewind+ and +feed+ also change state for external iteration,
+ * these methods may raise FrozenError too.
+ *
  * External iteration differs *significantly* from internal iteration
  * due to using a Fiber:
  * - The Fiber adds some overhead compared to internal enumeration.
@@ -165,7 +169,10 @@ static VALUE sym_each, sym_cycle, sym_yield;
 
 static VALUE lazy_use_super_method;
 
+extern ID ruby_static_id_cause;
+
 #define id_call idCall
+#define id_cause ruby_static_id_cause
 #define id_each idEach
 #define id_eqq idEqq
 #define id_initialize idInitialize
@@ -188,17 +195,18 @@ struct enumerator {
     int kw_splat;
 };
 
-RUBY_REFERENCES_START(enumerator_refs)
-    REF_EDGE(enumerator, obj),
-    REF_EDGE(enumerator, args),
-    REF_EDGE(enumerator, fib),
-    REF_EDGE(enumerator, dst),
-    REF_EDGE(enumerator, lookahead),
-    REF_EDGE(enumerator, feedvalue),
-    REF_EDGE(enumerator, stop_exc),
-    REF_EDGE(enumerator, size),
-    REF_EDGE(enumerator, procs),
-RUBY_REFERENCES_END
+RUBY_REFERENCES(enumerator_refs) = {
+    RUBY_REF_EDGE(struct enumerator, obj),
+    RUBY_REF_EDGE(struct enumerator, args),
+    RUBY_REF_EDGE(struct enumerator, fib),
+    RUBY_REF_EDGE(struct enumerator, dst),
+    RUBY_REF_EDGE(struct enumerator, lookahead),
+    RUBY_REF_EDGE(struct enumerator, feedvalue),
+    RUBY_REF_EDGE(struct enumerator, stop_exc),
+    RUBY_REF_EDGE(struct enumerator, size),
+    RUBY_REF_EDGE(struct enumerator, procs),
+    RUBY_REF_END
+};
 
 static VALUE rb_cGenerator, rb_cYielder, rb_cEnumProducer;
 
@@ -249,23 +257,15 @@ struct enum_product {
 
 VALUE rb_cArithSeq;
 
-#define enumerator_free RUBY_TYPED_DEFAULT_FREE
-
-static size_t
-enumerator_memsize(const void *p)
-{
-    return sizeof(struct enumerator);
-}
-
 static const rb_data_type_t enumerator_data_type = {
     "enumerator",
     {
-        REFS_LIST_PTR(enumerator_refs),
-        enumerator_free,
-        enumerator_memsize,
+        RUBY_REFS_LIST_PTR(enumerator_refs),
+        RUBY_TYPED_DEFAULT_FREE,
+        NULL, // Nothing allocated externally, so don't need a memsize function
         NULL,
     },
-    0, NULL, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_DECL_MARKING
+    0, NULL, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_DECL_MARKING | RUBY_TYPED_EMBEDDABLE
 };
 
 static struct enumerator *
@@ -296,22 +296,15 @@ proc_entry_compact(void *p)
     ptr->memo = rb_gc_location(ptr->memo);
 }
 
-#define proc_entry_free RUBY_TYPED_DEFAULT_FREE
-
-static size_t
-proc_entry_memsize(const void *p)
-{
-    return p ? sizeof(struct proc_entry) : 0;
-}
-
 static const rb_data_type_t proc_entry_data_type = {
     "proc_entry",
     {
         proc_entry_mark,
-        proc_entry_free,
-        proc_entry_memsize,
+        RUBY_TYPED_DEFAULT_FREE,
+        NULL, // Nothing allocated externally, so don't need a memsize function
         proc_entry_compact,
     },
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 };
 
 static struct proc_entry *
@@ -396,7 +389,7 @@ obj_to_enum(int argc, VALUE *argv, VALUE obj)
     }
     enumerator = rb_enumeratorize_with_size(obj, meth, argc, argv, 0);
     if (rb_block_given_p()) {
-        enumerator_ptr(enumerator)->size = rb_block_proc();
+        RB_OBJ_WRITE(enumerator, &enumerator_ptr(enumerator)->size, rb_block_proc());
     }
     return enumerator;
 }
@@ -425,15 +418,15 @@ enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, const VALUE *ar
         rb_raise(rb_eArgError, "unallocated enumerator");
     }
 
-    ptr->obj  = obj;
+    RB_OBJ_WRITE(enum_obj, &ptr->obj, obj);
     ptr->meth = rb_to_id(meth);
-    if (argc) ptr->args = rb_ary_new4(argc, argv);
+    if (argc) RB_OBJ_WRITE(enum_obj, &ptr->args, rb_ary_new4(argc, argv));
     ptr->fib = 0;
     ptr->dst = Qnil;
     ptr->lookahead = Qundef;
     ptr->feedvalue = Qundef;
     ptr->stop_exc = Qfalse;
-    ptr->size = size;
+    RB_OBJ_WRITE(enum_obj, &ptr->size, size);
     ptr->size_fn = size_fn;
     ptr->kw_splat = kw_splat;
 
@@ -512,13 +505,13 @@ enumerator_init_copy(VALUE obj, VALUE orig)
         rb_raise(rb_eArgError, "unallocated enumerator");
     }
 
-    ptr1->obj  = ptr0->obj;
-    ptr1->meth = ptr0->meth;
-    ptr1->args = ptr0->args;
+    RB_OBJ_WRITE(obj, &ptr1->obj, ptr0->obj);
+    RB_OBJ_WRITE(obj, &ptr1->meth, ptr0->meth);
+    RB_OBJ_WRITE(obj, &ptr1->args, ptr0->args);
     ptr1->fib  = 0;
     ptr1->lookahead  = Qundef;
     ptr1->feedvalue  = Qundef;
-    ptr1->size  = ptr0->size;
+    RB_OBJ_WRITE(obj, &ptr1->size, ptr0->size);
     ptr1->size_fn  = ptr0->size_fn;
 
     return obj;
@@ -566,11 +559,17 @@ enumerator_block_call(VALUE obj, rb_block_call_func *func, VALUE arg)
     const struct enumerator *e = enumerator_ptr(obj);
     ID meth = e->meth;
 
-    if (e->args) {
-        argc = RARRAY_LENINT(e->args);
-        argv = RARRAY_CONST_PTR(e->args);
+    VALUE args = e->args;
+    if (args) {
+        argc = RARRAY_LENINT(args);
+        argv = RARRAY_CONST_PTR(args);
     }
-    return rb_block_call_kw(e->obj, meth, argc, argv, func, arg, e->kw_splat);
+
+    VALUE ret = rb_block_call_kw(e->obj, meth, argc, argv, func, arg, e->kw_splat);
+
+    RB_GC_GUARD(args);
+
+    return ret;
 }
 
 /*
@@ -627,7 +626,7 @@ enumerator_each(int argc, VALUE *argv, VALUE obj)
         else {
             args = rb_ary_new4(argc, argv);
         }
-        e->args = args;
+        RB_OBJ_WRITE(obj, &e->args, args);
         e->size = Qnil;
         e->size_fn = 0;
     }
@@ -768,7 +767,7 @@ next_i(RB_BLOCK_CALL_FUNC_ARGLIST(_, obj))
     VALUE result;
 
     result = rb_block_call(obj, id_each, 0, 0, next_ii, obj);
-    e->stop_exc = rb_exc_new2(rb_eStopIteration, "iteration reached an end");
+    RB_OBJ_WRITE(obj, &e->stop_exc, rb_exc_new2(rb_eStopIteration, "iteration reached an end"));
     rb_ivar_set(e->stop_exc, id_result, result);
     return rb_fiber_yield(1, &nil);
 }
@@ -777,8 +776,8 @@ static void
 next_init(VALUE obj, struct enumerator *e)
 {
     VALUE curr = rb_fiber_current();
-    e->dst = curr;
-    e->fib = rb_fiber_new(next_i, obj);
+    RB_OBJ_WRITE(obj, &e->dst, curr);
+    RB_OBJ_WRITE(obj, &e->fib, rb_fiber_new(next_i, obj));
     e->lookahead = Qundef;
 }
 
@@ -787,8 +786,16 @@ get_next_values(VALUE obj, struct enumerator *e)
 {
     VALUE curr, vs;
 
-    if (e->stop_exc)
-        rb_exc_raise(e->stop_exc);
+    if (e->stop_exc) {
+        VALUE exc = e->stop_exc;
+        VALUE result = rb_attr_get(exc, id_result);
+        VALUE mesg = rb_attr_get(exc, idMesg);
+        if (!NIL_P(mesg)) mesg = rb_str_dup(mesg);
+        VALUE stop_exc = rb_exc_new_str(rb_eStopIteration, mesg);
+        rb_ivar_set(stop_exc, id_cause, exc);
+        rb_ivar_set(stop_exc, id_result, result);
+        rb_exc_raise(stop_exc);
+    }
 
     curr = rb_fiber_current();
 
@@ -858,6 +865,8 @@ enumerator_next_values(VALUE obj)
     struct enumerator *e = enumerator_ptr(obj);
     VALUE vs;
 
+    rb_check_frozen(obj);
+
     if (!UNDEF_P(e->lookahead)) {
         vs = e->lookahead;
         e->lookahead = Qundef;
@@ -919,9 +928,12 @@ enumerator_peek_values(VALUE obj)
 {
     struct enumerator *e = enumerator_ptr(obj);
 
+    rb_check_frozen(obj);
+
     if (UNDEF_P(e->lookahead)) {
-        e->lookahead = get_next_values(obj, e);
+        RB_OBJ_WRITE(obj, &e->lookahead, get_next_values(obj, e));
     }
+
     return e->lookahead;
 }
 
@@ -1043,10 +1055,12 @@ enumerator_feed(VALUE obj, VALUE v)
 {
     struct enumerator *e = enumerator_ptr(obj);
 
+    rb_check_frozen(obj);
+
     if (!UNDEF_P(e->feedvalue)) {
         rb_raise(rb_eTypeError, "feed value already set");
     }
-    e->feedvalue = v;
+    RB_OBJ_WRITE(obj, &e->feedvalue, v);
 
     return Qnil;
 }
@@ -1065,6 +1079,8 @@ enumerator_rewind(VALUE obj)
 {
     struct enumerator *e = enumerator_ptr(obj);
 
+    rb_check_frozen(obj);
+
     rb_check_funcall(e->obj, id_rewind, 0, 0);
 
     e->fib = 0;
@@ -1278,23 +1294,15 @@ yielder_compact(void *p)
     ptr->proc = rb_gc_location(ptr->proc);
 }
 
-#define yielder_free RUBY_TYPED_DEFAULT_FREE
-
-static size_t
-yielder_memsize(const void *p)
-{
-    return sizeof(struct yielder);
-}
-
 static const rb_data_type_t yielder_data_type = {
     "yielder",
     {
         yielder_mark,
-        yielder_free,
-        yielder_memsize,
+        RUBY_TYPED_DEFAULT_FREE,
+        NULL,
         yielder_compact,
     },
-    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 };
 
 static struct yielder *
@@ -1333,7 +1341,7 @@ yielder_init(VALUE obj, VALUE proc)
         rb_raise(rb_eArgError, "unallocated yielder");
     }
 
-    ptr->proc = proc;
+    RB_OBJ_WRITE(obj, &ptr->proc, proc);
 
     return obj;
 }
@@ -1418,23 +1426,15 @@ generator_compact(void *p)
     ptr->obj = rb_gc_location(ptr->obj);
 }
 
-#define generator_free RUBY_TYPED_DEFAULT_FREE
-
-static size_t
-generator_memsize(const void *p)
-{
-    return sizeof(struct generator);
-}
-
 static const rb_data_type_t generator_data_type = {
     "generator",
     {
         generator_mark,
-        generator_free,
-        generator_memsize,
+        RUBY_TYPED_DEFAULT_FREE,
+        NULL,
         generator_compact,
     },
-    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 };
 
 static struct generator *
@@ -1474,7 +1474,7 @@ generator_init(VALUE obj, VALUE proc)
         rb_raise(rb_eArgError, "unallocated generator");
     }
 
-    ptr->proc = proc;
+    RB_OBJ_WRITE(obj, &ptr->proc, proc);
 
     return obj;
 }
@@ -1522,7 +1522,7 @@ generator_init_copy(VALUE obj, VALUE orig)
         rb_raise(rb_eArgError, "unallocated generator");
     }
 
-    ptr1->proc = ptr0->proc;
+    RB_OBJ_WRITE(obj, &ptr1->proc, ptr0->proc);
 
     return obj;
 }
@@ -1689,7 +1689,7 @@ lazy_generator_init(VALUE enumerator, VALUE procs)
                   lazy_init_block, rb_ary_new3(2, obj, procs));
 
     gen_ptr = generator_ptr(generator);
-    gen_ptr->obj = obj;
+    RB_OBJ_WRITE(generator, &gen_ptr->obj, obj);
 
     return generator;
 }
@@ -1874,10 +1874,10 @@ lazy_add_method(VALUE obj, int argc, VALUE *argv, VALUE args, VALUE memo,
     VALUE entry_obj = TypedData_Make_Struct(rb_cObject, struct proc_entry,
                                             &proc_entry_data_type, entry);
     if (rb_block_given_p()) {
-        entry->proc = rb_block_proc();
+        RB_OBJ_WRITE(entry_obj, &entry->proc, rb_block_proc());
     }
     entry->fn = fn;
-    entry->memo = args;
+    RB_OBJ_WRITE(entry_obj, &entry->memo, args);
 
     lazy_set_args(entry_obj, memo);
 
@@ -1886,9 +1886,9 @@ lazy_add_method(VALUE obj, int argc, VALUE *argv, VALUE args, VALUE memo,
     rb_ary_push(new_procs, entry_obj);
 
     new_obj = enumerator_init_copy(enumerator_allocate(rb_cLazy), obj);
-    new_e = DATA_PTR(new_obj);
-    new_e->obj = new_generator;
-    new_e->procs = new_procs;
+    new_e = RTYPEDDATA_GET_DATA(new_obj);
+    RB_OBJ_WRITE(new_obj, &new_e->obj, new_generator);
+    RB_OBJ_WRITE(new_obj, &new_e->procs, new_procs);
 
     if (argc > 0) {
         new_e->meth = rb_to_id(*argv++);
@@ -1897,7 +1897,9 @@ lazy_add_method(VALUE obj, int argc, VALUE *argv, VALUE args, VALUE memo,
     else {
         new_e->meth = id_each;
     }
-    new_e->args = rb_ary_new4(argc, argv);
+
+    RB_OBJ_WRITE(new_obj, &new_e->args, rb_ary_new4(argc, argv));
+
     return new_obj;
 }
 
@@ -1983,7 +1985,7 @@ lazy_to_enum(int argc, VALUE *argv, VALUE self)
     }
     lazy = lazy_to_enum_i(self, meth, argc, argv, 0, rb_keyword_given_p());
     if (rb_block_given_p()) {
-        enumerator_ptr(lazy)->size = rb_block_proc();
+        RB_OBJ_WRITE(lazy, &enumerator_ptr(lazy)->size, rb_block_proc());
     }
     return lazy;
 }
@@ -2369,7 +2371,6 @@ lazy_zip_arrays_func(VALUE proc_entry, struct MEMO *result, VALUE memos, long me
         rb_ary_push(ary, rb_ary_entry(RARRAY_AREF(arrays, i), count));
     }
     LAZY_MEMO_SET_VALUE(result, ary);
-    LAZY_MEMO_SET_PACKED(result);
     rb_ary_store(memos, memo_index, LONG2NUM(++count));
     return result;
 }
@@ -2944,7 +2945,7 @@ static const rb_data_type_t producer_data_type = {
         producer_memsize,
         producer_compact,
     },
-    0, 0, RUBY_TYPED_FREE_IMMEDIATELY
+    0, 0, RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE
 };
 
 static struct producer *
@@ -2984,8 +2985,8 @@ producer_init(VALUE obj, VALUE init, VALUE proc)
         rb_raise(rb_eArgError, "unallocated producer");
     }
 
-    ptr->init = init;
-    ptr->proc = proc;
+    RB_OBJ_WRITE(obj, &ptr->init, init);
+    RB_OBJ_WRITE(obj, &ptr->proc, proc);
 
     return obj;
 }
@@ -3535,10 +3536,19 @@ static VALUE
 enum_product_total_size(VALUE enums)
 {
     VALUE total = INT2FIX(1);
+    VALUE sizes = rb_ary_hidden_new(RARRAY_LEN(enums));
     long i;
 
     for (i = 0; i < RARRAY_LEN(enums); i++) {
         VALUE size = enum_size(RARRAY_AREF(enums, i));
+        if (size == INT2FIX(0)) {
+            rb_ary_resize(sizes, 0);
+            return size;
+        }
+        rb_ary_push(sizes, size);
+    }
+    for (i = 0; i < RARRAY_LEN(sizes); i++) {
+        VALUE size = RARRAY_AREF(sizes, i);
 
         if (NIL_P(size) || (RB_TYPE_P(size, T_FLOAT) && isinf(NUM2DBL(size)))) {
             return size;
@@ -4561,7 +4571,7 @@ InitVM_Enumerator(void)
     rb_hash_aset(lazy_use_super_method, sym("uniq"), sym("_enumerable_uniq"));
     rb_hash_aset(lazy_use_super_method, sym("with_index"), sym("_enumerable_with_index"));
     rb_obj_freeze(lazy_use_super_method);
-    rb_gc_register_mark_object(lazy_use_super_method);
+    rb_vm_register_global_object(lazy_use_super_method);
 
 #if 0 /* for RDoc */
     rb_define_method(rb_cLazy, "to_a", lazy_to_a, 0);