* array.c (ary_iter_check): should check modification (size

  change) during iteration.

* array.c (rb_ary_initialize, rb_ary_shift, rb_ary_unshift,
  rb_ary_splice, rb_ary_reverse, rb_ary_sort, rb_ary_delete,
  rb_ary_delete_at, rb_ary_reject_bang, rb_ary_replace,
  rb_ary_clear, rb_ary_fill, rb_ary_uniq_bang, rb_ary_compact,
  rb_ary_shuffle): add iteration check.

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@11429 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

1 files changed, 142 insertions, 60 deletions

diff --git a/array.c b/array.c
index 18d8f84580..a31c10de59 100644
--- a/array.c
+++ b/array.c
@@ -39,7 +39,15 @@ memfill(register VALUE *mem, register long size, register VALUE val)
     }
 }
 
-#define ARY_TMPLOCK FL_USER1
+#define ARY_ITERLOCK FL_USER1
+static void
+ary_iter_check(VALUE ary)
+{
+    if (FL_TEST(ary, ARY_ITERLOCK)) {
+	rb_raise(rb_eRuntimeError, "can't modify array during iteration");
+    }
+}
+#define ARY_SORTLOCK FL_USER3
 #define ARY_SHARED_P(a) FL_TEST(a, ELTS_SHARED)
 
 #define ARY_SET_LEN(ary, n) do { \
@@ -52,6 +60,11 @@ memfill(register VALUE *mem, register long size, register VALUE val)
     RARRAY(ary)->aux.capa = (capacity);\
 } while (0)
 
