[stdlib] Fix buffer size of small-capacity strings (#67929)

Author: benrimmington

stdlib/public/core/SmallString.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -243,7 +243,9 @@ extension _SmallString {
   fileprivate mutating func withMutableCapacity(
     _ f: (UnsafeMutableRawBufferPointer) throws -> Int
   ) rethrows {
-    let len = try withUnsafeMutableBytes(of: &_storage, f)
+    let len = try withUnsafeMutableBytes(of: &_storage) {
+      try f(.init(start: $0.baseAddress, count: _SmallString.capacity))
+    }
 
     if len <= 0 {
       _debugPrecondition(len == 0)

stdlib/public/core/String.swift

@@ -2,7 +2,7 @@
 //
 // This source file is part of the Swift.org open source project
 //
-// Copyright (c) 2014 - 2020 Apple Inc. and the Swift project authors
+// Copyright (c) 2014 - 2023 Apple Inc. and the Swift project authors
 // Licensed under Apache License v2.0 with Runtime Library Exception
 //
 // See https://swift.org/LICENSE.txt for license information
@@ -558,7 +558,9 @@ extension String {
     ) throws -> Int
   ) rethrows {
     if _fastPath(capacity <= _SmallString.capacity) {
-      let smol = try _SmallString(initializingUTF8With: initializer)
+      let smol = try _SmallString(initializingUTF8With: {
+        try initializer(.init(start: $0.baseAddress, count: capacity))
+      })
       // Fast case where we fit in a _SmallString and don't need UTF8 validation
       if _fastPath(smol.isASCII) {
         self = String(_StringGuts(smol))

test/stdlib/StringCreate.swift

@@ -70,6 +70,9 @@ if #available(SwiftStdlib 5.3, *) {
     func test(bufferSize: Int, input: [UInt8], expected: String) {
       let strs = (0..<100).map { _ in
         String(unsafeUninitializedCapacity: bufferSize) { buffer in
+          if #available(SwiftStdlib 5.10, *) {
+            expectEqual(bufferSize, buffer.count)
+          }
           _ = buffer.initialize(from: input)
           return input.count
         }