1 files changed, 38 insertions, 20 deletions

diff --git a/src/string.c b/src/string.c
index ce27cdaa1..d6bdd6975 100644
--- a/src/string.c
+++ b/src/string.c
@@ -121,6 +121,9 @@ mrb_str_buf_new(mrb_state *mrb, size_t capa)
 static void
 resize_capa(mrb_state *mrb, struct RString *s, size_t capacity)
 {
+#if SIZE_MAX > MRB_INT_MAX
+    mrb_assert(capacity < MRB_INT_MAX);
+#endif
   if (RSTR_EMBED_P(s)) {
     if (RSTRING_EMBED_LEN_MAX < capacity) {
       char *const tmp = (char *)mrb_malloc(mrb, capacity+1);
@@ -129,15 +132,12 @@ resize_capa(mrb_state *mrb, struct RString *s, size_t capacity)
       RSTR_UNSET_EMBED_FLAG(s);
       s->as.heap.ptr = tmp;
       s->as.heap.len = len;
-      s->as.heap.aux.capa = capacity;
+      s->as.heap.aux.capa = (mrb_int)capacity;
     }
   }
   else {
-#if SIZE_MAX > MRB_INT_MAX
-    mrb_assert(capacity <= MRB_INT_MAX);
-#endif
-    s->as.heap.ptr = (char *)mrb_realloc(mrb, RSTR_PTR(s), capacity+1);
-    s->as.heap.aux.capa = capacity;
+    s->as.heap.ptr = (char*)mrb_realloc(mrb, RSTR_PTR(s), capacity+1);
+    s->as.heap.aux.capa = (mrb_int)capacity;
   }
 }
 
@@ -154,22 +154,26 @@ str_buf_cat(mrb_state *mrb, struct RString *s, const char *ptr, size_t len)
       off = ptr - RSTR_PTR(s);
   }
 
-  if (RSTR_EMBED_P(s))
-    capa = RSTRING_EMBED_LEN_MAX;
-  else
-    capa = s->as.heap.aux.capa;
+  capa = RSTR_CAPA(s);
+  if (capa <= RSTRING_EMBED_LEN_MAX)
+    capa = RSTRING_EMBED_LEN_MAX+1;
 
   total = RSTR_LEN(s)+len;
   if (total >= MRB_INT_MAX) {
+  size_error:
     mrb_raise(mrb, E_ARGUMENT_ERROR, "string size too big");
   }
   if (capa <= total) {
     while (total > capa) {
-      if (capa + 1 >= MRB_INT_MAX / 2) {
-        capa = MRB_INT_MAX;
-        break;
+      if (capa <= MRB_INT_MAX / 2) {
+        capa *= 2;
       }
-      capa = (capa + 1) * 2;
+      else {
+        capa = total;
+      }
+    }
+    if (capa < total || capa > MRB_INT_MAX) {
+      goto size_error;
     }
     resize_capa(mrb, s, capa);
   }
@@ -515,6 +519,7 @@ str_replace(mrb_state *mrb, struct RString *s1, struct RString *s2)
   long len;
 
   check_frozen(mrb, s1);
+  if (s1 == s2) return mrb_obj_value(s1);
   len = RSTR_LEN(s2);
   if (RSTR_SHARED_P(s1)) {
     str_decref(mrb, s1->as.heap.aux.shared);
@@ -665,7 +670,7 @@ mrb_str_modify(mrb_state *mrb, struct RString *s)
   if (RSTR_SHARED_P(s)) {
     mrb_shared_string *shared = s->as.heap.aux.shared;
 
-    if (shared->refcnt == 1 && s->as.heap.ptr == shared->ptr) {
+    if (shared->nofree == 0 && shared->refcnt == 1 && s->as.heap.ptr == shared->ptr) {
       s->as.heap.ptr = shared->ptr;
       s->as.heap.aux.capa = shared->len;
       RSTR_PTR(s)[s->as.heap.len] = '\0';
@@ -691,14 +696,21 @@ mrb_str_modify(mrb_state *mrb, struct RString *s)
   }
   if (RSTR_NOFREE_P(s)) {
     char *p = s->as.heap.ptr;
+    mrb_int len = s->as.heap.len;
 
-    s->as.heap.ptr = (char *)mrb_malloc(mrb, (size_t)s->as.heap.len+1);
+    RSTR_UNSET_NOFREE_FLAG(s);
+    if (len < RSTRING_EMBED_LEN_MAX) {
+      RSTR_SET_EMBED_FLAG(s);
+      RSTR_SET_EMBED_LEN(s, len);
+    }
+    else {
+      s->as.heap.ptr = (char *)mrb_malloc(mrb, (size_t)len+1);
+      s->as.heap.aux.capa = len;
+    }
     if (p) {
-      memcpy(RSTR_PTR(s), p, s->as.heap.len);
+      memcpy(RSTR_PTR(s), p, len);
     }
-    RSTR_PTR(s)[s->as.heap.len] = '\0';
-    s->as.heap.aux.capa = s->as.heap.len;
-    RSTR_UNSET_NOFREE_FLAG(s);
+    RSTR_PTR(s)[len] = '\0';
     return;
   }
 }
@@ -754,8 +766,14 @@ mrb_str_concat(mrb_state *mrb, mrb_value self, mrb_value other)
     other = mrb_str_to_str(mrb, other);
   }
   s2 = mrb_str_ptr(other);
+  if (RSTR_LEN(s2) == 0) {
+    return;
+  }
   len = RSTR_LEN(s1) + RSTR_LEN(s2);
 
+  if (len < 0 || len >= MRB_INT_MAX) {
+    mrb_raise(mrb, E_ARGUMENT_ERROR, "string size too big");
+  }
   if (RSTRING_CAPA(self) < len) {
     resize_capa(mrb, s1, len);
   }