From e9b313eab7870f1d056e36d26862cbb279fd4200 Mon Sep 17 00:00:00 2001
From: Nate Cook <natecook@apple.com>
Date: Wed, 17 May 2023 12:11:23 -0500
Subject: [PATCH] [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
---
 .../core/StringGutsRangeReplaceable.swift     | 18 +++++++-
 validation-test/stdlib/String.swift           | 41 +++++++++++++++++++
 2 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/stdlib/public/core/StringGutsRangeReplaceable.swift b/stdlib/public/core/StringGutsRangeReplaceable.swift
index 0a23a1c61dda2..6fa75b682e57b 100644
--- a/stdlib/public/core/StringGutsRangeReplaceable.swift
+++ b/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 {
diff --git a/validation-test/stdlib/String.swift b/validation-test/stdlib/String.swift
index 4e1bb42abfe6c..5a9788ae352be 100644
--- a/validation-test/stdlib/String.swift
+++ b/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()