[stdlib] Fix bug in small string uninitialized init

Fix a bug and enforce the invariant that all bits between the last code unit
and the descriminator in a small string should be unset.

Author: milseman

stdlib/public/core/SmallString.swift

@@ -131,6 +131,13 @@ extension _SmallString {
   internal func _invariantCheck() {
     _internalInvariant(count <= _SmallString.capacity)
     _internalInvariant(isASCII == computeIsASCII())
+
+    // No bits should be set between the last code unit and the discriminator
+    var copy = self
+    withUnsafeBytes(of: &copy._storage) {
+      _internalInvariant(
+        $0[count..<_SmallString.capacity].allSatisfy { $0 == 0 })
+    }
   }
   #endif // INTERNAL_CHECKS_ENABLED
 
@@ -213,14 +220,19 @@ extension _SmallString {
   }
 
   // Overwrite stored code units, including uninitialized. `f` should return the
-  // new count.
+  // new count. This will re-establish the invariant after `f` that all bits
+  // between the last code unit and the discriminator are unset.
   @inline(__always)
-  internal mutating func withMutableCapacity(
+  fileprivate mutating func withMutableCapacity(
     _ f: (UnsafeMutableRawBufferPointer) throws -> Int
   ) rethrows {
     let len = try withUnsafeMutableBytes(of: &self._storage) {
       (rawBufPtr: UnsafeMutableRawBufferPointer) -> Int in
-      return try f(rawBufPtr)
+      let len = try f(rawBufPtr)
+      UnsafeMutableRawBufferPointer(
+        rebasing: rawBufPtr[len...]
+      ).initializeMemory(as: UInt8.self, repeating: 0)
+      return len
     }
     if len == 0 {
       self = _SmallString()

test/stdlib/StringCreate.swift

@@ -80,3 +80,27 @@ if #available(macOS 10.16, iOS 14.0, watchOS 7.0, tvOS 14.0, *) {
     expectTrue(empty.isEmpty)
   }
 }
+
+if #available(macOS 10.16, iOS 14.0, watchOS 7.0, tvOS 14.0, *) {
+  StringCreateTests.test("Small string unsafeUninitializedCapacity") {
+    let str1 = "42"
+    let str2 = String(42)
+    expectEqual(str1, str2)
+
+    let str3 = String(unsafeUninitializedCapacity: 2) {
+      $0[0] = UInt8(ascii: "4")
+      $0[1] = UInt8(ascii: "2")
+      return 2
+    }
+    expectEqual(str1, str3)
+
+    // Write into uninitialized space
+    let str4 = String(unsafeUninitializedCapacity: 3) {
+      $0[0] = UInt8(ascii: "4")
+      $0[1] = UInt8(ascii: "2")
+      $0[2] = UInt8(ascii: "X")
+      return 2
+    }
+    expectEqual(str1, str4)
+  }
+}