[stdlib] Fix String.reserveCapacity underallocation (#65902)

When called on a string that is not uniquely referenced,
`String.reserveCapacity(_:)` ignores the current capacity, using
the passed-in capacity for the size of its new storage. This can
result in an underallocation and write past the end of the new
buffer.

This fix changes the new size calculation to use the current UTF-8
count as the minimum. Non-native or non-unique strings
now allocate the requested capacity (or space enough for the
current contents, if that's larger than what's requested).

rdar://109275875
Fixes #53483

Author: natecook1000

stdlib/public/core/StringGutsRangeReplaceable.swift

@@ -100,14 +100,28 @@ extension _StringGuts {
   // Grow to accommodate at least `n` code units
   @usableFromInline
   internal mutating func grow(_ n: Int) {
-    defer { self._invariantCheck() }
+    defer {
+      self._invariantCheck()
+      _internalInvariant(
+        self.uniqueNativeCapacity != nil && self.uniqueNativeCapacity! >= n)
+    }
 
     _internalInvariant(
       self.uniqueNativeCapacity == nil || self.uniqueNativeCapacity! < n)
 
+    // If unique and native, apply a 2x growth factor to avoid problematic
+    // performance when used in a loop. If one if those doesn't apply, we
+    // can just use the requested capacity (at least the current utf-8 count).
     // TODO: Don't do this! Growth should only happen for append...
-    let growthTarget = Swift.max(n, (self.uniqueNativeCapacity ?? 0) * 2)
+    let growthTarget: Int
+    if let capacity = self.uniqueNativeCapacity {
+      growthTarget = Swift.max(n, capacity * 2)
+    } else {
+      growthTarget = Swift.max(n, self.utf8Count)
+    }
 
+    // `isFastUTF8` is not the same as `isNative`. It can include small
+    // strings or foreign strings that provide contiguous UTF-8 access.
     if _fastPath(isFastUTF8) {
       let isASCII = self.isASCII
       let storage = self.withFastUTF8 {

validation-test/stdlib/String.swift

@@ -2370,4 +2370,45 @@ StringTests.test("SmallString.zeroTrailingBytes") {
   }
 }
 
+StringTests.test("String.CoW.reserveCapacity") {
+  // Test that reserveCapacity(_:) doesn't actually shrink capacity
+  var str = String(repeating: "a", count: 20)
+  str.reserveCapacity(10)
+  expectGE(str.capacity, 20)
+  str.reserveCapacity(30)
+  expectGE(str.capacity, 30)
+  let preGrowCapacity = str.capacity
+  
+  // Growth shouldn't be linear
+  str.append(contentsOf: String(repeating: "z", count: 20))
+  expectGE(str.capacity, preGrowCapacity * 2)
+  
+  // Capacity can shrink when copying, but not below the count
+  var copy = str
+  copy.reserveCapacity(30)
+  expectGE(copy.capacity, copy.count)
+  expectLT(copy.capacity, str.capacity)
+}
+
+StringTests.test("NSString.CoW.reserveCapacity") {
+#if _runtime(_ObjC)
+  func newNSString() -> NSString {
+    NSString(string: String(repeating: "a", count: 20))
+  }
+  
+  // Test that reserveCapacity(_:) doesn't actually shrink capacity
+  var str = newNSString() as String
+  var copy = str
+  copy.reserveCapacity(10)
+  copy.reserveCapacity(30)
+  expectGE(copy.capacity, 30)
+  
+  var str2 = newNSString() as String
+  var copy2 = str2
+  copy2.append(contentsOf: String(repeating: "z", count: 10))
+  expectGE(copy2.capacity, 30)
+  expectLT(copy2.capacity, 40)
+#endif
+}
+
 runAllTests()