1 files changed, 106 insertions, 62 deletions

diff --git a/src/array.c b/src/array.c
index df953832b..c7499ab36 100644
--- a/src/array.c
+++ b/src/array.c
@@ -16,28 +16,16 @@
 #define ARY_C_MAX_SIZE (SIZE_MAX / sizeof(mrb_value))
 #define ARY_MAX_SIZE ((ARY_C_MAX_SIZE < (size_t)MRB_INT_MAX) ? (mrb_int)ARY_C_MAX_SIZE : MRB_INT_MAX-1)
 
-static inline mrb_value
-ary_elt(mrb_value ary, mrb_int offset)
-{
-  if (offset < 0 || RARRAY_LEN(ary) <= offset) {
-    return mrb_nil_value();
-  }
-  return RARRAY_PTR(ary)[offset];
-}
-
 static struct RArray*
 ary_new_capa(mrb_state *mrb, mrb_int capa)
 {
   struct RArray *a;
-  mrb_int blen;
+  size_t blen;
 
   if (capa > ARY_MAX_SIZE) {
     mrb_raise(mrb, E_ARGUMENT_ERROR, "array size too big");
   }
   blen = capa * sizeof(mrb_value);
-  if (blen < capa) {
-    mrb_raise(mrb, E_ARGUMENT_ERROR, "array size too big");
-  }
 
   a = (struct RArray*)mrb_obj_alloc(mrb, MRB_TT_ARRAY, mrb->array_class);
   a->ptr = (mrb_value *)mrb_malloc(mrb, blen);
@@ -120,6 +108,10 @@ ary_fill_with_nil(mrb_value *ptr, mrb_int size)
 static void
 ary_modify(mrb_state *mrb, struct RArray *a)
 {
+  if (MRB_FROZEN_P(a)) {
+    mrb_raise(mrb, E_RUNTIME_ERROR, "can't modify frozen array");
+  }
+
   if (ARY_SHARED_P(a)) {
     mrb_shared_array *shared = a->aux.shared;
 
@@ -130,7 +122,7 @@ ary_modify(mrb_state *mrb, struct RArray *a)
     }
     else {
       mrb_value *ptr, *p;
-      mrb_int len;
+      size_t len;
 
       p = a->ptr;
       len = a->len * sizeof(mrb_value);
@@ -173,11 +165,12 @@ ary_make_shared(mrb_state *mrb, struct RArray *a)
 }
 
 static void
-ary_expand_capa(mrb_state *mrb, struct RArray *a, mrb_int len)
+ary_expand_capa(mrb_state *mrb, struct RArray *a, size_t len)
 {
-  mrb_int capa = a->aux.capa;
+  size_t capa = a->aux.capa;
 
   if (len > ARY_MAX_SIZE) {
+  size_error:
     mrb_raise(mrb, E_ARGUMENT_ERROR, "array size too big");
   }
 
@@ -185,12 +178,18 @@ ary_expand_capa(mrb_state *mrb, struct RArray *a, mrb_int len)
     capa = ARY_DEFAULT_LEN;
   }
   while (capa < len) {
-    capa *= 2;
+    if (capa <= ARY_MAX_SIZE / 2) {
+      capa *= 2;
+    }
+    else {
+      capa = len;
+    }
+  }
+  if (capa < len || capa > ARY_MAX_SIZE) {
+    goto size_error;
   }
 
-  if (capa > ARY_MAX_SIZE) capa = ARY_MAX_SIZE; /* len <= capa <= ARY_MAX_SIZE */
-
-  if (capa > a->aux.capa) {
+  if (capa > (size_t)a->aux.capa) {
     mrb_value *expanded_ptr = (mrb_value *)mrb_realloc(mrb, a->ptr, sizeof(mrb_value)*capa);
 
     a->aux.capa = capa;
@@ -254,13 +253,20 @@ mrb_ary_s_create(mrb_state *mrb, mrb_value self)
 }
 
 static void
-ary_concat(mrb_state *mrb, struct RArray *a, mrb_value *ptr, mrb_int blen)
+ary_concat(mrb_state *mrb, struct RArray *a, struct RArray *a2)
 {
-  mrb_int len = a->len + blen;
+  mrb_int len;
+
+  if (a2->len > ARY_MAX_SIZE - a->len) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "array size too big");
+  }
+  len = a->len + a2->len;
 
   ary_modify(mrb, a);
-  if (a->aux.capa < len) ary_expand_capa(mrb, a, len);
-  array_copy(a->ptr+a->len, ptr, blen);
+  if (a->aux.capa < len) {
+    ary_expand_capa(mrb, a, len);
+  }
+  array_copy(a->ptr+a->len, a2->ptr, a2->len);
   mrb_write_barrier(mrb, (struct RBasic*)a);
   a->len = len;
 }
@@ -270,17 +276,16 @@ mrb_ary_concat(mrb_state *mrb, mrb_value self, mrb_value other)
 {
   struct RArray *a2 = mrb_ary_ptr(other);
 
-  ary_concat(mrb, mrb_ary_ptr(self), a2->ptr, a2->len);
+  ary_concat(mrb, mrb_ary_ptr(self), a2);
 }
 
 static mrb_value
 mrb_ary_concat_m(mrb_state *mrb, mrb_value self)
 {
-  mrb_value *ptr;
-  mrb_int blen;
+  mrb_value ary;
 
-  mrb_get_args(mrb, "a", &ptr, &blen);
-  ary_concat(mrb, mrb_ary_ptr(self), ptr, blen);
+  mrb_get_args(mrb, "A", &ary);
+  mrb_ary_concat(mrb, self, ary);
   return self;
 }
 
@@ -318,9 +323,12 @@ ary_replace(mrb_state *mrb, struct RArray *a, mrb_value *argv, mrb_int len)
 MRB_API void
 mrb_ary_replace(mrb_state *mrb, mrb_value self, mrb_value other)
 {
+  struct RArray *a1 = mrb_ary_ptr(self);
   struct RArray *a2 = mrb_ary_ptr(other);
 
-  ary_replace(mrb, mrb_ary_ptr(self), a2->ptr, a2->len);
+  if (a1 != a2) {
+    ary_replace(mrb, a1, a2->ptr, a2->len);
+  }
 }
 
 static mrb_value
@@ -432,6 +440,7 @@ mrb_ary_pop(mrb_state *mrb, mrb_value ary)
 {
   struct RArray *a = mrb_ary_ptr(ary);
 
+  ary_modify(mrb, a);
   if (a->len == 0) return mrb_nil_value();
   return a->ptr[--a->len];
 }
@@ -444,6 +453,7 @@ mrb_ary_shift(mrb_state *mrb, mrb_value self)
   struct RArray *a = mrb_ary_ptr(self);
   mrb_value val;
 
+  ary_modify(mrb, a);
   if (a->len == 0) return mrb_nil_value();
   if (ARY_SHARED_P(a)) {
   L_SHIFT:
@@ -506,6 +516,9 @@ mrb_ary_unshift_m(mrb_state *mrb, mrb_value self)
   mrb_int len;
 
   mrb_get_args(mrb, "*", &vals, &len);
+  if (len > ARY_MAX_SIZE - a->len) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "array size too big");
+  }
   if (ARY_SHARED_P(a)
       && a->aux.shared->refcnt == 1 /* shared only referenced from this array */
       && a->ptr - a->aux.shared->ptr >= len) /* there's room for unshifted item */ {
@@ -563,13 +576,22 @@ mrb_ary_set(mrb_state *mrb, mrb_value ary, mrb_int n, mrb_value val)
   mrb_field_write_barrier_value(mrb, (struct RBasic*)a, val);
 }
 