+#define ITERATE(func, ary) do { \
+    FL_SET(ary, ARY_ITERLOCK); \
+    return rb_ensure(func, (ary), each_unlock, (ary));\
+} while (0)
+
 static VALUE *
 rb_ary_ptr(VALUE ary)
 {
@@ -62,8 +75,8 @@ static inline void
 rb_ary_modify_check(VALUE ary)
 {
     if (OBJ_FROZEN(ary)) rb_error_frozen("array");
-    if (FL_TEST(ary, ARY_TMPLOCK))
-	rb_raise(rb_eRuntimeError, "can't modify array during iteration");
+    if (FL_TEST(ary, ARY_SORTLOCK))
+	rb_raise(rb_eRuntimeError, "can't modify array during sort");
     if (!OBJ_TAINTED(ary) && rb_safe_level() >= 4)
 	rb_raise(rb_eSecurityError, "Insecure: can't modify array");
 }
@@ -101,7 +114,7 @@ static VALUE
 rb_ary_frozen_p(VALUE ary)
 {
     if (OBJ_FROZEN(ary)) return Qtrue;
-    if (FL_TEST(ary, ARY_TMPLOCK)) return Qtrue;
+    if (FL_TEST(ary, ARY_SORTLOCK)) return Qtrue;
     return Qfalse;
 }
 
@@ -284,6 +297,7 @@ rb_ary_initialize(int argc, VALUE *argv, VALUE ary)
     VALUE size, val;
 
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (rb_scan_args(argc, argv, "02", &size, &val) == 0) {
 	if (RARRAY_PTR(ary) && !ARY_SHARED_P(ary)) {
 	    free(RARRAY(ary)->ptr);
@@ -524,6 +538,7 @@ rb_ary_shift(VALUE ary)
     VALUE top;
 
     rb_ary_modify_check(ary);
+    ary_iter_check(ary);
     if (RARRAY_LEN(ary) == 0) return Qnil;
     top = RARRAY_PTR(ary)[0];
     if (!ARY_SHARED_P(ary)) {
@@ -569,6 +584,7 @@ rb_ary_shift_m(int argc, VALUE *argv, VALUE ary)
     }
 
     rb_ary_modify_check(ary);
+    ary_iter_check(ary);
     result = ary_shared_first(argc, argv, ary, Qfalse);
     n = RARRAY_LEN(result);
     if (ARY_SHARED_P(ary)) {
@@ -602,6 +618,7 @@ rb_ary_unshift_m(int argc, VALUE *argv, VALUE ary)
 
     if (argc == 0) return ary;
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (RARRAY(ary)->aux.capa <= RARRAY_LEN(ary)+argc) {
 	RESIZE_CAPA(ary, RARRAY(ary)->aux.capa + ARY_DEFAULT_SIZE);
     }
@@ -914,8 +931,8 @@ rb_ary_rindex(int argc, VALUE *argv, VALUE ary)
     long i = RARRAY_LEN(ary);
 
     if (rb_scan_args(argc, argv, "01", &val) == 0) {
+	RETURN_ENUMERATOR(ary, 0, 0);
 	while (i--) {
-	    RETURN_ENUMERATOR(ary, 0, 0);
 	    if (RTEST(rb_yield(RARRAY_PTR(ary)[i])))
 		return LONG2NUM(i);
 	    if (i > RARRAY_LEN(ary)) {
@@ -972,6 +989,7 @@ rb_ary_splice(VALUE ary, long beg, long len, VALUE rpl)
 	rlen = RARRAY_LEN(rpl);
     }
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (beg >= RARRAY_LEN(ary)) {
 	len = beg + rlen;
 	if (len >= ARY_CAPA(ary)) {
@@ -1094,6 +1112,24 @@ rb_ary_insert(int argc, VALUE *argv, VALUE ary)
     return ary;
 }
 
+static VALUE
+each_unlock(VALUE ary)
+{
+    FL_UNSET(ary, ARY_ITERLOCK);
+    return ary;
+}
+
+static VALUE
+each_i(VALUE ary)
+{
+    long i;
+
+    for (i=0; i<RARRAY_LEN(ary); i++) {
+	rb_yield(RARRAY_PTR(ary)[i]);
+    }
+    return ary;
+}
+
 /*
  *  call-seq:
  *     array.each {|item| block }   ->   array
@@ -1112,11 +1148,17 @@ rb_ary_insert(int argc, VALUE *argv, VALUE ary)
 VALUE
 rb_ary_each(VALUE ary)
 {
+    RETURN_ENUMERATOR(ary, 0, 0);
+    ITERATE(each_i, ary);
+}
+
+static VALUE
+each_index_i(VALUE ary)
+{
     long i;
 
-    RETURN_ENUMERATOR(ary, 0, 0);
     for (i=0; i<RARRAY_LEN(ary); i++) {
-	rb_yield(RARRAY_PTR(ary)[i]);
+	rb_yield(LONG2NUM(i));
     }
     return ary;
 }
@@ -1139,11 +1181,20 @@ rb_ary_each(VALUE ary)
 static VALUE
 rb_ary_each_index(VALUE ary)
 {
-    long i;
-
     RETURN_ENUMERATOR(ary, 0, 0);
-    for (i=0; i<RARRAY_LEN(ary); i++) {
-	rb_yield(LONG2NUM(i));
+    ITERATE(each_index_i, ary);
+}
+
+static VALUE
+reverse_each_i(VALUE ary)
+{
+    long len = RARRAY_LEN(ary);
+
+    while (len--) {
+	rb_yield(RARRAY_PTR(ary)[len]);
+	if (RARRAY_LEN(ary) < len) {
+	    len = RARRAY_LEN(ary);
+	}
     }
     return ary;
 }
@@ -1166,16 +1217,8 @@ rb_ary_each_index(VALUE ary)
 static VALUE
 rb_ary_reverse_each(VALUE ary)
 {
-    long len = RARRAY_LEN(ary);
-
     RETURN_ENUMERATOR(ary, 0, 0);
-    while (len--) {
-	rb_yield(RARRAY_PTR(ary)[len]);
-	if (RARRAY_LEN(ary) < len) {
-	    len = RARRAY_LEN(ary);
-	}
-    }
-    return ary;
+    ITERATE(reverse_each_i, ary);
 }
 
 /*
@@ -1383,6 +1426,7 @@ rb_ary_reverse(VALUE ary)
     VALUE tmp;
 
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (RARRAY_LEN(ary) > 1) {
 	p1 = RARRAY_PTR(ary);
 	p2 = p1 + RARRAY_LEN(ary) - 1;	/* points last item */
@@ -1479,7 +1523,7 @@ sort_2(const void *ap, const void *bp, void *data)
 }
 
 static VALUE
-sort_internal(VALUE ary)
+sort_i(VALUE ary)
 {
     struct ary_sort_data data;
 
@@ -1493,7 +1537,7 @@ sort_internal(VALUE ary)
 static VALUE
 sort_unlock(VALUE ary)
 {
-    FL_UNSET(ary, ARY_TMPLOCK);
+    FL_UNSET(ary, ARY_SORTLOCK);
     return ary;
 }
 
@@ -1517,9 +1561,10 @@ VALUE
 rb_ary_sort_bang(VALUE ary)
 {
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (RARRAY_LEN(ary) > 1) {
-	FL_SET(ary, ARY_TMPLOCK);	/* prohibit modification during sort */
-	rb_ensure(sort_internal, ary, sort_unlock, ary);
+	FL_SET(ary, ARY_SORTLOCK);	/* prohibit modification during sort */
+	rb_ensure(sort_i, ary, sort_unlock, ary);
     }
     return ary;
 }
@@ -1548,6 +1593,20 @@ rb_ary_sort(VALUE ary)
     return ary;
 }
 
+
+static VALUE
+collect_i(VALUE ary)
+{
+    long i;
+    VALUE collect;
+
+    collect = rb_ary_new2(RARRAY_LEN(ary));
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+	rb_ary_push(collect, rb_yield(RARRAY_PTR(ary)[i]));
+    }
+    return collect;
+}
+
 /*
  *  call-seq:
  *     array.collect {|item| block }  -> an_array
@@ -1565,15 +1624,21 @@ rb_ary_sort(VALUE ary)
 static VALUE
 rb_ary_collect(VALUE ary)
 {
+    RETURN_ENUMERATOR(ary, 0, 0);
+    ITERATE(collect_i, ary);
+}
+
+
+static VALUE
+collect_bang_i(VALUE ary)
+{
     long i;
-    VALUE collect;
 
-    RETURN_ENUMERATOR(ary, 0, 0);
-    collect = rb_ary_new2(RARRAY_LEN(ary));
+    rb_ary_modify(ary);
     for (i = 0; i < RARRAY_LEN(ary); i++) {
-	rb_ary_push(collect, rb_yield(RARRAY_PTR(ary)[i]));
+	RARRAY_PTR(ary)[i] = rb_yield(RARRAY_PTR(ary)[i]);
     }
-    return collect;
+    return ary;
 }
 
 /* 
@@ -1593,14 +1658,8 @@ rb_ary_collect(VALUE ary)
 static VALUE
 rb_ary_collect_bang(VALUE ary)
 {
-    long i;
-
     RETURN_ENUMERATOR(ary, 0, 0);
-    rb_ary_modify(ary);
-    for (i = 0; i < RARRAY_LEN(ary); i++) {
-	rb_ary_store(ary, i, rb_yield(RARRAY_PTR(ary)[i]));
-    }
-    return ary;
+    ITERATE(collect_bang_i, ary);
 }
 
 VALUE
@@ -1653,6 +1712,22 @@ rb_ary_values_at(int argc, VALUE *argv, VALUE ary)
     return rb_get_values_at(ary, RARRAY_LEN(ary), argc, argv, rb_ary_entry);
 }
 
+
+static VALUE
+select_i(VALUE ary)
+{
+    VALUE result;
+    long i;
+
+    result = rb_ary_new2(RARRAY_LEN(ary));
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+	if (RTEST(rb_yield(RARRAY_PTR(ary)[i]))) {
+	    rb_ary_push(result, rb_ary_elt(ary, i));
+	}
+    }
+    return result;
+}
+
 /*
  *  call-seq:
  *     array.select {|item| block } -> an_array
@@ -1668,17 +1743,8 @@ rb_ary_values_at(int argc, VALUE *argv, VALUE ary)
 static VALUE
 rb_ary_select(VALUE ary)
 {
-    VALUE result;
-    long i;
-
     RETURN_ENUMERATOR(ary, 0, 0);
-    result = rb_ary_new2(RARRAY_LEN(ary));
-    for (i = 0; i < RARRAY_LEN(ary); i++) {
-	if (RTEST(rb_yield(RARRAY_PTR(ary)[i]))) {
-	    rb_ary_push(result, rb_ary_elt(ary, i));
-	}
-    }
-    return result;
+    ITERATE(select_i, ary);
 }
 
 /*
@@ -1720,6 +1786,7 @@ rb_ary_delete(VALUE ary, VALUE item)
     }
 
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     if (RARRAY_LEN(ary) > i2) {
 	RARRAY(ary)->len = i2;
 	if (i2 * 2 < ARY_CAPA(ary) &&
@@ -1744,6 +1811,7 @@ rb_ary_delete_at(VALUE ary, long pos)
     }
 
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     del = RARRAY_PTR(ary)[pos];
     MEMMOVE(RARRAY_PTR(ary)+pos, RARRAY_PTR(ary)+pos+1, VALUE,
 	    RARRAY_LEN(ary)-pos-1);
@@ -1822,22 +1890,12 @@ rb_ary_slice_bang(int argc, VALUE *argv, VALUE ary)
     return rb_ary_delete_at(ary, NUM2LONG(arg1));
 }
 
-/*
- *  call-seq:
- *     array.reject! {|item| block }  -> array or nil
- *  
- *  Equivalent to <code>Array#delete_if</code>, deleting elements from
- *  _self_ for which the block evaluates to true, but returns
- *  <code>nil</code> if no changes were made. Also see
- *  <code>Enumerable#reject</code>.
- */
 
 static VALUE
-rb_ary_reject_bang(VALUE ary)
+reject_bang_i(VALUE ary)
 {
     long i1, i2;
 
-    RETURN_ENUMERATOR(ary, 0, 0);
     rb_ary_modify(ary);
     for (i1 = i2 = 0; i1 < RARRAY_LEN(ary); i1++) {
 	VALUE v = RARRAY_PTR(ary)[i1];
@@ -1856,6 +1914,24 @@ rb_ary_reject_bang(VALUE ary)
 
 /*
  *  call-seq:
+ *     array.reject! {|item| block }  -> array or nil
+ *  
+ *  Equivalent to <code>Array#delete_if</code>, deleting elements from
+ *  _self_ for which the block evaluates to true, but returns
+ *  <code>nil</code> if no changes were made. Also see
+ *  <code>Enumerable#reject</code>.
+ */
+
+static VALUE
+rb_ary_reject_bang(VALUE ary)
+{
+    RETURN_ENUMERATOR(ary, 0, 0);
+    ary_iter_check(ary);
+    ITERATE(reject_bang_i, ary);
+}
+
+/*
+ *  call-seq:
  *     array.reject {|item| block }  -> an_array
  *  
  *  Returns a new array containing the items in _self_
@@ -1866,9 +1942,7 @@ static VALUE
 rb_ary_reject(VALUE ary)
 {
     RETURN_ENUMERATOR(ary, 0, 0);
-    ary = rb_ary_dup(ary);
-    rb_ary_reject_bang(ary);
-    return ary;
+    return rb_ary_reject_bang(rb_ary_dup(ary));
 }
 
 /*
@@ -1918,6 +1992,7 @@ rb_ary_zip(int argc, VALUE *argv, VALUE ary)
     long len;
     VALUE result;
 
+    RETURN_ENUMERATOR(ary, argc, argv);
     for (i=0; i<argc; i++) {
 	argv[i] = to_a(argv[i]);
     }
@@ -2005,6 +2080,7 @@ rb_ary_replace(VALUE copy, VALUE orig)
     VALUE *ptr;
 
     rb_ary_modify(copy);
+    ary_iter_check(copy);
     orig = to_ary(orig);
     if (copy == orig) return copy;
     shared = ary_make_shared(orig);
@@ -2032,6 +2108,7 @@ VALUE
 rb_ary_clear(VALUE ary)
 {
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     RARRAY(ary)->len = 0;
     if (ARY_DEFAULT_SIZE * 2 < ARY_CAPA(ary)) {
 	RESIZE_CAPA(ary, ARY_DEFAULT_SIZE * 2);
@@ -2099,6 +2176,7 @@ rb_ary_fill(int argc, VALUE *argv, VALUE ary)
 	break;
     }
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     end = beg + len;
     if (end > RARRAY_LEN(ary)) {
 	if (end >= ARY_CAPA(ary)) {
@@ -2580,6 +2658,7 @@ rb_ary_uniq_bang(VALUE ary)
     VALUE hash, v, vv;
     long i, j;
 
+    ary_iter_check(ary);
     hash = ary_make_hash(ary, 0);
 
     if (RARRAY_LEN(ary) == RHASH(hash)->tbl->num_entries) {
@@ -2632,6 +2711,7 @@ rb_ary_compact_bang(VALUE ary)
     long n;
 
     rb_ary_modify(ary);
+    ary_iter_check(ary);
     p = t = RARRAY_PTR(ary);
     end = p + RARRAY_LEN(ary);
     
@@ -2825,6 +2905,8 @@ rb_ary_shuffle_bang(VALUE ary)
 {
     long i = RARRAY_LEN(ary);
 
+    rb_ary_modify(ary);
+    ary_iter_check(ary);
     while (i) {
 	long j = genrand_real()*i;
 	VALUE tmp = RARRAY_PTR(ary)[--i];