[stdlib] String: Fix more potential UB, and rework access patterns

Author: lorentey

stdlib/public/core/StringGuts.swift

@@ -123,7 +123,7 @@ extension _StringGuts {
   // Whether this string has breadcrumbs
   internal var hasBreadcrumbs: Bool {
     return hasSharedStorage
-      || (hasNativeStorage && _object.nativeStorage.hasBreadcrumbs)
+      || (hasNativeStorage && _object.withNativeStorage { $0.hasBreadcrumbs })
   }
 }
 

stdlib/public/core/StringGutsRangeReplaceable.swift

@@ -13,21 +13,21 @@
 // COW helpers
 extension _StringGuts {
   internal var nativeCapacity: Int? {
-      guard hasNativeStorage else { return nil }
-      return _object.nativeStorage.capacity
+    guard hasNativeStorage else { return nil }
+    return _object.withNativeStorage { $0.capacity }
   }
 
   internal var nativeUnusedCapacity: Int? {
-      guard hasNativeStorage else { return nil }
-      return _object.nativeStorage.unusedCapacity
+    guard hasNativeStorage else { return nil }
+    return _object.withNativeStorage { $0.unusedCapacity }
   }
 
   // If natively stored and uniquely referenced, return the storage's total
   // capacity. Otherwise, nil.
   internal var uniqueNativeCapacity: Int? {
     @inline(__always) mutating get {
       guard isUniqueNative else { return nil }
-      return _object.nativeStorage.capacity
+      return _object.withNativeStorage { $0.capacity }
     }
   }
 
@@ -36,7 +36,7 @@ extension _StringGuts {
   internal var uniqueNativeUnusedCapacity: Int? {
     @inline(__always) mutating get {
       guard isUniqueNative else { return nil }
-      return _object.nativeStorage.unusedCapacity
+      return _object.withNativeStorage { $0.unusedCapacity }
     }
   }
 
@@ -54,6 +54,20 @@ extension _StringGuts {
 
 // Range-replaceable operation support
 extension _StringGuts {
+  @inline(__always)
+  internal mutating func updateNativeStorage<R>(
+    _ body: (__StringStorage) -> R
+  ) -> R {
+    let storage = self._object.nativeStorage
+    self = _StringGuts()
+    defer {
+      // We re-initialize from the modified storage to pick up new count, flags,
+      // etc.
+      self = _StringGuts(storage)
+    }
+    return body(storage)
+  }
+
   @inlinable
   internal init(_initialCapacity capacity: Int) {
     self.init()
@@ -243,11 +257,7 @@ extension _StringGuts {
   internal mutating func appendInPlace(
     _ other: UnsafeBufferPointer<UInt8>, isASCII: Bool
   ) {
-    self._object.nativeStorage.appendInPlace(other, isASCII: isASCII)
-
-    // We re-initialize from the modified storage to pick up new count, flags,
-    // etc.
-    self = _StringGuts(self._object.nativeStorage)
+    updateNativeStorage { $0.appendInPlace(other, isASCII: isASCII) }
   }
 
   @inline(never) // slow-path
@@ -256,11 +266,7 @@ extension _StringGuts {
     _internalInvariant(self.uniqueNativeUnusedCapacity != nil)
 
     var iter = Substring(other).utf8.makeIterator()
-    self._object.nativeStorage.appendInPlace(&iter, isASCII: other.isASCII)
-
-    // We re-initialize from the modified storage to pick up new count, flags,
-    // etc.
-    self = _StringGuts(self._object.nativeStorage)
+    updateNativeStorage { $0.appendInPlace(&iter, isASCII: other.isASCII) }
   }
 
   internal mutating func clear() {
@@ -270,8 +276,7 @@ extension _StringGuts {
     }
 
     // Reset the count
-    _object.nativeStorage.clear()
-    self = _StringGuts(_object.nativeStorage)
+    updateNativeStorage { $0.clear() }
   }
 
   internal mutating func remove(from lower: Index, to upper: Index) {
@@ -281,10 +286,7 @@ extension _StringGuts {
     _internalInvariant(lowerOffset <= upperOffset && upperOffset <= self.count)
 
     if isUniqueNative {
-      _object.nativeStorage.remove(from: lowerOffset, to: upperOffset)
-      // We re-initialize from the modified storage to pick up new count, flags,
-      // etc.
-      self = _StringGuts(self._object.nativeStorage)
+      updateNativeStorage { $0.remove(from: lowerOffset, to: upperOffset) }
       return
     }
 
@@ -412,8 +414,9 @@ extension _StringGuts {
 
     let start = bounds.lowerBound._encodedOffset
     let end = bounds.upperBound._encodedOffset
-    _object.nativeStorage.replace(from: start, to: end, with: codeUnits)
-    self = _StringGuts(_object.nativeStorage)
+    updateNativeStorage {
+      $0.replace(from: start, to: end, with: codeUnits)
+    }
     return Range(_uncheckedBounds: (start, start + codeUnits.count))
   }
 
@@ -435,9 +438,10 @@ extension _StringGuts {
 
     let start = bounds.lowerBound._encodedOffset
     let end = bounds.upperBound._encodedOffset
-    _object.nativeStorage.replace(
-      from: start, to: end, with: codeUnits, replacementCount: replCount)
-    self = _StringGuts(_object.nativeStorage)
+    updateNativeStorage {
+      $0.replace(
+        from: start, to: end, with: codeUnits, replacementCount: replCount)
+    }
     return Range(_uncheckedBounds: (start, start + replCount))
   }
 

stdlib/public/core/StringObject.swift

@@ -894,6 +894,13 @@ extension _StringObject {
       discriminatedObjectRawBits & Nibbles.largeAddressMask)
   }
 
+  @inline(__always)
+  @_alwaysEmitIntoClient
+  internal var largeAddress: UnsafeRawPointer {
+    UnsafeRawPointer(bitPattern: largeAddressBits)
+      ._unsafelyUnwrappedUnchecked
+  }
+
   @inlinable @inline(__always)
   internal var nativeUTF8Start: UnsafePointer<UInt8> {
     _internalInvariant(largeFastIsTailAllocated)
@@ -940,10 +947,23 @@ extension _StringObject {
     return _unsafeUncheckedDowncast(storage, to: __StringStorage.self)
 #else
     _internalInvariant(hasNativeStorage)
-    let storageAddress =
-       UnsafeRawPointer(bitPattern:largeAddressBits)._unsafelyUnwrappedUnchecked
-    let unmanagedRef = Unmanaged<__StringStorage>.fromOpaque(storageAddress)
-    return unmanagedRef.takeUnretainedValue()
+    let unmanaged = Unmanaged<__StringStorage>.fromOpaque(largeAddress)
+    return unmanaged.takeUnretainedValue()
+#endif
+  }
+
+  /// Call `body` with the native storage object of `self`, without retaining
+  /// it.
+  @inline(__always)
+  internal func withNativeStorage<R>(
+    _ body: (__StringStorage) -> R
+  ) -> R {
+#if arch(i386) || arch(arm) || arch(arm64_32) || arch(wasm32)
+    return body(nativeStorage)
+#else
+    _internalInvariant(hasNativeStorage)
+    let unmanaged = Unmanaged<__StringStorage>.fromOpaque(largeAddress)
+    return unmanaged._withUnsafeGuaranteedRef { body($0) }
 #endif
   }
 
@@ -957,7 +977,8 @@ extension _StringObject {
 #else
     _internalInvariant(largeFastIsShared && !largeIsCocoa)
     _internalInvariant(hasSharedStorage)
-    return Builtin.reinterpretCast(largeAddressBits)
+    let unmanaged = Unmanaged<__SharedStringStorage>.fromOpaque(largeAddress)
+    return unmanaged.takeUnretainedValue()
 #endif
   }
 
@@ -970,7 +991,8 @@ extension _StringObject {
     return object
 #else
     _internalInvariant(largeIsCocoa && !isImmortal)
-    return Builtin.reinterpretCast(largeAddressBits)
+    let unmanaged = Unmanaged<AnyObject>.fromOpaque(largeAddress)
+    return unmanaged.takeUnretainedValue()
 #endif
   }
 
@@ -979,7 +1001,8 @@ extension _StringObject {
   @inline(__always)
   internal var owner: AnyObject? {
     guard self.isMortal else { return nil }
-    return Builtin.reinterpretCast(largeAddressBits)
+    let unmanaged = Unmanaged<AnyObject>.fromOpaque(largeAddress)
+    return unmanaged.takeUnretainedValue()
   }
 }
 
@@ -1058,7 +1081,8 @@ extension _StringObject {
   @inline(__always)
   internal var objCBridgeableObject: AnyObject {
     _internalInvariant(hasObjCBridgeableObject)
-    return Builtin.reinterpretCast(largeAddressBits)
+    let unmanaged = Unmanaged<AnyObject>.fromOpaque(largeAddress)
+    return unmanaged.takeUnretainedValue()
   }
 
   // Whether the object provides fast UTF-8 contents that are nul-terminated
@@ -1222,7 +1246,8 @@ extension _StringObject {
         }
       }
       if _countAndFlags.isNativelyStored {
-        let anyObj = Builtin.reinterpretCast(largeAddressBits) as AnyObject
+        let unmanaged = Unmanaged<AnyObject>.fromOpaque(largeAddress)
+        let anyObj = unmanaged.takeUnretainedValue()
         _internalInvariant(anyObj is __StringStorage)
       }
     }

stdlib/public/core/StringStorage.swift

@@ -735,7 +735,7 @@ extension _StringGuts {
 
     let mutPtr: UnsafeMutablePointer<_StringBreadcrumbs?>
     if hasNativeStorage {
-      mutPtr = _object.nativeStorage._breadcrumbsAddress
+      mutPtr = _object.withNativeStorage { $0._breadcrumbsAddress }
     } else {
       mutPtr = UnsafeMutablePointer(
         Builtin.addressof(&_object.sharedStorage._breadcrumbs))