[5.8] [stdlib] Fix String.reserveCapacity underallocation

stdlib/public/core/StringGutsRangeReplaceable.swift

@@ -74,14 +74,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
+  let newElementCount = (preGrowCapacity - str.count) + 10
+  str.append(contentsOf: String(repeating: "z", count: newElementCount))
+  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)
+#endif
+}
+
 runAllTests()