Re: [PATCH] Refine string and array slice assignments

diff --git a/Src/params.c b/Src/params.c
index 1620b2a11..cee10091e 100644
--- a/Src/params.c
+++ b/Src/params.c
@@ -2718,31 +2718,30 @@ assignstrvalue(Value v, char *val, int flags)
 		!v->pm->width)
 		v->pm->width = strlen(val);
 	} else {
-	    char *z, *x;
-            int zlen, vlen, newsize;
-
-            z = v->pm->gsu.s->getfn(v->pm);
-            zlen = strlen(z);
+	    char *const old = v->pm->gsu.s->getfn(v->pm);
+	    char *new;
+	    const int oldlen = strlen(old), vallen = strlen(val);
+	    int newlen;
 
 	    if ((v->valflags & VALFLAG_INV) && unset(KSHARRAYS))
 		v->start--, v->end--;
 	    if (v->start < 0) {
-		v->start += zlen;
+		v->start += oldlen;
 		if (v->start < 0)
 		    v->start = 0;
 	    }
-	    if (v->start > zlen)
-		v->start = zlen;
+	    if (v->start > oldlen)
+		v->start = oldlen;
 	    if (v->end < 0) {
-		v->end += zlen;
+		v->end += oldlen;
 		if (v->end < 0) {
 		    v->end = 0;
-		} else if (v->end >= zlen) {
-		    v->end = zlen;
+		} else if (v->end >= oldlen) {
+		    v->end = oldlen;
 		} else {
 #ifdef MULTIBYTE_SUPPORT
 		    if (isset(MULTIBYTE)) {
-			v->end += MB_METACHARLEN(z + v->end);
+			v->end += MB_METACHARLEN(old + v->end);
 		    } else {
 			v->end++;
 		    }
@@ -2753,30 +2752,33 @@ assignstrvalue(Value v, char *val, int flags)
 	    }
 	    if (v->end < v->start)
 		v->end = v->start;
-	    else if (v->end > zlen)
-		v->end = zlen;
-
-            vlen = strlen(val);
-            /* Characters preceding start index +
-               characters of what is assigned +
-               characters following end index */
-            newsize = v->start + vlen + (zlen - v->end);
-
-            /* Does new size differ? */
-            if (newsize != zlen ||
-		v->pm->gsu.s->setfn != strsetfn || z != v->pm->u.str) {
-                x = (char *) zalloc(newsize + 1);
-                strncpy(x, z, v->start);
-                strcpy(x + v->start, val);
-                strcat(x + v->start, z + v->end);
-                v->pm->gsu.s->setfn(v->pm, x);
-            } else {
-                /* Size doesn't change, can limit actions to only
-                 * overwriting bytes in already allocated string */
-		memcpy(z + v->start, val, vlen);
-		v->pm->gsu.s->setfn(v->pm, z);
-            }
-            zsfree(val);
+	    else if (v->end > oldlen)
+		v->end = oldlen;
+
+	    /* Chars before slice + chars from val + chars after slice */
+	    newlen = v->start + vallen + (oldlen - v->end);
+
+	    if (v->pm->gsu.s->setfn == strsetfn && old == v->pm->u.str) {
+		v->pm->u.str = NULL; /* Steal the old string */
+		/* If the string shrinks, move left chars after the slice */
+		if (newlen < oldlen && v->end < oldlen)
+		    memmove(old + v->start + vallen, old + v->end,
+			    oldlen - v->end);
+		/* If the string shrinks or expands, reallocate it */
+		new = oldlen == newlen ? old : (char *) zrealloc(old, newlen + 1);
+		/* If the string expands, move right chars after the slice */
+		if (newlen > oldlen && v->end < oldlen)
+		    memmove(new + v->start + vallen, new + v->end,
+			    oldlen - v->end);
+	    } else {
+		new = (char *) zalloc(newlen + 1);
+		strncpy(new, old, v->start);
+		strncpy(new + v->start + vallen, old + v->end, oldlen - v->end);
+	    }
+	    strncpy(new + v->start, val, vallen);
+	    zsfree(val);
+	    new[newlen] = '\0';
+	    v->pm->gsu.s->setfn(v->pm, new);
 	}
 	break;
     case PM_INTEGER:
@@ -2923,13 +2925,9 @@ setarrvalue(Value v, char **val)
 	return;
     } else {
 	char **const old = v->pm->gsu.a->getfn(v->pm);
-	char **new;
-	char **p, **q, **r; /* index variables */
-	const int pre_assignment_length = arrlen(old);
-	int post_assignment_length;
-	int i;
-
-	q = old;
+	char **new, **p, **q;
+	const int oldlen = arrlen(old), vallen = arrlen(val);
+	int newlen, i;
 
 	if ((v->valflags & VALFLAG_INV) && unset(KSHARRAYS)) {
 	    if (v->start > 0)
@@ -2937,95 +2935,55 @@ setarrvalue(Value v, char **val)
 	    v->end--;
 	}
 	if (v->start < 0) {
-	    v->start += pre_assignment_length;
+	    v->start += oldlen;
 	    if (v->start < 0)
 		v->start = 0;
 	}
 	if (v->end < 0) {
-	    v->end += pre_assignment_length + 1;
+	    v->end += oldlen + 1;
 	    if (v->end < 0)
 		v->end = 0;
 	}
-	if (v->end < v->start)
+	if (v->end < v->start || v->start > oldlen)
 	    v->end = v->start;
-
-	post_assignment_length = v->start + arrlen(val);
-	if (v->end < pre_assignment_length) {
-	    /* 
-	     * Allocate room for array elements between the end of the slice `v'
-	     * and the original array's end.
-	     */
-	    post_assignment_length += pre_assignment_length - v->end;
-	}
-
-	if (pre_assignment_length == post_assignment_length
-	    && v->pm->gsu.a->setfn == arrsetfn && old == v->pm->u.arr)
-	{
-	    /* v->start is 0-based */
-	    p = old + v->start;
-	    for (r = val; *r;) {
-		/* Free previous string */
-		zsfree(*p);
-		/* Give away ownership of the string */
-		*p++ = *r++;
-	    }
-	    v->pm->gsu.a->setfn(v->pm, old);
+	else if (v->end > oldlen)
+	    v->end = oldlen;
+
+	/* Strings before slice + strings from val + strings after slice */
+	newlen = v->start + vallen + MAX(0, oldlen - v->end);
+
+	if (v->pm->gsu.a->setfn == arrsetfn && old == v->pm->u.arr) {
+	    v->pm->u.arr = NULL; /* Steal the old array */
+	    /* Free strings that are part of the slice */
+	    for (q = old + v->start, i = v->end - v->start; i > 0; i--)
+		zsfree(*q++);
+	    /* If the array shrinks, move left strings after the slice */
+	    if (newlen < oldlen && v->end < oldlen)
+		memmove(old + v->start + vallen, old + v->end,
+			sizeof(char *) * (oldlen - v->end));
+	    /* If the array shrinks or expands, reallocate it */
+	    new = oldlen == newlen ? old :
+		(char **) zrealloc(old, sizeof(char *) * (newlen + 1));
+	    /* If the array expands, move right strings after the slice */
+	    if (newlen > oldlen && v->end < oldlen)
+		memmove(new + v->start + vallen, new + v->end,
+			sizeof(char *) * (oldlen - v->end));
 	} else {
-            /* arr+=( ... )
-             * arr[${#arr}+x,...]=( ... ) */
-            if (post_assignment_length > pre_assignment_length &&
-                    pre_assignment_length <= v->start &&
-                    v->pm->gsu.a->setfn == arrsetfn && old == v->pm->u.arr)
-            {
-                p = new = (char **) zrealloc(old, sizeof(char *)
-                                           * (post_assignment_length + 1));
-
-                p += pre_assignment_length; /* after old elements */
-
-                /* Consider 1 < 0, case for a=( 1 ); a[1,..] =
-                 *          1 < 1, case for a=( 1 ); a[2,..] = */
-                if (pre_assignment_length < v->start) {
-                    for (i = pre_assignment_length; i < v->start; i++) {
-                        *p++ = ztrdup("");
-                    }
-                }
-
-                for (r = val; *r;) {
-                    /* Give away ownership of the string */
-                    *p++ = *r++;
-                }
-
-                /* v->end doesn't matter:
-                 * a=( 1 2 ); a[4,100]=( a b ); echo "${(q@)a}"
-                 * 1 2 '' a b */
-                *p = NULL;
-
-                v->pm->u.arr = NULL;
-                v->pm->gsu.a->setfn(v->pm, new);
-            } else {
-                p = new = (char **) zalloc(sizeof(char *)
-                                           * (post_assignment_length + 1));
-                for (i = 0; i < v->start; i++)
-                    *p++ = i < pre_assignment_length ? ztrdup(*q++) : ztrdup("");
-                for (r = val; *r;) {
-                    /* Give away ownership of the string */
-                    *p++ = *r++;
-                }
-                if (v->end < pre_assignment_length)
-                    for (q = old + v->end; *q;)
-                        *p++ = ztrdup(*q++);
-                *p = NULL;
-
-                v->pm->gsu.a->setfn(v->pm, new);
-            }
-
-	    DPUTS2(p - new != post_assignment_length, "setarrvalue: wrong allocation: %d 1= %lu",
-		   post_assignment_length, (unsigned long)(p - new));
+	    new = (char **) zalloc(sizeof(char *) * (newlen + 1));
+	    for (p = new, q = old, i = MIN(v->start, oldlen); i > 0; i--)
+		*p++ = ztrdup(*q++);
+	    if (v->end < oldlen)
+		for (p = new + v->start + vallen, q = old + v->end; *q;)
+		    *p++ = ztrdup(*q++);
 	}
-
-        /* Ownership of all strings has been
-         * given away, can plainly free */
+	/* Copy and give away ownership of the strings from val */
+	memcpy(new + v->start, val, sizeof(char *) * vallen);
 	free(val);
+	/* Fill any gap between the old array end and the slice start */
+	for (p = new + oldlen, i = v->start - oldlen; i > 0; i--)
+	    *p++ = ztrdup("");
+	new[newlen] = NULL;
+	v->pm->gsu.a->setfn(v->pm, new);
     }
 }
 
diff --git a/Test/A06assign.ztst b/Test/A06assign.ztst
index 6f4f68efa..d653d75b3 100644
--- a/Test/A06assign.ztst
+++ b/Test/A06assign.ztst
@@ -769,3 +769,17 @@
     a=( a b c )
   )
 0:regression workers/54340
+
+ (
+   typeset s  # Creates a string with pm->u.str = NULL
+   s[1]=""    # Value from pm->gsu.s->getfn() may not be stored in pm->u.str
+   s="abcde"  # Frees pm->u.str => Crashes if it contained value from pm->gsu.s->getfn()
+ )
+0:string assignments don't take ownership of non-freeable pointers
+
+ (
+   typeset -a a  # Creates an array with pm->u.arr = NULL
+   a[1]=()       # Value from pm->gsu.a->getfn() may not be stored in pm->u.arr
+   a=( a b c )   # Frees pm->u.arr => Crashes if it contained value from pm->gsu.a->getfn()
+ )
+0:array assignments don't take ownership of non-freeable pointers