1 files changed, 35 insertions, 28 deletions

diff --git a/Objects/stringobject.c b/Objects/stringobject.c
index 4de8b1f..25a7c92 100644
--- a/Objects/stringobject.c
+++ b/Objects/stringobject.c
@@ -1531,7 +1531,7 @@ mymemreplace(const char *str, int len,		/* input string */
              const char *pat, int pat_len,	/* pattern string to find */
              const char *sub, int sub_len,	/* substitution string */
              int count,				/* number of replacements */
-             int *out_len)
+	     int *out_len)
 {
 	char *out_s;
 	char *new_s;
@@ -1548,41 +1548,48 @@ mymemreplace(const char *str, int len,		/* input string */
 		nfound = count;
 	if (nfound == 0)
 		goto return_same;
+
 	new_len = len + nfound*(sub_len - pat_len);
+	if (new_len == 0) {
+		/* Have to allocate something for the caller to free(). */
+		out_s = (char *)PyMem_MALLOC(1);
+		if (out_s == NULL)
+			return NULL;
+		out_s[0] = '\0';
+	}
+	else {
+		assert(new_len > 0);
+		new_s = (char *)PyMem_MALLOC(new_len);
+		if (new_s == NULL)
+			return NULL;
+		out_s = new_s;
 
-	new_s = (char *)PyMem_MALLOC(new_len);
-	if (new_s == NULL) return NULL;
+		for (; count > 0 && len > 0; --count) {
+			/* find index of next instance of pattern */
+			offset = mymemfind(str, len, pat, pat_len);
+			if (offset == -1)
+				break;
+
+			/* copy non matching part of input string */
+			memcpy(new_s, str, offset);
+			str += offset + pat_len;
+			len -= offset + pat_len;
 
+			/* copy substitute into the output string */
+			new_s += offset;
+			memcpy(new_s, sub, sub_len);
+			new_s += sub_len;
+		}
+		/* copy any remaining values into output string */
+		if (len > 0)
+			memcpy(new_s, str, len);
+	}
 	*out_len = new_len;
-	out_s = new_s;
-
-	while (len > 0) {
-		/* find index of next instance of pattern */
-		offset = mymemfind(str, len, pat, pat_len);
-		/* if not found,  break out of loop */
-		if (offset == -1) break;
-
-		/* copy non matching part of input string */
-		memcpy(new_s, str, offset); /* copy part of str before pat */
-		str += offset + pat_len; /* move str past pattern */
-		len -= offset + pat_len; /* reduce length of str remaining */
-
-		/* copy substitute into the output string */
-		new_s += offset; /* move new_s to dest for sub string */
-		memcpy(new_s, sub, sub_len); /* copy substring into new_s */
-		new_s += sub_len; /* offset new_s past sub string */
-
-		/* break when we've done count replacements */
-		if (--count == 0) break;
-	}
-	/* copy any remaining values into output string */
-	if (len > 0)
-		memcpy(new_s, str, len);
 	return out_s;
 
   return_same:
 	*out_len = -1;
-	return (char*)str;	/* have to cast away constness here */
+	return (char *)str; /* cast away const */
 }