Zero trailing bytes of SmallStrings more efficiently (#59823)

Author: Catfish-Man

stdlib/public/core/SmallString.swift

@@ -241,22 +241,29 @@ extension _SmallString {
   fileprivate mutating func withMutableCapacity(
     _ f: (UnsafeMutableRawBufferPointer) throws -> Int
   ) rethrows {
-    let len = try withUnsafeMutableBytes(of: &self._storage) {
-      (rawBufPtr: UnsafeMutableRawBufferPointer) -> Int in
-      let len = try f(rawBufPtr)
-      UnsafeMutableRawBufferPointer(
-        rebasing: rawBufPtr[len...]
-      ).initializeMemory(as: UInt8.self, repeating: 0)
-      return len
-    }
-    if len == 0 {
+    let len = try withUnsafeMutableBytes(of: &_storage, f)
+
+    if len <= 0 {
+      _debugPrecondition(len == 0)
       self = _SmallString()
       return
     }
-    _internalInvariant(len <= _SmallString.capacity)
+    _SmallString.zeroTrailingBytes(of: &_storage, from: len)
+    self = _SmallString(leading: _storage.0, trailing: _storage.1, count: len)
+  }
 
-    let (leading, trailing) = self.zeroTerminatedRawCodeUnits
-    self = _SmallString(leading: leading, trailing: trailing, count: len)
+  @inlinable
+  @_alwaysEmitIntoClient
+  internal static func zeroTrailingBytes(
+    of storage: inout RawBitPattern, from index: Int
+  ) {
+    _internalInvariant(index > 0)
+    _internalInvariant(index <= _SmallString.capacity)
+    //FIXME: Verify this on big-endian architecture
+    let mask0 = (UInt64(bitPattern: ~0) &>> (8 &* ( 8 &- Swift.min(index, 8))))
+    let mask1 = (UInt64(bitPattern: ~0) &>> (8 &* (16 &- Swift.max(index, 8))))
+    storage.0 &= (index <= 0) ? 0 : mask0.littleEndian
+    storage.1 &= (index <= 8) ? 0 : mask1.littleEndian
   }
 }
 

validation-test/stdlib/String.swift

@@ -2347,4 +2347,27 @@ if #available(SwiftStdlib 5.6, *) {
   }
 }
 
+StringTests.test("SmallString.zeroTrailingBytes") {
+  if #available(SwiftStdlib 5.8, *) {
+    let full: _SmallString.RawBitPattern = (.max, .max)
+    withUnsafeBytes(of: full) {
+      expectTrue($0.allSatisfy({ $0 == 0xff }))
+    }
+
+    let testIndices = [1, 7, 8, _SmallString.capacity]
+    for i in testIndices {
+      // The internal invariants in `_zeroTrailingBytes(of:from:)`
+      expectTrue(0 < i && i <= _SmallString.capacity)
+      print(i)
+      var bits = full
+      _SmallString.zeroTrailingBytes(of: &bits, from: i)
+      withUnsafeBytes(of: bits) {
+        expectTrue($0[..<i].allSatisfy({ $0 == 0xff }))
+        expectTrue($0[i...].allSatisfy({ $0 == 0 }))
+      }
+      bits = (0, 0)
+    }
+  }
+}
+
 runAllTests()