+static struct RArray*
+ary_dup(mrb_state *mrb, struct RArray *a)
+{
+  struct RArray *d = ary_new_capa(mrb, a->len);
+
+  ary_replace(mrb, d, a->ptr, a->len);
+  return d;
+}
+
 MRB_API mrb_value
 mrb_ary_splice(mrb_state *mrb, mrb_value ary, mrb_int head, mrb_int len, mrb_value rpl)
 {
   struct RArray *a = mrb_ary_ptr(ary);
-  mrb_int tail, size;
   const mrb_value *argv;
-  mrb_int i, argc;
+  mrb_int argc;
+  size_t tail;
 
   ary_modify(mrb, a);
 
@@ -583,40 +605,59 @@ mrb_ary_splice(mrb_state *mrb, mrb_value ary, mrb_int head, mrb_int len, mrb_val
       mrb_raise(mrb, E_INDEX_ERROR, "index is out of array");
     }
   }
-  if (a->len < len || a->len < head + len) {
+  tail = head + len;
+  if (a->len < len || (size_t)a->len < tail) {
     len = a->len - head;
   }
-  tail = head + len;
 
   /* size check */
   if (mrb_array_p(rpl)) {
     argc = RARRAY_LEN(rpl);
     argv = RARRAY_PTR(rpl);
+    if (argv == a->ptr) {
+      struct RArray *r = ary_dup(mrb, a);
+      argv = r->ptr;
+    }
   }
   else {
     argc = 1;
     argv = &rpl;
   }
-  size = head + argc;
-
-  if (tail < a->len) size += a->len - tail;
-  if (size > a->aux.capa)
-    ary_expand_capa(mrb, a, size);
-
-  if (head > a->len) {
+  if (head >= a->len) {
+    if (head > ARY_MAX_SIZE - argc) {
+      mrb_raisef(mrb, E_INDEX_ERROR, "index %S too big", mrb_fixnum_value(head));
+    }
+    len = head + argc;
+    if (len > a->aux.capa) {
+      ary_expand_capa(mrb, a, head + argc);
+    }
     ary_fill_with_nil(a->ptr + a->len, head - a->len);
+    if (argc > 0) {
+      array_copy(a->ptr + head, argv, argc);
+    }
+    a->len = len;
   }
-  else if (head < a->len) {
-    value_move(a->ptr + head + argc, a->ptr + tail, a->len - tail);
-  }
-
-  for (i = 0; i < argc; i++) {
-    *(a->ptr + head + i) = *(argv + i);
-    mrb_field_write_barrier_value(mrb, (struct RBasic*)a, argv[i]);
-  }
+  else {
+    mrb_int alen;
 
-  a->len = size;
+    if (a->len - len > ARY_MAX_SIZE - argc) {
+      mrb_raisef(mrb, E_INDEX_ERROR, "index %S too big", mrb_fixnum_value(a->len + argc - len));
+    }
+    alen = a->len + argc - len;
+    if (alen > a->aux.capa) {
+      ary_expand_capa(mrb, a, alen);
+    }
 
+    if (len != argc) {
+      tail = head + len;
+      value_move(a->ptr + head + argc, a->ptr + tail, a->len - tail);
+      a->len = alen;
+    }
+    if (argc > 0) {
+      value_move(a->ptr + head, argv, argc);
+    }
+  }
+  mrb_write_barrier(mrb, (struct RBasic*)a);
   return ary;
 }
 
@@ -702,7 +743,7 @@ mrb_ary_aget(mrb_state *mrb, mrb_value self)
     switch (mrb_type(index)) {
       /* a[n..m] */
     case MRB_TT_RANGE:
-      if (mrb_range_beg_len(mrb, index, &i, &len, a->len)) {
+      if (mrb_range_beg_len(mrb, index, &i, &len, a->len, TRUE) == 1) {
         return ary_subseq(mrb, a, i, len);
       }
       else {
@@ -766,20 +807,18 @@ mrb_ary_aset(mrb_state *mrb, mrb_value self)
   mrb_value v1, v2, v3;
   mrb_int i, len;
 
+  mrb_ary_modify(mrb, mrb_ary_ptr(self));
   if (mrb_get_args(mrb, "oo|o", &v1, &v2, &v3) == 2) {
-    switch (mrb_type(v1)) {
     /* a[n..m] = v */
-    case MRB_TT_RANGE:
-      if (mrb_range_beg_len(mrb, v1, &i, &len, RARRAY_LEN(self))) {
-        mrb_ary_splice(mrb, self, i, len, v2);
-      }
+    switch (mrb_range_beg_len(mrb, v1, &i, &len, RARRAY_LEN(self), FALSE)) {
+    case 0:                   /* not range */
+      mrb_ary_set(mrb, self, aget_index(mrb, v1), v2);
       break;
-    /* a[n] = v */
-    case MRB_TT_FIXNUM:
-      mrb_ary_set(mrb, self, mrb_fixnum(v1), v2);
+    case 1:                   /* range */
+      mrb_ary_splice(mrb, self, i, len, v2);
       break;
-    default:
-      mrb_ary_set(mrb, self, aget_index(mrb, v1), v2);
+    case 2:                   /* out of range */
+      mrb_raisef(mrb, E_RANGE_ERROR, "%S out of range", v1);
       break;
     }
     return v2;
@@ -877,13 +916,16 @@ static mrb_value
 mrb_ary_rindex_m(mrb_state *mrb, mrb_value self)
 {
   mrb_value obj;
-  mrb_int i;
+  mrb_int i, len;
 
   mrb_get_args(mrb, "o", &obj);
   for (i = RARRAY_LEN(self) - 1; i >= 0; i--) {
     if (mrb_equal(mrb, RARRAY_PTR(self)[i], obj)) {
       return mrb_fixnum_value(i);
     }
+    if (i > (len = RARRAY_LEN(self))) {
+      i = len;
+    }
   }
   return mrb_nil_value();
 }
@@ -915,6 +957,7 @@ mrb_ary_splat(mrb_state *mrb, mrb_value v)
       recv_class,
       mrb_obj_value(mrb_obj_class(mrb, a))
     );
+    /* not reached */
     return mrb_undef_value();
   }
 }
@@ -932,6 +975,7 @@ mrb_ary_clear(mrb_state *mrb, mrb_value self)
 {
   struct RArray *a = mrb_ary_ptr(self);
 
+  ary_modify(mrb, a);
   if (ARY_SHARED_P(a)) {
     mrb_ary_decref(mrb, a->aux.shared);
     ARY_UNSET_SHARED_FLAG(a);